From fc82425e519617559d2032df3df4a6253293bb42 Mon Sep 17 00:00:00 2001
From: Maxence Boutet <mboutet@explorance.com>
Date: Tue, 17 Nov 2020 18:14:13 -0500
Subject: [PATCH] Add doc strings

---
 locust/dispatch.py          | 33 ++++++++++++++++++++++++++++++---
 locust/distribution.py      | 11 +++++++++--
 locust/test/test_runners.py |  4 ++--
 3 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/locust/dispatch.py b/locust/dispatch.py
index d3de20a858..9c34938ed0 100644
--- a/locust/dispatch.py
+++ b/locust/dispatch.py
@@ -20,6 +20,31 @@ def dispatch_users(
     user_class_occurrences: Dict[str, int],
     spawn_rate: float,
 ) -> Generator[Dict[str, Dict[str, int]], None, None]:
+    """
+    Generator function that dispatches the users
+    in `user_class_occurrences` to the workers.
+    The currently running users is also taken into
+    account.
+
+    It waits an appropriate amount of time between each iteration
+    in order for the spawn rate to be respected, whether running in
+    local or distributed mode.
+
+    The spawn rate is only applicable when additional users are needed.
+    Hence, if `user_class_occurrences` contains less users than there are
+    currently running, this function won't wait and will only run for
+    one iteration. The logic for not stopping users at a rate of `spawn_rate`
+    is that stopping them is a blocking operation, especially when
+    having a `stop_timeout` and users with tasks running for a few seconds or
+    more. If we were to dispatch multiple spawn messages to have a ramp down,
+    we'd run into problems where the previous spawning would be killed
+    by the new message. See the call to `self.spawning_greenlet.kill()` in
+    `:py:meth:`locust.runners.LocalRunner.start` and `:py:meth:`locust.runners.WorkerRunner.worker`.
+
+    :param worker_nodes: List of worker nodes
+    :param user_class_occurrences: Desired number of users for each class
+    :param spawn_rate: The spawn rate
+    """
     initial_dispatched_users = {
         worker_node.id: {
             user_class: worker_node.user_class_occurrences.get(user_class, 0)
@@ -119,10 +144,8 @@ def dispatch_users(
             gevent.sleep(max(0.0, wait_between_dispatch - delta))
 
         # If we are here, it means we have an excess of users for one or more user classes.
-        # Hence, we need to dispatch a last set of users that will bring the desired users
+        # Hence, we need to dispatch a last set of users that will bring the users
         # distribution to the desired one.
-        # TODO: Explain why we don't stop the users at "spawn_rate"
-        #       and why we stop the excess users once at the end.
         yield balanced_users
 
 
@@ -188,6 +211,10 @@ def balance_users_among_workers(
     worker_nodes,  # type: List[WorkerNode]
     user_class_occurrences: Dict[str, int],
 ) -> Dict[str, Dict[str, int]]:
+    """
+    Balance the users among the workers so that
+    each worker gets around the same number of users of each user class
+    """
     balanced_users = {
         worker_node.id: {user_class: 0 for user_class in sorted(user_class_occurrences.keys())}
         for worker_node in worker_nodes
diff --git a/locust/distribution.py b/locust/distribution.py
index 3acdfae239..05aff62b0e 100644
--- a/locust/distribution.py
+++ b/locust/distribution.py
@@ -14,11 +14,18 @@ def weight_users(
     number_of_users: int,
 ) -> Dict[str, int]:
     """
-    Compute users to spawn
+    Compute the desired state of users using the weight of each user class.
+
+    If `number_of_users` is less than `len(user_classes)`, at most one user of each user class
+    is chosen. User classes with higher weight are chosen first.
+
+    If `number_of_users` is greater than or equal to `len(user_classes)`, at least one user of each
+    user class will be chosen. The greater `number_of_users` is, the better the actual distribution
+    of users will match the desired one (as dictated by the weight attributes).
 
     :param user_classes: the list of user class
     :param number_of_users: total number of users
-    :return: the new set of users to run
+    :return: the set of users to run
     """
     assert number_of_users >= 0
 
diff --git a/locust/test/test_runners.py b/locust/test/test_runners.py
index 48ab351a03..e7d1263c72 100644
--- a/locust/test/test_runners.py
+++ b/locust/test/test_runners.py
@@ -387,7 +387,7 @@ def my_task(self):
         runner.spawning_greenlet.join()
         delta = time.time() - ts
         self.assertTrue(
-            0 <= delta <= 0.01, "Expected user count to increase to 10 instantaneously, instead it took %f" % delta
+            0 <= delta <= 0.05, "Expected user count to increase to 10 instantaneously, instead it took %f" % delta
         )
         self.assertTrue(
             runner.user_count == 10, "User count has not decreased correctly to 2, it is : %i" % runner.user_count
@@ -398,7 +398,7 @@ def my_task(self):
         runner.spawning_greenlet.join()
         delta = time.time() - ts
         self.assertTrue(
-            0 <= delta <= 0.01, "Expected user count to decrease to 2 instantaneously, instead it took %f" % delta
+            0 <= delta <= 0.05, "Expected user count to decrease to 2 instantaneously, instead it took %f" % delta
         )
         self.assertTrue(
             runner.user_count == 2, "User count has not decreased correctly to 2, it is : %i" % runner.user_count