From 1393f6a0304c16b8f47c938f15f11d27dcb1ff1a Mon Sep 17 00:00:00 2001
From: Neil Booth <kyuupichan@gmail.com>
Date: Tue, 8 Nov 2016 00:04:10 +0900
Subject: [PATCH] Move signal handling out of the controller

---
 electrumx_server.py  | 37 ++++++++++++++++++++++++++++++-------
 server/controller.py | 13 -------------
 2 files changed, 30 insertions(+), 20 deletions(-)

diff --git a/electrumx_server.py b/electrumx_server.py
index 60948fa..9429883 100755
--- a/electrumx_server.py
+++ b/electrumx_server.py
@@ -9,16 +9,29 @@
 
 '''Script to kick off the server.'''
 
-
 import asyncio
 import logging
 import os
+import signal
 import traceback
+from functools import partial
 
 from server.env import Env
 from server.controller import Controller
 
 
+def cancel_tasks(loop):
+    # Cancel and collect the remaining tasks
+    tasks = asyncio.Task.all_tasks()
+    for task in tasks:
+        task.cancel()
+
+    try:
+        loop.run_until_complete(asyncio.gather(*tasks))
+    except asyncio.CancelledError:
+        pass
+
+
 def main_loop():
     '''Get tasks; loop until complete.'''
     if os.geteuid() == 0:
@@ -33,16 +46,26 @@ def main_loop():
     #loop.set_debug(True)
 
     controller = Controller(loop, env)
-    controller.start()
 
-    tasks = asyncio.Task.all_tasks(loop)
+    # Signal handlers
+    def on_signal(signame):
+        '''Call on receipt of a signal to cleanly shutdown.'''
+        logging.warning('received {} signal, preparing to shut down'
+                        .format(signame))
+        loop.stop()
+
+    for signame in ('SIGINT', 'SIGTERM'):
+        loop.add_signal_handler(getattr(signal, signame),
+                                partial(on_signal, signame))
+
+    controller.start()
     try:
-        loop.run_until_complete(asyncio.gather(*tasks))
-    except asyncio.CancelledError:
-        logging.warning('task cancelled; asyncio event loop closing')
+        loop.run_forever()
     finally:
         controller.stop()
-        loop.close()
+        cancel_tasks(loop)
+
+    loop.close()
 
 
 def main():
diff --git a/server/controller.py b/server/controller.py
index 6cfd030..a30aeb4 100644
--- a/server/controller.py
+++ b/server/controller.py
@@ -12,7 +12,6 @@ client-serving data such as histories.
 '''
 
 import asyncio
-import signal
 import ssl
 from functools import partial
 
@@ -46,11 +45,6 @@ class Controller(LoggedClass):
         for coro in coros:
             asyncio.ensure_future(coro)
 
-        # Signal handlers
-        for signame in ('SIGINT', 'SIGTERM'):
-            self.loop.add_signal_handler(getattr(signal, signame),
-                                         partial(self.on_signal, signame))
-
     async def on_update(self, height, touched):
         if not self.servers:
             self.servers = await self.start_servers()
@@ -98,10 +92,3 @@ class Controller(LoggedClass):
         '''Close the listening servers.'''
         for server in self.servers:
             server.close()
-
-    def on_signal(self, signame):
-        '''Call on receipt of a signal to cleanly shutdown.'''
-        self.logger.warning('received {} signal, preparing to shut down'
-                            .format(signame))
-        for task in asyncio.Task.all_tasks(self.loop):
-            task.cancel()