Implement BaseService.run_task(awaitable)

It runs the given awaitable in the background, catches/ignores
OperationCancelled and ensures the task is finished before the
service is finished.

Also removes trinity.Node.add_service() as the same
functionality is provided by BaseService.run_child_service()

Closes: #1118

Author: gsalgado

p2p/DEVELOPMENT.md

@@ -17,18 +17,17 @@ library.
 
 ## BaseService
 
-- If your service runs coroutines in the background (e.g. via `asyncio.ensure_future`), you must
-  ensure they exit when `is_running` is False or when the cancel token is triggered
+- If your service needs to run coroutines in the background, you should use the `BaseService.run_task()` method and
+  ensure they exit when `is_running` is False or when the cancel token is triggered.
 - If your service runs other services in the background, you should pass your CancelToken down to
-  those services and ensure your `_cleanup()` waits for them to cleanup as well
+  those services and run those using `BaseService.run_child_service()`.
 
 ```Python
 class Node(BaseService):
     async def _run(self):
         self.discovery = DiscoveryService(token=self.cancel_token)
-        asyncio.ensure_future(self.discovery.run())
+        self.run_child_service(self.discovery)
+        self.run_task(self.discovery.bootstrap())
         # Node's run logic goes here...
 
-    async def _cleanup(self):
-        await self.discovery.cleaned_up.wait()
 ```

p2p/discovery.py

@@ -335,18 +335,18 @@ def __init__(
 
     async def _run(self) -> None:
         connect_loop_sleep = 2
-        asyncio.ensure_future(self.proto.bootstrap())
+        self.run_task(self.proto.bootstrap())
         while True:
             await self.maybe_connect_to_more_peers()
-            await self.wait(asyncio.sleep(connect_loop_sleep))
+            await self.sleep(connect_loop_sleep)
 
     async def maybe_connect_to_more_peers(self) -> None:
         """Connect to more peers if we're not yet maxed out to max_peers"""
         if self.peer_pool.is_full:
             self.logger.debug("Already connected to %s peers; sleeping", len(self.peer_pool))
             return
 
-        asyncio.ensure_future(self.maybe_lookup_random_node())
+        self.run_task(self.maybe_lookup_random_node())
 
         await self.peer_pool.connect_to_nodes(
             self.proto.get_nodes_to_connect(self.peer_pool.max_peers))

p2p/peer.py

@@ -720,7 +720,7 @@ def unsubscribe(self, subscriber: PeerSubscriber) -> None:
             peer.remove_subscriber(subscriber)
 
     async def start_peer(self, peer: BasePeer) -> None:
-        asyncio.ensure_future(peer.run())
+        self.run_task(peer.run())
         await self.wait(peer.events.started.wait(), timeout=1)
         try:
             # Although connect() may seem like a more appropriate place to perform the DAO fork
@@ -756,7 +756,7 @@ def _add_peer(self,
     async def _run(self) -> None:
         # FIXME: PeerPool should probably no longer be a BaseService, but for now we're keeping it
         # so in order to ensure we cancel all peers when we terminate.
-        asyncio.ensure_future(self._periodically_report_stats())
+        self.run_task(self._periodically_report_stats())
         await self.cancel_token.wait()
 
     async def stop_all_peers(self) -> None:
@@ -933,7 +933,7 @@ async def _periodically_report_stats(self) -> None:
                     self.logger.debug("    %s", line)
             self.logger.debug("== End peer details == ")
             try:
-                await self.wait(asyncio.sleep(self._report_interval))
+                await self.sleep(self._report_interval)
             except OperationCancelled:
                 break
 
@@ -1053,7 +1053,8 @@ def _test() -> None:
         ROPSTEN_VM_CONFIGURATION,
     )
 
-    asyncio.ensure_future(connect_to_peers_loop(peer_pool, nodes))
+    asyncio.ensure_future(peer_pool.run())
+    peer_pool.run_task(connect_to_peers_loop(peer_pool, nodes))
 
     async def request_stuff() -> None:
         # Request some stuff from ropsten's block 2440319
@@ -1086,7 +1087,6 @@ async def exit_on_sigint() -> None:
 
     asyncio.ensure_future(exit_on_sigint())
     asyncio.ensure_future(request_stuff())
-    asyncio.ensure_future(peer_pool.run())
     loop.set_debug(True)
     loop.run_forever()
     loop.close()

p2p/service.py

@@ -4,11 +4,13 @@
 import logging
 from typing import (
     Any,
+    Awaitable,
     Callable,
     List,
     Optional,
     cast,
 )
+from weakref import WeakSet
 
 from eth.tools.logging import TraceLogger
 
@@ -30,6 +32,8 @@ def __init__(self) -> None:
 class BaseService(ABC, CancellableMixin):
     logger: TraceLogger = None
     _child_services: List['BaseService']
+    # Use a WeakSet so that we don't have to bother updating it when tasks finish.
+    _tasks: 'WeakSet[asyncio.Future[Any]]'
     _finished_callbacks: List[Callable[['BaseService'], None]]
     # Number of seconds cancel() will wait for run() to finish.
     _wait_until_finished_timeout = 5
@@ -45,6 +49,7 @@ def __init__(self,
         self.events = ServiceEvents()
         self._run_lock = asyncio.Lock()
         self._child_services = []
+        self._tasks = WeakSet()
         self._finished_callbacks = []
 
         self._loop = loop
@@ -114,12 +119,26 @@ async def run(
     def add_finished_callback(self, finished_callback: Callable[['BaseService'], None]) -> None:
         self._finished_callbacks.append(finished_callback)
 
-    def run_child_service(self, child_service: 'BaseService') -> 'asyncio.Future[Any]':
+    def run_task(self, awaitable: Awaitable[Any]) -> None:
+        """Run the given awaitable in the background.
+
+        The awaitable should return whenever this service's cancel token is triggered.
+
+        If it raises OperationCancelled, that is caught and ignored.
+        """
+        async def f() -> None:
+            try:
+                await awaitable
+            except OperationCancelled:
+                pass
+        self._tasks.add(asyncio.ensure_future(f()))
+
+    def run_child_service(self, child_service: 'BaseService') -> None:
         """
         Run a child service and keep a reference to it to be considered during the cleanup.
         """
         self._child_services.append(child_service)
-        return asyncio.ensure_future(child_service.run())
+        self.run_task(child_service.run())
 
     async def _run_in_executor(self, callback: Callable[..., Any], *args: Any) -> Any:
         loop = self.get_event_loop()
@@ -136,6 +155,7 @@ async def cleanup(self) -> None:
         await asyncio.gather(*[
             child_service.events.cleaned_up.wait()
             for child_service in self._child_services],
+            *[task for task in self._tasks],
             self._cleanup()
         )
         self.events.cleaned_up.set()
@@ -155,10 +175,20 @@ async def cancel(self) -> None:
             await asyncio.wait_for(
                 self.events.cleaned_up.wait(), timeout=self._wait_until_finished_timeout)
         except asyncio.futures.TimeoutError:
-            self.logger.info("Timed out waiting for %s to finish its cleanup, exiting anyway", self)
+            self.logger.info(
+                "Timed out waiting for %s to finish its cleanup, forcibly cancelling pending "
+                "tasks and exiting anyway", self)
+            self._forcibly_cancel_all_tasks()
+            # Sleep a bit because the Future.cancel() method just schedules the callbacks, so we
+            # need to give the event loop a chance to actually call them.
+            await asyncio.sleep(0.5)
         else:
             self.logger.debug("%s finished cleanly", self)
 
+    def _forcibly_cancel_all_tasks(self) -> None:
+        for task in self._tasks:
+            task.cancel()
+
     @property
     def is_running(self) -> bool:
         return self._run_lock.locked()
@@ -187,13 +217,12 @@ async def _run(self) -> None:
         """
         raise NotImplementedError()
 
-    @abstractmethod
     async def _cleanup(self) -> None:
         """Clean up any resources held by this service.
 
         Called after the service's _run() method returns.
         """
-        raise NotImplementedError()
+        pass
 
 
 def service_timeout(timeout: int) -> Callable[..., Any]:

trinity/nodes/base.py

@@ -6,7 +6,6 @@
 )
 from threading import Thread
 from typing import (
-    List,
     Type,
 )
 
@@ -62,7 +61,6 @@ def __init__(self, plugin_manager: PluginManager, chain_config: ChainConfig) ->
         self._headerdb = self._db_manager.get_headerdb()  # type: ignore
 
         self._jsonrpc_ipc_path: Path = chain_config.jsonrpc_ipc_path
-        self._auxiliary_services: List[BaseService] = []
 
     @abstractmethod
     def get_chain(self) -> BaseChain:
@@ -91,12 +89,6 @@ def db_manager(self) -> BaseManager:
     def headerdb(self) -> BaseAsyncHeaderDB:
         return self._headerdb
 
-    def add_service(self, service: BaseService) -> None:
-        if self.is_running:
-            raise RuntimeError("Cannot add an auxiliary service while the node is running")
-        else:
-            self._auxiliary_services.append(service)
-
     def notify_resource_available(self) -> None:
 
         # We currently need this to give plugins the chance to start as soon
@@ -139,9 +131,6 @@ async def _run(self) -> None:
 
             asyncio.run_coroutine_threadsafe(self._ipc_server.run(), loop=ipc_loop)
 
-        for service in self._auxiliary_services:
-            asyncio.ensure_future(service.run())
-
         await self.get_p2p_server().run()
 
     async def _cleanup(self) -> None:
@@ -151,12 +140,6 @@ async def _cleanup(self) -> None:
             # Stop the this IPCServer-specific event loop, so that the IPCServer thread will exit
             self._ipc_loop.stop()
 
-        await asyncio.gather(*[
-            service.events.cleaned_up.wait()
-            for service
-            in self._auxiliary_services
-        ])
-
     def _make_new_loop_thread(self) -> asyncio.AbstractEventLoop:
         new_loop = asyncio.new_event_loop()
 

trinity/nodes/light.py

@@ -51,13 +51,13 @@ def __init__(self, plugin_manager: PluginManager, chain_config: ChainConfig) ->
         self._discovery = DiscoveryService(
             self._discovery_proto, self._peer_pool, self.cancel_token)
         self._peer_chain = LightPeerChain(self.headerdb, self._peer_pool, self.cancel_token)
-        self.add_service(self._discovery)
-        self.add_service(self._peer_pool)
-        self.add_service(self._peer_chain)
         self.notify_resource_available()
 
     async def _run(self) -> None:
-        # TODO add a datagram endpoint service that can be added with self.add_service
+        self.run_child_service(self._discovery)
+        self.run_child_service(self._peer_pool)
+        self.run_child_service(self._peer_chain)
+        # TODO add a datagram endpoint service that can be added with self.run_child_service
         self.logger.info(
             "enode://%s@%s:%s",
             self.nodekey.public_key.to_hex()[2:],

trinity/protocol/common/handlers.py

@@ -38,9 +38,6 @@ async def _run(self) -> None:
 
         await self.cancel_token.wait()
 
-    async def _cleanup(self) -> None:
-        pass
-
     def get_stats(self) -> List[str]:
         manager_attrs = self._managers.keys()
         return [getattr(self, attr).get_stats() for attr in manager_attrs]

trinity/protocol/common/managers.py

@@ -102,9 +102,6 @@ async def _run(self) -> None:
                 else:
                     self.logger.warning("Unexpected message type: %s", cmd.__class__.__name__)
 
-    async def _cleanup(self) -> None:
-        pass
-
     async def _handle_msg(self, msg: TMsg) -> None:
         if self.pending_request is None:
             self.logger.debug(

trinity/sync/__main__.py

@@ -82,7 +82,7 @@ def _test() -> None:
         nodes = DEFAULT_PREFERRED_NODES[network_id]
 
     asyncio.ensure_future(peer_pool.run())
-    asyncio.ensure_future(connect_to_peers_loop(peer_pool, nodes))
+    peer_pool.run_task(connect_to_peers_loop(peer_pool, nodes))
     chain = chain_class(base_db)
     syncer: BaseHeaderChainSyncer = None
     if args.fast:

trinity/sync/common/chain.py

@@ -76,30 +76,25 @@ def register_peer(self, peer: BasePeer) -> None:
 
     async def _handle_msg_loop(self) -> None:
         while self.is_running:
-            try:
-                peer, cmd, msg = await self.wait(self.msg_queue.get())
-            except OperationCancelled:
-                break
-
+            peer, cmd, msg = await self.wait(self.msg_queue.get())
             # Our handle_msg() method runs cpu-intensive tasks in sub-processes so that the main
-            # loop can keep processing msgs, and that's why we use ensure_future() instead of
+            # loop can keep processing msgs, and that's why we use self.run_task() instead of
             # awaiting for it to finish here.
-            asyncio.ensure_future(self.handle_msg(cast(HeaderRequestingPeer, peer), cmd, msg))
+            self.run_task(self.handle_msg(cast(HeaderRequestingPeer, peer), cmd, msg))
 
     async def handle_msg(self, peer: HeaderRequestingPeer, cmd: protocol.Command,
                          msg: protocol._DecodedMsgType) -> None:
         try:
             await self._handle_msg(peer, cmd, msg)
         except OperationCancelled:
-            # Silently swallow OperationCancelled exceptions because we run unsupervised (i.e.
-            # with ensure_future()). Our caller will also get an OperationCancelled anyway, and
-            # there it will be handled.
+            # Silently swallow OperationCancelled exceptions because otherwise they'll be caught
+            # by the except below and treated as unexpected.
             pass
         except Exception:
             self.logger.exception("Unexpected error when processing msg from %s", peer)
 
     async def _run(self) -> None:
-        asyncio.ensure_future(self._handle_msg_loop())
+        self.run_task(self._handle_msg_loop())
         with self.subscribe(self.peer_pool):
             while True:
                 peer_or_finished: Any = await self.wait_first(
@@ -114,12 +109,7 @@ async def _run(self) -> None:
 
                 # Since self._sync_complete is not set, peer_or_finished can only be a Peer
                 # instance.
-                asyncio.ensure_future(self.sync(peer_or_finished))
-
-    async def _cleanup(self) -> None:
-        # We don't need to cancel() anything, but we yield control just so that the coroutines we
-        # run in the background notice the cancel token has been triggered and return.
-        await asyncio.sleep(0)
+                self.run_task(self.sync(peer_or_finished))
 
     async def sync(self, peer: HeaderRequestingPeer) -> None:
         if self._syncing: