Merge pull request #1100 from pipermerriam/piper/use-cancel-token-library

convert to use CancelToken from libary

Author: pipermerriam

p2p/DEVELOPMENT.md

@@ -7,10 +7,12 @@ package, but for now it's just a collection of random notes/recommendations.
 ## Task cancellation
 
 In order to make sure we stop all pending asyncio tasks upon exit (or when a service terminates),
-we use `CancelToken`s, which are heavily inspired by https://vorpus.org/blog/timeouts-and-cancellation-for-humans/
+we use `CancelToken`s from the
+[asyncio-cancel-token](https://asyncio-cancel-token.readthedocs.io/en/latest/index.html)
+library.
 
 - A `CancelToken` must be available to all our async APIs. Either as an instance attribute or as an explicit argument.
-- When one of our async APIs `await` for stdlib/third-party coroutines, it must use `wait_with_token()` to ensure the scheduled task is cancelled when the token is triggered.
+- When one of our async APIs `await` for stdlib/third-party coroutines, it must use `CancelToken.cancellable_wait()` to ensure the scheduled task is cancelled when the token is triggered.
 
 
 ## BaseService

p2p/auth.py

@@ -17,9 +17,10 @@
 
 from eth_hash.auto import keccak
 
+from cancel_token import CancelToken
+
 from p2p import ecies
 from p2p import kademlia
-from p2p.cancel_token import CancelToken, wait_with_token
 from p2p.constants import REPLY_TIMEOUT
 from p2p.exceptions import (
     BadAckMessage,
@@ -74,9 +75,8 @@ async def _handshake(initiator: 'HandshakeInitiator', reader: asyncio.StreamRead
     auth_init = initiator.encrypt_auth_message(auth_msg)
     writer.write(auth_init)
 
-    auth_ack = await wait_with_token(
+    auth_ack = await token.cancellable_wait(
         reader.read(ENCRYPTED_AUTH_ACK_LEN),
-        token=token,
         timeout=REPLY_TIMEOUT)
 
     if reader.at_eof():
@@ -118,9 +118,8 @@ def pubkey(self) -> datatypes.PublicKey:
         return self.privkey.public_key
 
     async def connect(self) -> Tuple[asyncio.StreamReader, asyncio.StreamWriter]:
-        return await wait_with_token(
+        return await self.cancel_token.cancellable_wait(
             asyncio.open_connection(host=self.remote.address.ip, port=self.remote.address.tcp_port),
-            token=self.cancel_token,
             timeout=REPLY_TIMEOUT)
 
     def derive_secrets(self,

p2p/cancel_token.py

@@ -0,0 +1,41 @@
+from typing import (
+    Awaitable,
+    TypeVar,
+)
+
+from cancel_token import CancelToken
+
+
+class CancellableMixin:
+    cancel_token: CancelToken = None
+
+    _TReturn = TypeVar('_TReturn')
+
+    async def wait(self,
+                   awaitable: Awaitable[_TReturn],
+                   token: CancelToken = None,
+                   timeout: float = None) -> _TReturn:
+        """See wait_first()"""
+        return await self.wait_first(awaitable, token=token, timeout=timeout)
+
+    async def wait_first(self,
+                         *awaitables: Awaitable[_TReturn],
+                         token: CancelToken = None,
+                         timeout: float = None) -> _TReturn:
+        """
+        Wait for the first awaitable to complete, unless we timeout or the token chain is triggered.
+
+        The given token is chained with this service's token, so triggering either will cancel
+        this.
+
+        Returns the result of the first one to complete.
+
+        Raises TimeoutError if we timeout or OperationCancelled if the token chain is triggered.
+
+        All pending futures are cancelled before returning.
+        """
+        if token is None:
+            token_chain = self.cancel_token
+        else:
+            token_chain = token.chain(self.cancel_token)
+        return await token_chain.cancellable_wait(*awaitables, timeout=timeout)

p2p/cancellable.py

@@ -23,6 +23,8 @@
 
 from eth_typing import BlockNumber, Hash32
 
+from cancel_token import CancelToken
+
 from eth.constants import (
     BLANK_ROOT_HASH, EMPTY_UNCLE_HASH, GENESIS_BLOCK_NUMBER, GENESIS_PARENT_HASH)
 from eth.chains import AsyncChain
@@ -36,7 +38,7 @@
 from p2p import protocol
 from p2p import eth
 from p2p import les
-from p2p.cancel_token import CancellableMixin, CancelToken
+from p2p.cancellable import CancellableMixin
 from p2p.constants import MAX_REORG_DEPTH, SEAL_CHECK_RANDOM_SAMPLE_RATE
 from p2p.exceptions import NoEligiblePeers, OperationCancelled
 from p2p.p2p_proto import DisconnectReason

p2p/chain.py

@@ -42,7 +42,8 @@
 
 from eth_hash.auto import keccak
 
-from p2p.cancel_token import CancelToken
+from cancel_token import CancelToken
+
 from p2p.exceptions import NoEligibleNodes, OperationCancelled
 from p2p import kademlia
 from p2p.peer import PeerPool

p2p/discovery.py

@@ -40,7 +40,7 @@
 
 from eth_hash.auto import keccak
 
-from p2p.cancel_token import CancelToken, wait_with_token
+from cancel_token import CancelToken
 
 # Workaround for import cycles caused by type annotations:
 # http://mypy.readthedocs.io/en/latest/common_issues.html#import-cycles
@@ -504,8 +504,8 @@ async def wait_ping(self, remote: Node, cancel_token: CancelToken) -> bool:
         with self.ping_callbacks.acquire(remote, event.set):
             got_ping = False
             try:
-                got_ping = await wait_with_token(
-                    event.wait(), token=cancel_token, timeout=k_request_timeout)
+                got_ping = await cancel_token.cancellable_wait(
+                    event.wait(), timeout=k_request_timeout)
                 self.logger.debug('got expected ping from %s', remote)
             except TimeoutError:
                 self.logger.debug('timed out waiting for ping from %s', remote)
@@ -525,8 +525,8 @@ async def wait_pong(self, remote: Node, token: bytes, cancel_token: CancelToken)
         with self.pong_callbacks.acquire(pingid, event.set):
             got_pong = False
             try:
-                got_pong = await wait_with_token(
-                    event.wait(), token=cancel_token, timeout=k_request_timeout)
+                got_pong = await cancel_token.cancellable_wait(
+                    event.wait(), timeout=k_request_timeout)
                 self.logger.debug('got expected pong with token %s', encode_hex(token))
             except TimeoutError:
                 self.logger.debug(
@@ -555,8 +555,8 @@ def process(response: List[Node]) -> None:
 
         with self.neighbours_callbacks.acquire(remote, process):
             try:
-                await wait_with_token(
-                    event.wait(), token=cancel_token, timeout=k_request_timeout)
+                await cancel_token.cancellable_wait(
+                    event.wait(), timeout=k_request_timeout)
                 self.logger.debug('got expected neighbours response from %s', remote)
             except TimeoutError:
                 self.logger.debug('timed out waiting for neighbours response from %s', remote)

p2p/kademlia.py

@@ -29,6 +29,8 @@
 from trie import HexaryTrie
 from trie.exceptions import BadTrieProof
 
+from cancel_token import CancelToken
+
 from eth.exceptions import (
     BlockNotFound,
     HeaderNotFound,
@@ -42,7 +44,6 @@
     NoConnectedPeers,
     NoEligiblePeers,
 )
-from p2p.cancel_token import CancelToken
 from p2p import protocol
 from p2p.constants import (
     COMPLETION_TIMEOUT,

p2p/lightchain.py

@@ -9,7 +9,7 @@
 )
 from urllib.parse import urlparse
 
-from p2p.cancel_token import (
+from cancel_token import (
     CancelToken,
 )
 from p2p.exceptions import (

p2p/nat.py

@@ -48,6 +48,8 @@
     keys,
 )
 
+from cancel_token import CancelToken
+
 from eth.chains.mainnet import MAINNET_NETWORK_ID
 from eth.chains.ropsten import ROPSTEN_NETWORK_ID
 from eth.constants import GENESIS_BLOCK_NUMBER
@@ -75,7 +77,6 @@
     UnknownProtocolCommand,
     UnreachablePeer,
 )
-from p2p.cancel_token import CancelToken
 from p2p.service import BaseService
 from p2p.utils import (
     gen_request_id,