utils, examples: Add ExperimentManager CM

Adds an ExperimentManager context manager for scripts:

 - No more ugly stack traces after execution if something goes wrong

 - No more ugly and easy to forget try-finally blocks

7 files changed, 44 insertions, 20 deletions

diff --git a/examples/example.py b/examples/example.py
index 11e8331..9f0ce03 100755
--- a/examples/example.py
+++ b/examples/example.py
@@ -3,6 +3,7 @@
 # An example script using the rumba package
 
 from rumba.model import *
+from rumba.utils import ExperimentManager
 
 # import testbed plugins
 import rumba.testbeds.emulab as emulab
@@ -44,12 +45,10 @@ exp = rl.Experiment(tb, nodes = [a, b])
 
 print(exp)
 
-try:
+with ExperimentManager(exp):
     exp.swap_in()
     exp.bootstrap_prototype()
     c1 = Client("rinaperf", options ="-t perf -s 1000 -c 10000")
     s1 = Server("rinaperf", arrival_rate=2, mean_duration=5, options = "-l", nodes = [a], clients = [c1])
     sb = StoryBoard(exp, duration=3600, servers = [s1])
     sb.start()
-finally:
-    exp.swap_out()
diff --git a/examples/jfed-rlite.py b/examples/jfed-rlite.py
index d80b56e..5cc087f 100755
--- a/examples/jfed-rlite.py
+++ b/examples/jfed-rlite.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 
 from rumba.model import *
+from rumba.utils import ExperimentManager
 
 import rumba.testbeds.jfed as jfed
 import rumba.prototypes.rlite as rlite
@@ -46,9 +47,7 @@ tb = jfed.Testbed(exp_name = args.expname,
 
 exp = rlite.Experiment(tb, nodes = [a, b])
 
-try:
+with ExperimentManager(exp):
     exp.swap_in()
     exp.install_prototype()
     exp.bootstrap_prototype()
-finally:
-    exp.swap_out()
diff --git a/examples/mouse.py b/examples/mouse.py
index 5ff5c6a..026b3bc 100755
--- a/examples/mouse.py
+++ b/examples/mouse.py
@@ -3,6 +3,7 @@
 # An example script using the rumba package
 
 from rumba.model import *
+from rumba.utils import ExperimentManager
 
 # import testbed plugins
 import rumba.testbeds.emulab as emulab
@@ -15,7 +16,6 @@ import rumba.prototypes.ouroboros as our
 import rumba.prototypes.rlite as rl
 import rumba.prototypes.irati as irati
 
-
 log.set_logging_level('DEBUG')
 
 
@@ -102,8 +102,6 @@ exp = rl.Experiment(tb, nodes = [a, b, c, d, e, f, g, h, i, j, k, l, m, n])
 
 print(exp)
 
-try:
+with ExperimentManager(exp):
     exp.swap_in()
     exp.bootstrap_prototype()
-finally:
-    exp.swap_out()
diff --git a/examples/snake.py b/examples/snake.py
index c6b67b3..6f9c48d 100755
--- a/examples/snake.py
+++ b/examples/snake.py
@@ -3,6 +3,7 @@
 # An example script using the rumba package
 
 from rumba.model import *
+from rumba.utils import ExperimentManager
 
 # import testbed plugins
 import rumba.testbeds.emulab as emulab
@@ -18,6 +19,7 @@ import rumba.prototypes.irati as irati
 import argparse
 import sys
 
+
 description = "Script to create a snake"
 
 argparser = argparse.ArgumentParser(description = description)
@@ -62,8 +64,6 @@ exp = rl.Experiment(tb, nodes = nodes)
 
 print(exp)
 
-try:
+with ExperimentManager(exp):
     exp.swap_in()
     exp.bootstrap_prototype()
-finally:
-    exp.swap_out()
diff --git a/examples/two-layers.py b/examples/two-layers.py
index 6d22673..6e4671f 100755
--- a/examples/two-layers.py
+++ b/examples/two-layers.py
@@ -16,7 +16,7 @@ import rumba.prototypes.rlite as rl
 import rumba.prototypes.irati as irati
 
 import rumba.log as log
-
+from rumba.utils import ExperimentManager
 
 log.set_logging_level('DEBUG')
 
@@ -54,8 +54,10 @@ exp = rl.Experiment(tb, nodes = [a, b, c, d])
 
 print(exp)
 
-try:
+with ExperimentManager(exp):
     exp.swap_in()
     exp.bootstrap_prototype()
-finally:
-    exp.swap_out()
+    sb = StoryBoard(exp, duration=15, servers=[])
+    sb.run_command(7.5, a, 'echo "7.5 secs in. We are at $(hostname)"')
+    sb.run_command(12, b, 'echo "12 secs in. We are at $(hostname)"')
+    sb.start()
diff --git a/examples/vpn.py b/examples/vpn.py
index 4ad9f6d..b2f3c81 100755
--- a/examples/vpn.py
+++ b/examples/vpn.py
@@ -3,6 +3,7 @@
 # An example script using the rumba package
 
 from rumba.model import *
+from rumba.utils import ExperimentManager
 
 # import testbed plugins
 import rumba.testbeds.emulab as emulab
@@ -49,8 +50,6 @@ exp = our.Experiment(tb, nodes = [a, b, c, d])
 
 print(exp)
 
-#try:
+# with ExperimentManager(exp):
 #    exp.swap_in()
 #    exp.bootstrap_prototype()
-#finally:
-#    exp.swap_out()
diff --git a/rumba/utils.py b/rumba/utils.py
new file mode 100644
index 0000000..5348b0a
--- /dev/null
+++ b/rumba/utils.py
@@ -0,0 +1,27 @@
+import time
+
+import rumba.log as log
+import rumba.model as model
+
+logger = log.get_logger(__name__)
+
+
+class ExperimentManager(object):
+
+    def __init__(self, experiment, do_swap_out=True):
+        assert isinstance(experiment, model.Experiment), \
+            'An experiment instance is required.'
+        self.experiment = experiment
+        self.do_swap_out = do_swap_out
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.do_swap_out:
+            self.experiment.swap_out()
+        if exc_val is not None:
+            logger.error('Something went wrong. Got %s: %s',
+                         type(exc_val).__name__, str(exc_val))
+        time.sleep(0.1)  # Give the queue logger enough time to flush.
+        return True  # Suppress the exception we logged: no traceback.