async def test_call_unary_handler(should_cancel,
replicate,
handle_name="handle"):
handler_cancelled = False
async def ping_handler(request, context):
try:
await asyncio.sleep(2)
except asyncio.CancelledError:
nonlocal handler_cancelled
handler_cancelled = True
return dht_pb2.PingResponse(peer=dht_pb2.NodeInfo(
node_id=context.id.encode(), rpc_port=context.port),
sender_endpoint=context.handle_name,
available=True)
server_primary = await P2P.create()
server = await replicate_if_needed(server_primary, replicate)
server_pid = server_primary._child.pid
await server.add_unary_handler(handle_name, ping_handler,
dht_pb2.PingRequest, dht_pb2.PingResponse)
assert is_process_running(server_pid)
nodes = bootstrap_from([server])
client_primary = await P2P.create(bootstrap=True, bootstrap_peers=nodes)
client = await replicate_if_needed(client_primary, replicate)
client_pid = client_primary._child.pid
assert is_process_running(client_pid)
ping_request = dht_pb2.PingRequest(peer=dht_pb2.NodeInfo(
node_id=client.id.encode(), rpc_port=client._host_port),
validate=True)
expected_response = dht_pb2.PingResponse(peer=dht_pb2.NodeInfo(
node_id=server.id.encode(), rpc_port=server._host_port),
sender_endpoint=handle_name,
available=True)
await client.wait_for_at_least_n_peers(1)
libp2p_server_id = PeerID.from_base58(server.id)
stream_info, reader, writer = await client._client.stream_open(
libp2p_server_id, (handle_name, ))
await P2P.send_protobuf(ping_request, dht_pb2.PingRequest, writer)
if should_cancel:
writer.close()
await asyncio.sleep(1)
assert handler_cancelled
else:
result, err = await P2P.receive_protobuf(dht_pb2.PingResponse, reader)
assert err is None
assert result == expected_response
assert not handler_cancelled
await server.stop_listening()
await server_primary.shutdown()
assert not is_process_running(server_pid)
await client_primary.shutdown()
assert not is_process_running(client_pid)
async def create(cls,
node_id: DHTID,
bucket_size: int,
depth_modulo: int,
num_replicas: int,
wait_timeout: float,
parallel_rpc: Optional[int] = None,
cache_size: Optional[int] = None,
listen=True,
listen_on='0.0.0.0:*',
endpoint: Optional[Endpoint] = None,
channel_options: Sequence[Tuple[str, Any]] = (),
**kwargs) -> DHTProtocol:
"""
A protocol that allows DHT nodes to request keys/neighbors from other DHT nodes.
As a side-effect, DHTProtocol also maintains a routing table as described in
https://pdos.csail.mit.edu/~petar/papers/maymounkov-kademlia-lncs.pdf
See DHTNode (node.py) for a more detailed description.
:note: the rpc_* methods defined in this class will be automatically exposed to other DHT nodes,
for instance, def rpc_ping can be called as protocol.call_ping(endpoint, dht_id) from a remote machine
Only the call_* methods are meant to be called publicly, e.g. from DHTNode
Read more: https://github.com/bmuller/rpcudp/tree/master/rpcudp
"""
self = cls(_initialized_with_create=True)
self.node_id, self.bucket_size, self.num_replicas = node_id, bucket_size, num_replicas
self.wait_timeout, self.channel_options = wait_timeout, tuple(
channel_options)
self.storage, self.cache = DHTLocalStorage(), DHTLocalStorage(
maxsize=cache_size)
self.routing_table = RoutingTable(node_id, bucket_size, depth_modulo)
self.rpc_semaphore = asyncio.Semaphore(
parallel_rpc if parallel_rpc is not None else float('inf'))
if listen: # set up server to process incoming rpc requests
grpc.aio.init_grpc_aio()
self.server = grpc.aio.server(**kwargs,
options=GRPC_KEEPALIVE_OPTIONS)
dht_grpc.add_DHTServicer_to_server(self, self.server)
self.port = self.server.add_insecure_port(listen_on)
assert self.port != 0, f"Failed to listen to {listen_on}"
if endpoint is not None and endpoint.endswith('*'):
endpoint = replace_port(endpoint, self.port)
self.node_info = dht_pb2.NodeInfo(
node_id=node_id.to_bytes(),
rpc_port=self.port,
endpoint=endpoint
or dht_pb2.NodeInfo.endpoint.DESCRIPTOR.default_value)
await self.server.start()
else: # not listening to incoming requests, client-only mode
# note: use empty node_info so peers won't add you to their routing tables
self.node_info, self.server, self.port = dht_pb2.NodeInfo(
), None, None
if listen_on != '0.0.0.0:*' or len(kwargs) != 0:
logger.warning(
f"DHTProtocol has no server (due to listen=False), listen_on"
f"and kwargs have no effect (unused kwargs: {kwargs})")
return self
async def test_call_unary_handler_error(handle_name="handle"):
async def error_handler(request, context):
raise ValueError('boom')
server = await P2P.create()
server_pid = server._child.pid
await server.add_unary_handler(handle_name, error_handler,
dht_pb2.PingRequest, dht_pb2.PingResponse)
assert is_process_running(server_pid)
nodes = bootstrap_from([server])
client = await P2P.create(bootstrap=True, bootstrap_peers=nodes)
client_pid = client._child.pid
assert is_process_running(client_pid)
await client.wait_for_at_least_n_peers(1)
ping_request = dht_pb2.PingRequest(peer=dht_pb2.NodeInfo(
node_id=client.id.encode(), rpc_port=client._host_port),
validate=True)
libp2p_server_id = PeerID.from_base58(server.id)
stream_info, reader, writer = await client._client.stream_open(
libp2p_server_id, (handle_name, ))
await P2P.send_protobuf(ping_request, dht_pb2.PingRequest, writer)
result, err = await P2P.receive_protobuf(dht_pb2.PingResponse, reader)
assert result is None
assert err.message == 'boom'
await server.stop_listening()
await server.shutdown()
await client.shutdown()
async def ping_handler(request, context):
try:
await asyncio.sleep(2)
except asyncio.CancelledError:
nonlocal handler_cancelled
handler_cancelled = True
return dht_pb2.PingResponse(peer=dht_pb2.NodeInfo(
node_id=context.id.encode(), rpc_port=context.port),
sender_endpoint=context.handle_name,
available=True)