Rework server_base to spawn a single task

electrumx/lib/server_base.py

@@ -12,7 +12,6 @@ import sys
 import time
 from functools import partial
 
-from electrumx.lib.tasks import Tasks
 from electrumx.lib.util import class_logger
 
 
@@ -22,9 +21,8 @@ class ServerBase(object):
     Derived classes are expected to:
 
     - set PYTHON_MIN_VERSION and SUPPRESS_MESSAGES as appropriate
-    - implement the start_servers() coroutine, called from the run() method.
+    - implement the serve() coroutine, called from the run() method.
       Upon return the event loop runs until the shutdown signal is received.
-    - implement the shutdown() coroutine
     '''
 
     SUPPRESS_MESSAGES = [
@@ -45,7 +43,6 @@ class ServerBase(object):
         self.logger = class_logger(__name__, self.__class__.__name__)
         self.logger.info(f'Python version: {sys.version}')
         self.env = env
-        self.tasks = Tasks()
 
         # Sanity checks
         if sys.version_info < self.PYTHON_MIN_VERSION:
@@ -59,23 +56,13 @@ class ServerBase(object):
                                'To continue as root anyway, restart with '
                                'environment variable ALLOW_ROOT non-empty')
 
-        # Trigger this event to cleanly shutdown
-        self.shutdown_event = asyncio.Event()
+    async def serve(self, shutdown_event):
+        '''Override to provide the main server functionality.
+        Run as a task that will be cancelled to request shutdown.
 
-    async def start_servers(self):
-        '''Override to perform initialization that requires the event loop,
-        and start servers.'''
-        pass
-
-    async def shutdown(self):
-        '''Override to perform the shutdown sequence, if any.'''
-        pass
-
-    def on_signal(self, signame):
-        '''Call on receipt of a signal to cleanly shutdown.'''
-        self.logger.warning('received {} signal, initiating shutdown'
-                            .format(signame))
-        self.shutdown_event.set()
+        Setting the event also shuts down the server.
+        '''
+        shutdown_event.set()
 
     def on_exception(self, loop, context):
         '''Suppress spurious messages it appears we cannot control.'''
@@ -90,30 +77,34 @@ class ServerBase(object):
         '''Run the server application:
 
         - record start time
-        - set the event loop policy as specified by the environment
-        - install SIGINT and SIGKILL handlers to trigger shutdown_event
+        - install SIGINT and SIGTERM handlers to trigger shutdown_event
         - set loop's exception handler to suppress unwanted messages
-        - run the event loop until start_servers() completes
-        - run the event loop until shutdown is signalled
+        - run the event loop until serve() completes
         '''
-        self.start_time = time.time()
+        def on_signal(signame):
+            shutdown_event.set()
+            self.logger.warning(f'received {signame} signal, '
+                                f'initiating shutdown')
 
+        self.start_time = time.time()
         for signame in ('SIGINT', 'SIGTERM'):
             loop.add_signal_handler(getattr(signal, signame),
-                                    partial(self.on_signal, signame))
+                                    partial(on_signal, signame))
         loop.set_exception_handler(self.on_exception)
 
-        self.tasks.create_task(self.start_servers())
-
-        # Wait for shutdown to be signalled, and log it.
-        # Derived classes may want to provide a shutdown() coroutine.
-        await self.shutdown_event.wait()
-        self.logger.info('shutting down')
-        await self.shutdown()
-
-        # Let the loop clean itself up; prevents some silly logs
-        await asyncio.sleep(0.001)
-
+        shutdown_event = asyncio.Event()
+        task = loop.create_task(self.serve(shutdown_event))
+        try:
+            # Wait for shutdown to be signalled, and log it.
+            await shutdown_event.wait()
+            self.logger.info('shutting down')
+            task.cancel()
+            await task
+        finally:
+            await loop.shutdown_asyncgens()
+
+        # Prevent some silly logs
+        await asyncio.sleep(0)
         # Finally, work around an apparent asyncio bug that causes log
         # spew on shutdown for partially opened SSL sockets
         try: