[Nested TaskSet] Child TaskSet's on_stop method is not called when GreenletExit

[da-zhanglei]

Describe the bug

I have a locustfile like this. The LightSet and HeavySet 's on_stop method is never called.
The background is, I want to do some cleanup works in LightSet and HeavySet respectively.
Is there anyway to do it?

from locust import HttpLocust, task, TaskSet, constant


class LightSet(TaskSet):
    @task(1)
    def light_1(self):
        print('Light 1')

    def on_stop(self):
        print('Light stop')


class HeavySet(TaskSet):
    @task(1)
    def heavy_1(self):
        print('Heavy 1')

    def on_stop(self):
        print('Heavy stop')


class LightAndHeavySet(TaskSet):
    tasks = {LightSet: 2, HeavySet: 8}


class User(HttpLocust):
    task_set = LightAndHeavySet
    wait_time = constant(1)
    host = 'localhost'

Expected behavior

The LightSet and HeavySet 's on_stop should be called.

Or there is a workaround to do the cleanup for child TaskSet.

Actual behavior

The LightSet and HeavySet 's on_stop is not called.
I cannot find a way to do it(By looking at Docs and StackOverflow).

Steps to reproduce

Environment settings

• OS: Mac OS
• Python version: 3.7
• Locust version: 0.13.5

[ertanden]

Is there some kind of a workaround for this?

[da-zhanglei]

Finally I gave up using "Nested TaskSet". Instead I use different "User" as a workaround.
It worked as expected.

Like this:

from locust import HttpLocust, task, TaskSet, constant


class LightSet(TaskSet):

    @task(1)
    def light_1(self):
        print('Light 1')

    def on_stop(self):
        print('Light stop')


class HeavySet(TaskSet):
    @task(1)
    def heavy_1(self):
        print('Heavy 1')

    def on_stop(self):
        print('Heavy stop')


class HeavyUser(HttpLocust):
    task_set = HeavySet
    weight = 8
    wait_time = constant(1)
    host = 'localhost'


class LightUser(HttpLocust):
    task_set = LightSet
    weight = 2
    wait_time = constant(1)
    host = 'localhost'

[heyman]

I can confirm and reproduce this bug.

The current code for calling on_stop does not take nested TaskSets into account:

locust/locust/core.py

Lines 191 to 193 in ed3b263

	# Run the task_set on_stop method, if it has one
	if hasattr(task_set_instance, "on_stop"):
	task_set_instance.on_stop()

I think the call to on_stop must live in the TaskSet implementation, and not Locust.

Also, with the current implementation on_stop is called even when no --stop-timeout is specified, which I think is wrong (without --stop-timeout specified, it should die as soon as possible).