Browse Source

HUP handler now starts new worker atexit

Ask Solem 12 years ago
parent
commit
3ddad79207
2 changed files with 13 additions and 5 deletions
  1. 7 3
      celery/apps/worker.py
  2. 6 2
      celery/tests/bin/test_celeryd.py

+ 7 - 3
celery/apps/worker.py

@@ -323,15 +323,19 @@ install_worker_int_handler = partial(
 )
 
 
+def _clone_current_worker():
+    if os.fork() == 0:
+        os.execv(sys.executable, [sys.executable] + sys.argv)
+
+
 def install_worker_restart_handler(worker, sig='SIGHUP'):
 
     def restart_worker_sig_handler(signum, frame):
         """Signal handler restarting the current python program."""
         set_in_sighandler(True)
         safe_say('Restarting celeryd (%s)' % (' '.join(sys.argv), ))
-        pid = os.fork()
-        if pid == 0:
-            os.execv(sys.executable, [sys.executable] + sys.argv)
+        import atexit
+        atexit.register(_clone_current_worker)
         from celery.worker import state
         state.should_stop = True
     platforms.signals[sig] = restart_worker_sig_handler

+ 6 - 2
celery/tests/bin/test_celeryd.py

@@ -621,8 +621,9 @@ class test_signal_handlers(AppCase):
             state.should_stop = False
 
     @disable_stdouts
+    @patch('atexit.register')
     @patch('os.fork')
-    def test_worker_restart_handler(self, fork):
+    def test_worker_restart_handler(self, fork, register):
         fork.return_value = 0
         if getattr(os, 'execv', None) is None:
             raise SkipTest('platform does not have excv')
@@ -637,10 +638,13 @@ class test_signal_handlers(AppCase):
             handlers = self.psig(cd.install_worker_restart_handler, worker)
             handlers['SIGHUP']('SIGHUP', object())
             self.assertTrue(state.should_stop)
+            self.assertTrue(register.called)
+            callback = register.call_args[0][0]
+            callback()
             self.assertTrue(argv)
             argv[:] = []
             fork.return_value = 1
-            handlers['SIGHUP']('SIGHUP', object())
+            callback()
             self.assertFalse(argv)
         finally:
             os.execv = execv