Repository: locustio/locust
File: locust/user/task.py (Python)

Code before:
        method within a task makes it possible for a user to be killed mid-task, even if you've
        set a stop_timeout. If this behavour is not desired you should make the user wait using
        gevent.sleep() instead.
        """
        self._check_stop_condition()
        self.user._state = LOCUST_STATE_WAITING
        self._sleep(self.wait_time())
        self._check_stop_condition()
        self.user._state = LOCUST_STATE_RUNNING

    def _sleep(self, seconds):
        gevent.sleep(seconds)
    
    def _check_stop_condition(self):
        if self.user._state == LOCUST_STATE_STOPPING:
            raise StopUser()
    
    def interrupt(self, reschedule=True):
        """
        Interrupt the TaskSet and hand over execution control back to the parent TaskSet.
        
        If *reschedule* is True (default), the parent User will immediately re-schedule,
        and execute, a new task.
        
        This method should not be called by the root TaskSet (the one that is immediately, 
        attached to the User class's *tasks* attribute), but rather in nested TaskSet
        classes further down the hierarchy.
        """
        raise InterruptTaskSet(reschedule)
    
    @property
    def client(self):
        """
        Reference to the :py:attr:`client <locust.User.client>` attribute of the root
        User instance.
        """
        return self.user.client


class DefaultTaskSet(TaskSet):
    """
    Default root TaskSet that executes tasks in User.tasks.
    It executes tasks declared directly on the Locust with the user instance as the task argument.
    """
    def get_next_task(self):
        return random.choice(self.user.tasks)
    
    def execute_task(self, task):
        if hasattr(task, "tasks") and issubclass(task, TaskSet):
            # task is  (nested) TaskSet class
            task(self.user).run()

Review comment (refactor):
Actually we should be able to remove this whole sentence now, since the root TaskSet is an internal one (since the tasks-under-user change).

Code after:
        method within a task makes it possible for a user to be killed mid-task, even if you've
        set a stop_timeout. If this behavour is not desired you should make the user wait using
        gevent.sleep() instead.
        """
        self._check_stop_condition()
        self.user._state = LOCUST_STATE_WAITING
        self._sleep(self.wait_time())
        self._check_stop_condition()
        self.user._state = LOCUST_STATE_RUNNING

    def _sleep(self, seconds):
        gevent.sleep(seconds)
    
    def _check_stop_condition(self):
        if self.user._state == LOCUST_STATE_STOPPING:
            raise StopUser()
    
    def interrupt(self, reschedule=True):
        """
        Interrupt the TaskSet and hand over execution control back to the parent TaskSet.
        
        If *reschedule* is True (default), the parent User will immediately re-schedule,
        and execute, a new task.
        """
        raise InterruptTaskSet(reschedule)
    
    @property
    def client(self):
        """
        Reference to the :py:attr:`client <locust.User.client>` attribute of the root
        User instance.
        """
        return self.user.client


class DefaultTaskSet(TaskSet):
    """
    Default root TaskSet that executes tasks in User.tasks.
    It executes tasks declared directly on the Locust with the user instance as the task argument.
    """
    def get_next_task(self):
        return random.choice(self.user.tasks)
    
    def execute_task(self, task):
        if hasattr(task, "tasks") and issubclass(task, TaskSet):
            # task is  (nested) TaskSet class
            task(self.user).run()
        else:
            # task is a function
            task(self.user)