Safely handle code checking process being killed.

While waiting for a result, if the process is not alive, it returns an
error status and restarts another process to continue working.

2 files changed, 61 insertions, 3 deletions

diff --git a/yaksh/code_server.py b/yaksh/code_server.py
index d74d35b..75dd9b2 100755
--- a/yaksh/code_server.py
+++ b/yaksh/code_server.py
@@ -44,12 +44,12 @@ def run_as_nobody():
     os.seteuid(nobody.pw_uid)
 
 
-def check_code(job_queue, results):
+def check_code(pid, job_queue, results):
     """Check the code, this runs forever.
     """
     while True:
         uid, json_data, user_dir = job_queue.get(True)
-        results[uid] = dict(status='running', result=None)
+        results[uid] = dict(status='running', pid=pid, result=None)
         data = json.loads(json_data)
         grader = Grader(user_dir)
         result = grader.evaluate(data)
@@ -81,7 +81,7 @@ class ServerPool(object):
         self.job_queue = Queue()
         processes = []
         for i in range(n):
-            p = Process(target=check_code, args=(self.job_queue, self.results))
+            p = self._make_process(i)
             processes.append(p)
         self.processes = processes
         self.app = self._make_app()
@@ -93,11 +93,33 @@ class ServerPool(object):
         app.listen(self.my_port)
         return app
 
+    def _make_process(self, pid):
+        return Process(
+            target=check_code, args=(pid, self.job_queue, self.results)
+        )
+
     def _start_code_servers(self):
         for proc in self.processes:
             if proc.pid is None:
                 proc.start()
 
+    def _handle_dead_process(self, result):
+        if result.get('status') == 'running':
+            pid = result.get('pid')
+            proc = self.processes[pid]
+            if not proc.is_alive():
+                # If the processes is dead, something bad happened so
+                # restart that process.
+                new_proc = self._make_process(pid)
+                self.processes[pid] = new_proc
+                new_proc.start()
+                result['status'] = 'done'
+                result['result'] = json.dumps(dict(
+                    success=False, weight=0.0,
+                    error=['Process ended with exit code %s.'
+                           % proc.exitcode]
+                ))
+
     # Public Protocol ##########
 
     def get_status(self):
@@ -117,6 +139,7 @@ class ServerPool(object):
 
     def get_result(self, uid):
         result = self.results.get(uid, dict(status='unknown'))
+        self._handle_dead_process(result)
         if result.get('status') == 'done':
             self.results.pop(uid)
         return json.dumps(result)
@@ -239,6 +262,7 @@ def main(args=None):
 
     server_pool.run()
 
+
 if __name__ == '__main__':
     args = sys.argv[1:]
     main(args)
diff --git a/yaksh/tests/test_code_server.py b/yaksh/tests/test_code_server.py
index a73c12f..5f80f2d 100644
--- a/yaksh/tests/test_code_server.py
+++ b/yaksh/tests/test_code_server.py
@@ -157,6 +157,40 @@ class TestCodeServer(unittest.TestCase):
         expect = '5 processes, 0 running, 0 queued'
         self.assertTrue(expect in data)
 
+    def test_killing_process_revives_it(self):
+        # Given
+        testdata = {
+            'metadata': {
+                'user_answer': 'import sys; sys.exit()',
+                'language': 'python',
+                'partial_grading': False
+            },
+            'test_case_data': [{'test_case': '',
+                                'test_case_type': 'standardtestcase',
+                                'weight': 0.0}]
+        }
+
+        # When
+        submit(self.url, '0', json.dumps(testdata), '')
+        result = get_result(self.url, '0', block=True)
+
+        # Then
+        data = json.loads(result.get('result'))
+        self.assertFalse(data['success'])
+        self.assertTrue('Process ended with exit code' in data['error'][0])
+
+        # Now check the server status to see if the right number
+        # processes are running.
+        url = "http://localhost:%s/" % SERVER_POOL_PORT
+
+        # When
+        response = urllib.request.urlopen(url)
+        data = response.read().decode('utf-8')
+
+        # Then
+        expect = '5 processes, 0 running, 0 queued'
+        self.assertTrue(expect in data)
+
 
 if __name__ == '__main__':
     unittest.main()