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.

3 files changed, 109 insertions, 28 deletions

diff --git a/examples/rinaperf_sb-usage.py b/examples/rinaperf_sb-usage.py
new file mode 100644
index 0000000..0e7ca72
--- /dev/null
+++ b/examples/rinaperf_sb-usage.py
@@ -0,0 +1,58 @@
+#!/usr/bin/env python
+
+from rumba.model import *
+from rumba.storyboard import *
+import rumba.testbeds.qemu as qemu
+import rumba.prototypes.irati as irati
+import rumba.log as log
+from rumba.utils import ExperimentManager
+
+log.set_logging_level('DEBUG')
+
+
+n1 = NormalDIF("n1")
+
+n1.add_policy("rmt.pff", "lfa")
+n1.add_policy("security-manager", "passwd")
+
+e1 = ShimEthDIF("e1")
+
+a = Node("a",
+         difs=[n1, e1],
+         dif_registrations={n1: [e1]})
+
+b = Node("b",
+         difs=[e1, n1],
+         dif_registrations={n1: [e1]})
+
+tb = qemu.Testbed(exp_name="example1",
+                  username="root",
+                  password="root")
+
+exp = irati.Experiment(tb, nodes=[a, b])
+
+
+# General setup (can be reused in other scripts as-is)
+storyboard = StoryBoard(duration=30)
+client = Client("rinaperf",
+                options="-t perf -s 1000 -D <duration>",
+                shutdown="")
+server = Server("rinaperf", options="-l", arrival_rate=0.5,
+                mean_duration=5, clients=[client])
+
+
+# Experiment-specific configuration:
+# (This can be done anytime before storyboard.start())
+
+storyboard.set_experiment(exp)
+storyboard.add_server((server, a))
+client.add_node(b)
+
+
+print(exp)
+
+with ExperimentManager(exp):
+    exp.swap_in()
+    exp.bootstrap_prototype()
+    storyboard.start()
+    input('ENTER')
diff --git a/rumba/ssh_support.py b/rumba/ssh_support.py
index 39dbf5f..fbe68c8 100644
--- a/rumba/ssh_support.py
+++ b/rumba/ssh_support.py
@@ -51,7 +51,9 @@ def get_ssh_client():
 
 
 def _print_stream(stream):
-        o = str(stream.read()).strip('b\'\"\\n')
+        o = str(stream.read()).strip('b')
+        o = o.strip('\'\"')
+        o = o.replace('\\n', '')
         if o != "":
             o_array = o.split('\\n')
             for oi in o_array:
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: