BaseService now terminates itself if a daemon dies unexpectedly

If a daemon finishes while its parent is still running, the parent
terminates itself.

Closes: #679

Author: gsalgado

p2p/discovery.py

@@ -9,6 +9,7 @@
 import collections
 import logging
 import random
+import socket
 import time
 from typing import (
     Any,
@@ -327,19 +328,29 @@ class DiscoveryService(BaseService):
     _last_lookup: float = 0
     _lookup_interval: int = 30
 
-    def __init__(
-            self, proto: DiscoveryProtocol, peer_pool: PeerPool, token: CancelToken = None) -> None:
+    def __init__(self, proto: DiscoveryProtocol, peer_pool: PeerPool,
+                 port: int, token: CancelToken = None) -> None:
         super().__init__(token)
         self.proto = proto
         self.peer_pool = peer_pool
+        self.port = port
 
     async def _run(self) -> None:
+        await self._start_udp_listener()
         connect_loop_sleep = 2
         self.run_task(self.proto.bootstrap())
         while not self.cancel_token.triggered:
             await self.maybe_connect_to_more_peers()
             await self.sleep(connect_loop_sleep)
 
+    async def _start_udp_listener(self) -> None:
+        loop = asyncio.get_event_loop()
+        # TODO: Support IPv6 addresses as well.
+        await loop.create_datagram_endpoint(
+            lambda: self.proto,
+            local_addr=('0.0.0.0', self.port),
+            family=socket.AF_INET)
+
     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:

p2p/service.py

@@ -8,6 +8,7 @@
     Callable,
     List,
     Optional,
+    Set,
     cast,
 )
 from weakref import WeakSet
@@ -31,7 +32,7 @@ def __init__(self) -> None:
 
 class BaseService(ABC, CancellableMixin):
     logger: TraceLogger = None
-    _child_services: List['BaseService']
+    _child_services: Set['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]]
@@ -48,7 +49,7 @@ def __init__(self,
                  loop: asyncio.AbstractEventLoop = None) -> None:
         self.events = ServiceEvents()
         self._run_lock = asyncio.Lock()
-        self._child_services = []
+        self._child_services = set()
         self._tasks = WeakSet()
         self._finished_callbacks = []
 
@@ -126,20 +127,44 @@ def run_task(self, awaitable: Awaitable[Any]) -> None:
 
         If it raises OperationCancelled, that is caught and ignored.
         """
-        async def f() -> None:
+        async def _run_task_wrapper() -> None:
+            self.logger.debug("Running task %s", awaitable)
             try:
                 await awaitable
             except OperationCancelled:
                 pass
-        self._tasks.add(asyncio.ensure_future(f()))
+            except Exception as e:
+                self.logger.warning("Task %s finished unexpectedly: %s", awaitable, e)
+            else:
+                self.logger.debug("Task %s finished with no errors", awaitable)
+        self._tasks.add(asyncio.ensure_future(_run_task_wrapper()))
 
     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)
+        self._child_services.add(child_service)
         self.run_task(child_service.run())
 
+    def run_daemon(self, service: 'BaseService') -> None:
+        """
+        Run a service and keep a reference to it to be considered during the cleanup.
+
+        If the service finishes while we're still running, we'll terminate as well.
+        """
+        self._child_services.add(service)
+
+        async def _run_daemon_wrapper() -> None:
+            try:
+                await service.run()
+            finally:
+                if not self.cancel_token.triggered:
+                    self.logger.debug(
+                        "%s finished while we're still running, terminating as well", service)
+                    self.cancel_token.trigger()
+
+        self.run_task(_run_daemon_wrapper())
+
     async def _run_in_executor(self, callback: Callable[..., Any], *args: Any) -> Any:
         loop = self.get_event_loop()
         return await self.wait(loop.run_in_executor(self._executor, callback, *args))

tests/p2p/test_service.py

@@ -0,0 +1,34 @@
+import asyncio
+
+import pytest
+
+from p2p.service import BaseService
+
+
+class ParentService(BaseService):
+    """A Service which just runs WaitService with run_daemon() and waits for its cancel token to
+    be triggered.
+    """
+
+    async def _run(self):
+        self.daemon = WaitService(token=self.cancel_token)
+        self.run_daemon(self.daemon)
+        await self.cancel_token.wait()
+
+
+class WaitService(BaseService):
+
+    async def _run(self):
+        await self.cancel_token.wait()
+
+
+@pytest.mark.asyncio
+async def test_daemon_exit_causes_parent_cancellation():
+    service = ParentService()
+    asyncio.ensure_future(service.run())
+    await asyncio.sleep(0.01)
+    assert service.daemon.is_running
+    await service.daemon.cancel()
+    await asyncio.sleep(0.01)
+    assert not service.is_running
+    await service.events.cleaned_up.wait()

trinity/server.py

@@ -1,9 +1,7 @@
 import asyncio
 import logging
 import secrets
-import socket
 from typing import (
-    cast,
     Sequence,
     Tuple,
     Type,
@@ -28,7 +26,6 @@
     REPLY_TIMEOUT,
 )
 from p2p.discovery import (
-    DiscoveryProtocol,
     DiscoveryService,
     PreferredNodeDiscoveryProtocol,
 )
@@ -66,8 +63,6 @@
 class Server(BaseService):
     """Server listening for incoming connections"""
     _tcp_listener = None
-    _udp_listener = None
-    _udp_transport = None
 
     peer_pool: PeerPool = None
 
@@ -118,18 +113,6 @@ async def _close_tcp_listener(self) -> None:
             self._tcp_listener.close()
             await self._tcp_listener.wait_closed()
 
-    async def _start_udp_listener(self, discovery: DiscoveryProtocol) -> None:
-        loop = asyncio.get_event_loop()
-        # TODO: Support IPv6 addresses as well.
-        self._udp_transport, _ = await loop.create_datagram_endpoint(
-            lambda: discovery,
-            local_addr=('0.0.0.0', self.port),
-            family=socket.AF_INET)
-
-    async def _close_udp_listener(self) -> None:
-        if self._udp_transport:
-            cast(asyncio.DatagramTransport, self._udp_transport).abort()
-
     def _make_syncer(self, peer_pool: PeerPool) -> BaseService:
         # This method exists only so that ShardSyncer can provide a different implementation.
         return FullNodeSyncer(
@@ -166,8 +149,8 @@ async def _run(self) -> None:
         addr = Address(external_ip, self.port, self.port)
         discovery_proto = PreferredNodeDiscoveryProtocol(
             self.privkey, addr, self.bootstrap_nodes, self.preferred_nodes)
-        await self._start_udp_listener(discovery_proto)
-        self.discovery = DiscoveryService(discovery_proto, self.peer_pool, self.cancel_token)
+        self.discovery = DiscoveryService(
+            discovery_proto, self.peer_pool, self.port, self.cancel_token)
         self.run_child_service(self.peer_pool)
         self.run_child_service(self.discovery)
         self.run_child_service(self.upnp_service)
@@ -176,8 +159,7 @@ async def _run(self) -> None:
 
     async def _cleanup(self) -> None:
         self.logger.info("Closing server...")
-        await asyncio.gather(
-            self._close_tcp_listener(), self._close_udp_listener())
+        await self._close_tcp_listener()
 
     async def receive_handshake(
             self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: