support step load in both standalone run and distributed run

locust/runners.py

@@ -210,6 +210,30 @@ def start_hatching(self, locust_count=None, hatch_rate=None, wait=False):
             else:
                 self.spawn_locusts(wait=wait)
 
+    def start_stepload(self, locust_count, hatch_rate, step_locust_count, step_duration):
+        self.total_clients = locust_count
+        self.hatch_rate = hatch_rate
+        self.step_clients_growth = step_locust_count
+        self.step_duration = step_duration
+        if self.stepload_greenlet:
+            logger.info("There is an ongoing swarming in Step Load mode, will stop it now.")
+            self.greenlet.killone(self.stepload_greenlet)
+        logger.info("Start a new swarming in Step Load mode: total locust count of %d, hatch rate of %d, step locust count of %d, step duration of %d " % (locust_count, hatch_rate, step_locust_count, step_duration))
+        self.state = STATE_INIT
+        self.stepload_greenlet = self.greenlet.spawn(self.stepload_worker)
+        self.stepload_greenlet.link_exception(callback=self.noop)
+
+    def stepload_worker(self):
+        current_num_clients = 0
+        while self.state == STATE_INIT or self.state == STATE_HATCHING or self.state == STATE_RUNNING:
+            current_num_clients += self.step_clients_growth
+            if current_num_clients > int(self.total_clients):
+                logger.info('Step Load is finished.')
+                break
+            self.start_hatching(current_num_clients, self.hatch_rate)
+            logger.info('Step loading: start hatch job of %d locust.' % (current_num_clients))
+            gevent.sleep(self.step_duration)
+
     def stop(self):
         # if we are currently hatching locusts we need to kill the hatching greenlet first
         if self.hatching_greenlet and not self.hatching_greenlet.ready():
@@ -229,6 +253,10 @@ def log_exception(self, node_id, msg, formatted_tb):
         row["nodes"].add(node_id)
         self.exceptions[key] = row
 
+    def noop(self, *args, **kwargs):
+        """ Used to link() greenlets to in order to be compatible with gevent 1.0 """
+        pass
+
 class LocalLocustRunner(LocustRunner):
     def __init__(self, locust_classes, options):
         super(LocalLocustRunner, self).__init__(locust_classes, options)
@@ -252,10 +280,6 @@ def __init__(self, locust_classes, options):
         self.master_bind_port = options.master_bind_port
         self.heartbeat_liveness = options.heartbeat_liveness
         self.heartbeat_interval = options.heartbeat_interval
-
-    def noop(self, *args, **kwargs):
-        """ Used to link() greenlets to in order to be compatible with gevent 1.0 """
-        pass
 
 class SlaveNode(object):
     def __init__(self, id, state=STATE_INIT, heartbeat_liveness=3):
@@ -348,19 +372,6 @@ def start_hatching(self, locust_count, hatch_rate):
 
         self.state = STATE_HATCHING
 
-    def start_stepload(self, locust_count, hatch_rate, step_locust_count, step_duration):
-        self.total_clients = locust_count
-        self.hatch_rate = hatch_rate
-        self.step_clients_growth = step_locust_count
-        self.step_duration = step_duration
-        if self.stepload_greenlet:
-            logger.info("There is an ongoing swarming in Step Load mode, will stop it now.")
-            self.greenlet.killone(self.stepload_greenlet)
-        logger.info("Start a new swarming in Step Load mode: total locust count of %d, hatch rate of %d, step locust count of %d, step duration of %d " % (locust_count, hatch_rate, step_locust_count, step_duration))
-        self.state = STATE_INIT
-        self.stepload_greenlet = self.greenlet.spawn(self.stepload_worker)
-        self.stepload_greenlet.link_exception(callback=self.noop)
-
     def stop(self):
         self.state = STATE_STOPPING
         for client in self.clients.all:
@@ -384,17 +395,6 @@ def heartbeat_worker(self):
                 else:
                     client.heartbeat -= 1
 
-    def stepload_worker(self):
-        current_num_clients = 0
-        while self.state == STATE_INIT or self.state == STATE_HATCHING or self.state == STATE_RUNNING:
-            current_num_clients += self.step_clients_growth
-            if current_num_clients > int(self.total_clients):
-                logger.info('Step Load is finished.')
-                break
-            self.start_hatching(current_num_clients, self.hatch_rate)
-            logger.info('Step loading: start hatch job of %d locust.' % (current_num_clients))
-            gevent.sleep(self.step_duration)
-
     def client_listener(self):
         while True:
             client_id, msg = self.server.recv_from_client()