Merge pull request #9 from Nextdoor/kill_znode_on_term

Handle SIGTERMs gracefully by explicitly stopping all service registrations

zk_watcher/version.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = '0.3.2'  # http://semver.org/
+__version__ = '0.3.3'  # http://semver.org/

zk_watcher/zk_watcher.py

@@ -86,6 +86,7 @@ def __init__(self, server, config_file=None, verbose=False):
         self._config_file = config_file
         self._server = server
         self._verbose = verbose
+        self.done = False
 
         # Set up our threading environment
         self._event = threading.Event()
@@ -96,6 +97,7 @@ def __init__(self, server, config_file=None, verbose=False):
 
         # Watch for any signals
         signal.signal(signal.SIGHUP, self._signal_handler)
+        signal.signal(signal.SIGTERM, self._signal_handler)
 
         # Bring in our configuration options
         self._parse_config()
@@ -109,12 +111,17 @@ def __init__(self, server, config_file=None, verbose=False):
     def _signal_handler(self, signum, frame):
         """Watch for certain signals"""
         self.log.warning('Received signal: %s' % signum)
-        if signum == 1:
-            self.log.warning('Received SIGHUP. Reloading config.')
+
+        if signum == signal.SIGHUP:
+            self.log.warning('Reloading config.')
             self._parse_config()
             self._connect()
             self._setup_watchers()
 
+        if signum == signal.SIGTERM:
+            self.log.warning('Terminating all node registrations.')
+            self.stop()
+
     def _parse_config(self):
         """Read in the supplied config file and update our local settings."""
         self.log.debug('Loading config...')
@@ -252,9 +259,14 @@ def run(self):
             self._event.wait(1)
 
         # At this point we must be exiting. Kill off our above threads
+        self.log.info('Shutting down')
         for w in self._watchers:
+            self.log.info('Stopping watcher: %s' % w._path)
             w.stop()
 
+        # Finally, mark us as done
+        self.done = True
+
     def stop(self):
         self._event.set()
 
@@ -453,17 +465,18 @@ def setup_logger():
 
 def main():
     logger = setup_logger()
-    WatcherDaemon(
+    w = WatcherDaemon(
         config_file=options.config,
         server=options.server,
         verbose=options.verbose)
 
-    while True:
+    while True and not w.done:
         try:
             time.sleep(1)
         except KeyboardInterrupt:
-            logger.info('Exiting')
             break
 
+    logger.info('Exiting')
+
 if __name__ == '__main__':
     main()