Clean up temp dirs and processes at exit

nilmrun/processmanager.py

@@ -185,6 +185,7 @@ class ProcessManager(object):
     def __init__(self):
         self.processes = {}
         self.tmpdirs = {}
+        atexit.register(self._atexit)
 
     def _cleanup_tmpdir(self, pid):
         if pid in self.tmpdirs:
@@ -194,6 +195,17 @@ class ProcessManager(object):
                 pass
             del self.tmpdirs[pid]
 
+    def _atexit(self):
+        # Kill remaining processes, remove their dirs
+        for pid in self.processes.keys():
+            try:
+                self.processes[pid].terminate()
+                del self.processes[pid]
+                shutil.rmtree(self.tmpdirs[pid])
+                del self.tmpdirs[pid]
+            except Exception: # pragma: no cover
+                pass
+
     def __iter__(self):
         return iter(self.processes.keys())
 

nilmrun/server.py

@@ -210,6 +210,7 @@ class Server(object):
         manager = serializer_proxy(nilmrun.processmanager.ProcessManager)()
 
         # Build up the application and mount it
+        self._manager = manager
         root = App()
         root.process = AppProcess(manager)
         root.run = AppRun(manager)

tests/test_nilmrun.py

@@ -390,3 +390,20 @@ class TestClient(object):
         # Programs that spit out invalid UTF-8 should get replacement
         # markers
         verify("echo -ne \\\\xae", u"\ufffd")
+
+    def test_client_11_atexit(self):
+        # Leave a directory and running process behind, for the atexit
+        # handler to clean up.  Here we trigger the atexit manually,
+        # since it's hard to trigger it as part of the test suite.
+        client = HTTPClient(baseurl = testurl, post_json = True)
+        code = textwrap.dedent("""
+        import time
+        time.sleep(10)
+        """)
+        client.post("run/code", { "code": code, "args": [ "hello"] })
+
+        # Trigger atexit function
+        test_server._manager._atexit()
+
+        # Ensure no processes exit
+        eq_(client.get("process/list"), [])