From 1d14975bd56c341c6d7e13d0ba3114b569495ba4 Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Tue, 13 Oct 2015 01:11:57 +0800
Subject: [PATCH] worker: cleaner termination on exception in user code,
 improve unittest

---
 artiq/master/worker.py      |  6 ++++++
 artiq/master/worker_impl.py |  4 ++++
 artiq/test/worker.py        | 27 +++++++++++++++++++++++++--
 3 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/artiq/master/worker.py b/artiq/master/worker.py
index faa2dcf18..59f3f9b4a 100644
--- a/artiq/master/worker.py
+++ b/artiq/master/worker.py
@@ -21,6 +21,10 @@ class WorkerWatchdogTimeout(Exception):
     pass
 
 
+class WorkerException(Exception):
+    pass
+
+
 class WorkerError(Exception):
     pass
 
@@ -159,6 +163,8 @@ class Worker:
                 return True
             elif action == "pause":
                 return False
+            elif action == "exception":
+                raise WorkerException
             del obj["action"]
             if action == "create_watchdog":
                 func = self.create_watchdog
diff --git a/artiq/master/worker_impl.py b/artiq/master/worker_impl.py
index ad18d9149..8b22d3acf 100644
--- a/artiq/master/worker_impl.py
+++ b/artiq/master/worker_impl.py
@@ -1,6 +1,7 @@
 import sys
 import time
 import os
+import traceback
 
 from artiq.protocols import pyon
 from artiq.tools import file_import
@@ -214,6 +215,9 @@ def main():
                 put_object({"action": "completed"})
             elif action == "terminate":
                 break
+    except:
+        traceback.print_exc()
+        put_object({"action": "exception"})
     finally:
         device_mgr.close_devices()
 
diff --git a/artiq/test/worker.py b/artiq/test/worker.py
index 7be4304a0..59847f6a5 100644
--- a/artiq/test/worker.py
+++ b/artiq/test/worker.py
@@ -8,6 +8,22 @@ from artiq import *
 from artiq.master.worker import *
 
 
+class SimpleExperiment(EnvExperiment):
+    def build(self):
+        pass
+
+    def run(self):
+        pass
+
+
+class ExceptionTermination(EnvExperiment):
+    def build(self):
+        pass
+
+    def run(self):
+        raise TypeError
+
+
 class WatchdogNoTimeout(EnvExperiment):
     def build(self):
         pass
@@ -53,11 +69,11 @@ def _run_experiment(class_name):
         "arguments": dict()
     }
     loop = asyncio.get_event_loop()
-    worker = Worker()
+    worker = Worker(handlers={"log": lambda message: None})
     loop.run_until_complete(_call_worker(worker, expid))
 
 
-class WatchdogCase(unittest.TestCase):
+class WorkerCase(unittest.TestCase):
     def setUp(self):
         if os.name == "nt":
             self.loop = asyncio.ProactorEventLoop()
@@ -65,6 +81,13 @@ class WatchdogCase(unittest.TestCase):
             self.loop = asyncio.new_event_loop()
         asyncio.set_event_loop(self.loop)
 
+    def test_simple_run(self):
+        _run_experiment("SimpleExperiment")
+
+    def test_exception(self):
+        with self.assertRaises(WorkerException):
+            _run_experiment("ExceptionTermination")
+
     def test_watchdog_no_timeout(self):
         _run_experiment("WatchdogNoTimeout")