storyboard: force client to terminate within storyboard lifetime

Before, client could outlive the storyboard and consequently be
killed abruptly, leading to ugly logs and more.

1 files changed, 19 insertions, 4 deletions

diff --git a/rumba/storyboard.py b/rumba/storyboard.py
index 4a09ac8..5d1202f 100644
--- a/rumba/storyboard.py
+++ b/rumba/storyboard.py
@@ -50,7 +50,8 @@ except ImportError:
         return random.expovariate(1.0 / mean_duration)
 
     logger.debug("Falling back to simple implementations.")
-    # PROBLEM! These logs will almost never be printed... But we might not care
+    # PROBLEM! These logs will almost never be printed...
+    # But we might not care
 
 current_id = -1
 
@@ -79,7 +80,14 @@ class Client(object):
 
     def process(self, duration):
         node = random.choice(self.nodes) if len(self.nodes) > 0 else None
-        return ClientProcess(get_id(), self.ap, self.startup, duration, node, self.shutdown)
+        return ClientProcess(
+            get_id(),
+            self.ap,
+            self.startup,
+            duration,
+            node,
+            self.shutdown
+        )
 
 
 # Base class for client processes
@@ -323,10 +331,17 @@ class StoryBoard:
         try:
             for server in self.servers:
                 server.run()
-            while time.time() - self.start_time < self.duration:
+            while time.time() - self.start_time < (self.duration + 10):
+                # + 10 -> To account for last-minute terminations.
                 for server in self.servers:
                     clients = server.get_new_clients(self.DEFAULT_INTERVAL)
-                    for new_client in clients:
+                    for new_client in clients:  # type: ClientProcess
+                        new_client.duration = min(
+                            new_client.duration,
+                            self.duration - time.time()
+                        )
+                        # Make sure the duration of the client does not
+                        # go beyond the storyboard lifetime
                         new_client.run()
                         self.active_clients.append(new_client)
                 surviving = []