Added user-defined wait_function to locust and TaskSet

docs/api.rst

@@ -7,20 +7,20 @@ Locust class
 ============
 
 .. autoclass:: locust.core.Locust
-	:members: min_wait, max_wait, task_set, weight
+	:members: min_wait, max_wait, wait_function, task_set, weight
 
 HttpLocust class
 ================
 
 .. autoclass:: locust.core.HttpLocust
-	:members: min_wait, max_wait, task_set, client
+	:members: min_wait, max_wait, wait_function, task_set, client
 
 
 TaskSet class
 =============
 
 .. autoclass:: locust.core.TaskSet
-	:members: locust, parent, min_wait, max_wait, client, tasks, interrupt, schedule_task
+	:members: locust, parent, min_wait, max_wait, wait_function, client, tasks, interrupt, schedule_task
 
 task decorator
 ==============

examples/custom_wait_function.py

@@ -0,0 +1,51 @@
+from locust import HttpLocust, TaskSet, task
+import random
+
+def index(l):
+    l.client.get("/")
+
+def stats(l):
+    l.client.get("/stats/requests")
+
+class UserTasks(TaskSet):
+    # one can specify tasks like this
+    tasks = [index, stats]
+
+    # but it might be convenient to use the @task decorator
+    @task
+    def page404(self):
+        self.client.get("/does_not_exist")
+
+class WebsiteUser(HttpLocust):
+    """
+    Locust user class that does requests to the locust web server running on localhost
+    """
+    host = "http://127.0.0.1:8089"
+    # Most task inter-arrival times approximate to exponential distributions
+    # We will model this wait time as exponentially distributed with a mean of 1 second
+    wait_function = lambda: random.expovariate(1)*1000 # *1000 to convert to milliseconds
+    task_set = UserTasks
+
+def strictExp(min_wait,max_wait,mu=1):
+    """
+    Returns an exponentially distributed time strictly between two bounds.
+    """
+    while True:
+        x = random.expovariate(mu)
+        increment = (max_wait-min_wait)/(mu*6.0)
+        result = min_wait + (x*increment)
+        if result<max_wait:
+            break
+    return result
+
+class StrictWebsiteUser(HttpLocust):
+    """
+    Locust user class that makes exponential requests but strictly between two bounds.
+    """
+    host = "http://127.0.0.1:8089"
+    wait_function = lambda self: strictExp(self.min_wait, self.max_wait)*1000
+    task_set = UserTasks
+
+
+
+

locust/core.py

@@ -87,6 +87,9 @@ class Locust(object):
 
     max_wait = 1000
     """Maximum waiting time between the execution of locust tasks"""
+
+    wait_function = lambda self: random.randint(self.min_wait,self.max_wait) 
+    """Function used to calculate waiting time between the execution of locust tasks in milliseconds"""
 
     task_set = None
     """TaskSet class that defines the execution behaviour of this locust"""
@@ -204,9 +207,9 @@ class TaskSet(object):
     Class defining a set of tasks that a Locust user will execute. 
     
     When a TaskSet starts running, it will pick a task from the *tasks* attribute, 
-    execute it, call it's wait function which will sleep a random number between
-    *min_wait* and *max_wait* milliseconds. It will then schedule another task for 
-    execution and so on.
+    execute it, and call its *wait_function* which will define a time to sleep for. 
+    This defaults to a uniformly distributed random number between *min_wait* and 
+    *max_wait* milliseconds. It will then schedule another task for execution and so on.
     
     TaskSets can be nested, which means that a TaskSet's *tasks* attribute can contain 
     another TaskSet. If the nested TaskSet it scheduled to be executed, it will be 
@@ -246,6 +249,13 @@ class ForumPage(TaskSet):
     TaskSet.
     """
 
+    wait_function = None
+    """
+    Function used to calculate waiting time betwen the execution of locust tasks in milliseconds. 
+    Can be used to override the wait_function defined in the root Locust class, which will be used
+    if not set on the TaskSet.
+    """
+
     locust = None
     """Will refer to the root Locust class instance when the TaskSet has been instantiated"""
 
@@ -272,11 +282,13 @@ def __init__(self, parent):
 
         self.parent = parent
 
-        # if this class doesn't have a min_wait or max_wait defined, copy it from Locust
+        # if this class doesn't have a min_wait, max_wait or wait_function defined, copy it from Locust
         if not self.min_wait:
             self.min_wait = self.locust.min_wait
         if not self.max_wait:
             self.max_wait = self.locust.max_wait
+        if not self.wait_function:
+            self.wait_function = self.locust.wait_function
 
         self._lock.acquire()
         if hasattr(self, "setup") and self._setup_has_run is False:
@@ -377,10 +389,12 @@ def schedule_task(self, task_callable, args=None, kwargs=None, first=False):
     def get_next_task(self):
         return random.choice(self.tasks)
 
+    def get_wait_secs(self):
+        millis = self.wait_function()
+        return millis / 1000.0
+
     def wait(self):
-        millis = random.randint(self.min_wait, self.max_wait)
-        seconds = millis / 1000.0
-        self._sleep(seconds)
+        self._sleep(self.get_wait_secs())
 
     def _sleep(self, seconds):
         gevent.sleep(seconds)

locust/test/test_locust_class.py

@@ -177,6 +177,14 @@ def t1(self):
         taskset = MyTaskSet3(self.locust)
         self.assertEqual(len(taskset.tasks), 3)
 
+    def test_wait_function(self):
+        class MyTaskSet(TaskSet):
+            min_wait = 1000
+            max_wait = 2000
+            wait_function = lambda self: 1000 + (self.max_wait-self.min_wait)
+        taskset = MyTaskSet(self.locust)
+        self.assertEqual(taskset.get_wait_secs(), 2.0)
+
     def test_sub_taskset(self):
         class MySubTaskSet(TaskSet):
             min_wait = 1