storyboard: make Client shutdown an argument

Before the shutdown call for the Client/ClientProcess was fixed to kill <pid>.
Now it's passed as an argument to the Client constructor.

1 files changed, 48 insertions, 27 deletions

diff --git a/rumba/storyboard.py b/rumba/storyboard.py
index 8c73422..52bd875 100644
--- a/rumba/storyboard.py
+++ b/rumba/storyboard.py
@@ -52,14 +52,15 @@ except ImportError:
 
 
 class Client(object):
-    def __init__(self, ap, nodes=None, options=None):
+    def __init__(self, ap, nodes=None, options=None, shutdown="kill <pid>"):
         self.ap = ap
-        self.options = options
+        self.startup = (ap + ((" " + options) if options is not None else ""))
         if isinstance(nodes, model.Node):
             nodes = [nodes]
-        if nodes is None:
+        elif nodes is None:
             nodes = []
         self.nodes = nodes
+        self.shutdown = shutdown
 
     def add_node(self, node):
         if not isinstance(node, model.Node):
@@ -68,7 +69,7 @@ class Client(object):
 
     def process(self, duration):
         node = random.choice(self.nodes) if len(self.nodes) > 0 else None
-        return ClientProcess(self.ap, duration, node, self.options)
+        return ClientProcess(self.ap, self.startup, duration, node, self.shutdown)
 
 
 # Base class for client processes
@@ -78,17 +79,21 @@ class Client(object):
 # @start_time: The time at which this process is started.
 # @options: Options to pass to the binary
 #
-class ClientProcess(Client):
-    def __init__(self, ap, duration, node=None, options=None):
-        super(ClientProcess, self).__init__(ap, node, options=options)
+class ClientProcess(object):
+    def __init__(self, ap, startup, duration,
+                 node=None, shutdown="<kill <pid>"):
+        self.ap = ap
+        self.startup = startup
         self.duration = duration
         self.start_time = None
         self.running = False
         self.node = node
         self.pid = None
+        self.shutdown = shutdown
 
     def run(self, node=None):
-        self.node = node
+        if node is not None:
+            self.node = node
         if self.node is None:
             raise Exception('No node specified for client %s' % (self.ap,))
         self.start_time = time.time()
@@ -98,11 +103,12 @@ class ClientProcess(Client):
             self.ap, self.node.name, self.duration
         )
 
-        opt_str = self.options if self.options is not None else ""
-        cmd = "./startup.sh %s %s" % (self.ap, opt_str)
+        start_cmd = "./startup.sh %s" % (
+            self.startup.replace("<duration>", str(self.duration)),
+        )
         self.running = True
         try:
-            self.pid = self.node.execute_command(cmd)
+            self.pid = self.node.execute_command(start_cmd)
         except ssh_support.SSHException:
             logger.warning('Could not start client %s on node %s.',
                            self.ap, self.node.name)
@@ -110,15 +116,22 @@ class ClientProcess(Client):
                      self.ap, self.node.name, self.pid)
 
     def stop(self):
-        logger.debug(
-            'Killing client %s on node %s.',
-            self.ap, self.node.name
-        )
-        try:
-            self.node.execute_command("kill %s" % self.pid)
-        except ssh_support.SSHException:
-                logger.warn('Could not kill client %s on node %s.',
-                            self.ap, self.node.name)
+        if self.shutdown != "":
+            logger.debug(
+                'Killing client %s on node %s.',
+                self.ap, self.node.name
+            )
+            try:
+                kill_cmd = self.shutdown.replace('<pid>', str(self.pid))
+                self.node.execute_command(kill_cmd)
+            except ssh_support.SSHException:
+                    logger.warn('Could not kill client %s on node %s.',
+                                self.ap, self.node.name)
+        else:
+            logger.debug(
+                'Client %s on node %s has terminated.',
+                self.ap, self.node.name
+            )
 
     def check(self):
         """Check if the process should keep running, stop it if not,
@@ -153,6 +166,8 @@ class Server:
         if clients is None:
             clients = list()
         self.clients = clients
+        if nodes is None:
+            nodes = []
         self.nodes = nodes
         self.arrival_rate = arrival_rate  # mean requests/s
         self.mean_duration = mean_duration  # in seconds
@@ -187,22 +202,28 @@ class Server:
         if len(self.clients) == 0:
             raise Exception("Server %s has empty client list." % (self,))
         duration = exponential(self.mean_duration)
-        return random.choice(self.clients).process(duration=duration)
+        return random.choice(self.clients).process(
+            duration=float("%.2f" % (duration,))
+        )
 
     def run(self):
         for node in self.nodes:
-            opt_str = self.options
             logfile = "%s_server.log" % self.ap
             script = r'nohup "$@" > %s & echo "$!"' % (logfile,)
-            cmds = ["echo '%s' > startup.sh && chmod a+x startup.sh"
-                    % (script,),
-                    "./startup.sh %s %s" % (self.ap, opt_str)]
+            run_cmd = self.ap + (
+                (" " + self.options) if self.options is not None else ""
+            )
+            cmd_1 = "echo '%s' > startup.sh && chmod a+x startup.sh && cat startup.sh" \
+                    % (script,)
+            cmd_2 = "./startup.sh %s" % (run_cmd,)
             logger.debug(
                 'Starting server %s on node %s with logfile %s.',
                 self.ap, node.name, logfile
             )
             try:
-                self.pids[node] = (node.execute_commands(cmds))
+                node.execute_command(cmd_1)
+                self.pids[node] = (node.execute_command(cmd_2))
+                node.execute_command("cat startup.sh")
             except ssh_support.SSHException:
                 logger.warn('Could not start server %s on node %s.',
                             self.ap, node.name)
@@ -276,7 +297,7 @@ class StoryBoard:
 
     def start(self):
         self.start_time = time.time()
-        script = r'nohup "$@" > /tmp/ & echo "$!"'
+        script = r'nohup "$@" > /tmp/$1 & echo "$!"'
         logger.debug("Writing utility startup script on client nodes.")
         for server in self.servers:
             for client in server.clients: