def __init__(
self,
local_private_key: keys.PrivateKey,
listen_on: Endpoint,
enr_db: ENRDatabaseAPI,
events: EventsAPI = None,
message_type_registry: MessageTypeRegistry = v51_registry,
) -> None:
self._local_private_key = local_private_key
self.listen_on = listen_on
self._listening = trio.Event()
self.enr_manager = ENRManager(
private_key=local_private_key,
enr_db=enr_db,
)
self.enr_db = enr_db
self._registry = message_type_registry
# Datagrams
(
self._outbound_datagram_send_channel,
self._outbound_datagram_receive_channel,
) = trio.open_memory_channel[OutboundDatagram](256)
(
self._inbound_datagram_send_channel,
self._inbound_datagram_receive_channel,
) = trio.open_memory_channel[InboundDatagram](256)
# EnvelopePair
(
self._outbound_envelope_send_channel,
self._outbound_envelope_receive_channel,
) = trio.open_memory_channel[OutboundEnvelope](256)
(
self._inbound_envelope_send_channel,
self._inbound_envelope_receive_channel,
) = trio.open_memory_channel[InboundEnvelope](256)
# Messages
(
self._outbound_message_send_channel,
self._outbound_message_receive_channel,
) = trio.open_memory_channel[AnyOutboundMessage](256)
(
self._inbound_message_send_channel,
self._inbound_message_receive_channel,
) = trio.open_memory_channel[AnyInboundMessage](256)
if events is None:
events = Events()
self.events = events
self.pool = Pool(
local_private_key=self._local_private_key,
local_node_id=self.enr_manager.enr.node_id,
enr_db=self.enr_db,
outbound_envelope_send_channel=self.
_outbound_envelope_send_channel,
inbound_message_send_channel=self._inbound_message_send_channel,
message_type_registry=self._registry,
events=self.events,
)
self.dispatcher = Dispatcher(
self._inbound_envelope_receive_channel,
self._inbound_message_receive_channel,
self.pool,
self.enr_db,
self._registry,
self.events,
)
self.envelope_decoder = EnvelopeEncoder(
self._outbound_envelope_receive_channel,
self._outbound_datagram_send_channel,
)
self.envelope_encoder = EnvelopeDecoder(
self._inbound_datagram_receive_channel,
self._inbound_envelope_send_channel,
self.enr_manager.enr.node_id,
)
self._ready = trio.Event()
async def dispatcher_pair(
self, node_a: NodeAPI, node_b: NodeAPI
) -> AsyncIterator[Tuple[DispatcherAPI, DispatcherAPI]]:
if node_a.node_id < node_b.node_id:
left = node_a
right = node_b
elif node_b.node_id < node_a.node_id:
left = node_b
right = node_a
else:
raise Exception("Cannot pair with self")
key = (left.node_id, right.node_id)
if key in self._running_dispatchers:
raise Exception("Already running dispatchers for: "
f"{humanize_node_id(left.node_id)} <-> "
f"{humanize_node_id(right.node_id)}")
self.logger.info("setting up dispatcher pair: %s <> %s", left, right)
async with AsyncExitStack() as stack:
left_pool, left_channels = self._get_or_create_pool_for_node(left)
if left.node_id in self._managed_dispatchers:
self.logger.info("dispatcher already present for %s", left)
left_managed_dispatcher = self._managed_dispatchers[
left.node_id]
else:
self.logger.info("setting up new dispatcher for %s", left)
(
left_inbound_envelope_send_channel,
left_inbound_envelope_receive_channel,
) = trio.open_memory_channel[InboundEnvelope](256)
left_dispatcher = Dispatcher(
left_inbound_envelope_receive_channel,
left_channels.inbound_message_receive_channel,
left_pool,
left.enr_db,
events=left.events,
)
left_managed_dispatcher = ManagedDispatcher(
dispatcher=left_dispatcher,
send_channel=left_inbound_envelope_send_channel,
)
self._managed_dispatchers[
left.node_id] = left_managed_dispatcher
await stack.enter_async_context(
background_trio_service(left_dispatcher))
right_pool, right_channels = self._get_or_create_pool_for_node(
right)
if right.node_id in self._managed_dispatchers:
self.logger.info("dispatcher already present for %s", right)
right_managed_dispatcher = self._managed_dispatchers[
right.node_id]
else:
self.logger.info("setting up new dispatcher for %s", right)
(
right_inbound_envelope_send_channel,
right_inbound_envelope_receive_channel,
) = trio.open_memory_channel[InboundEnvelope](256)
right_dispatcher = Dispatcher(
right_inbound_envelope_receive_channel,
right_channels.inbound_message_receive_channel,
right_pool,
right.enr_db,
events=right.events,
)
right_managed_dispatcher = ManagedDispatcher(
dispatcher=right_dispatcher,
send_channel=right_inbound_envelope_send_channel,
)
self._managed_dispatchers[
right.node_id] = right_managed_dispatcher
await stack.enter_async_context(
background_trio_service(right_dispatcher))
if left is node_a:
dispatchers = (
left_managed_dispatcher.dispatcher,
right_managed_dispatcher.dispatcher,
)
elif left is node_b:
dispatchers = (
right_managed_dispatcher.dispatcher,
left_managed_dispatcher.dispatcher,
)
else:
raise Exception("Invariant")
async with staple(
left,
left_channels.outbound_envelope_receive_channel,
right_managed_dispatcher.send_channel.clone(
), # type: ignore
):
async with staple(
right,
right_channels.outbound_envelope_receive_channel,
left_managed_dispatcher.send_channel.clone(
), # type: ignore
):
self._running_dispatchers.add(key)
try:
yield dispatchers
finally:
self._running_dispatchers.remove(key)
class Client(Service, ClientAPI):
logger = logging.getLogger("ddht.Client")
def __init__(
self,
local_private_key: keys.PrivateKey,
listen_on: Endpoint,
enr_db: ENRDatabaseAPI,
events: EventsAPI = None,
message_type_registry: MessageTypeRegistry = v51_registry,
) -> None:
self._local_private_key = local_private_key
self.listen_on = listen_on
self._listening = trio.Event()
self.enr_manager = ENRManager(
private_key=local_private_key,
enr_db=enr_db,
)
self.enr_db = enr_db
self._registry = message_type_registry
# Datagrams
(
self._outbound_datagram_send_channel,
self._outbound_datagram_receive_channel,
) = trio.open_memory_channel[OutboundDatagram](256)
(
self._inbound_datagram_send_channel,
self._inbound_datagram_receive_channel,
) = trio.open_memory_channel[InboundDatagram](256)
# EnvelopePair
(
self._outbound_envelope_send_channel,
self._outbound_envelope_receive_channel,
) = trio.open_memory_channel[OutboundEnvelope](256)
(
self._inbound_envelope_send_channel,
self._inbound_envelope_receive_channel,
) = trio.open_memory_channel[InboundEnvelope](256)
# Messages
(
self._outbound_message_send_channel,
self._outbound_message_receive_channel,
) = trio.open_memory_channel[AnyOutboundMessage](256)
(
self._inbound_message_send_channel,
self._inbound_message_receive_channel,
) = trio.open_memory_channel[AnyInboundMessage](256)
if events is None:
events = Events()
self.events = events
self.pool = Pool(
local_private_key=self._local_private_key,
local_node_id=self.enr_manager.enr.node_id,
enr_db=self.enr_db,
outbound_envelope_send_channel=self.
_outbound_envelope_send_channel,
inbound_message_send_channel=self._inbound_message_send_channel,
message_type_registry=self._registry,
events=self.events,
)
self.dispatcher = Dispatcher(
self._inbound_envelope_receive_channel,
self._inbound_message_receive_channel,
self.pool,
self.enr_db,
self._registry,
self.events,
)
self.envelope_decoder = EnvelopeEncoder(
self._outbound_envelope_receive_channel,
self._outbound_datagram_send_channel,
)
self.envelope_encoder = EnvelopeDecoder(
self._inbound_datagram_receive_channel,
self._inbound_envelope_send_channel,
self.enr_manager.enr.node_id,
)
self._ready = trio.Event()
@property
def local_node_id(self) -> NodeID:
return self.pool.local_node_id
async def run(self) -> None:
self.manager.run_daemon_child_service(self.dispatcher)
self.manager.run_daemon_child_service(self.envelope_decoder)
self.manager.run_daemon_child_service(self.envelope_encoder)
self.manager.run_daemon_task(self._do_listen, self.listen_on)
await self.manager.wait_finished()
async def wait_listening(self) -> None:
await self._listening.wait()
async def _do_listen(self, listen_on: Endpoint) -> None:
sock = trio.socket.socket(
family=trio.socket.AF_INET,
type=trio.socket.SOCK_DGRAM,
)
ip_address, port = listen_on
await sock.bind((socket.inet_ntoa(ip_address), port))
self._listening.set()
await self.events.listening.trigger(listen_on)
self.logger.debug("Network connection listening on %s", listen_on)
# TODO: the datagram services need to use the `EventsAPI`
datagram_sender = DatagramSender(
self._outbound_datagram_receive_channel,
sock) # type: ignore # noqa: E501
self.manager.run_daemon_child_service(datagram_sender)
datagram_receiver = DatagramReceiver(
sock,
self._inbound_datagram_send_channel) # type: ignore # noqa: E501
self.manager.run_daemon_child_service(datagram_receiver)
await self.manager.wait_finished()
#
# Message API
#
@contextmanager
def _get_request_id(self,
node_id: NodeID,
request_id: Optional[int] = None) -> Iterator[int]:
request_id_context: ContextManager[int]
if request_id is None:
request_id_context = self.dispatcher.reserve_request_id(node_id)
else:
request_id_context = nullcontext(request_id)
with request_id_context as message_request_id:
yield message_request_id
async def send_ping(
self,
endpoint: Endpoint,
node_id: NodeID,
*,
request_id: Optional[int] = None,
) -> int:
with self._get_request_id(node_id, request_id) as message_request_id:
message = AnyOutboundMessage(
PingMessage(message_request_id,
self.enr_manager.enr.sequence_number),
endpoint,
node_id,
)
await self.dispatcher.send_message(message)
return message_request_id
async def send_pong(
self,
endpoint: Endpoint,
node_id: NodeID,
*,
request_id: int,
) -> None:
message = AnyOutboundMessage(
PongMessage(
request_id,
self.enr_manager.enr.sequence_number,
endpoint.ip_address,
endpoint.port,
),
endpoint,
node_id,
)
await self.dispatcher.send_message(message)
async def send_find_nodes(
self,
endpoint: Endpoint,
node_id: NodeID,
*,
distances: Collection[int],
request_id: Optional[int] = None,
) -> int:
with self._get_request_id(node_id, request_id) as message_request_id:
message = AnyOutboundMessage(
FindNodeMessage(message_request_id, tuple(distances)),
endpoint,
node_id,
)
await self.dispatcher.send_message(message)
return message_request_id
async def send_found_nodes(
self,
endpoint: Endpoint,
node_id: NodeID,
*,
enrs: Sequence[ENRAPI],
request_id: int,
) -> int:
enr_batches = partition_enrs(
enrs, max_payload_size=FOUND_NODES_MAX_PAYLOAD_SIZE)
num_batches = len(enr_batches)
for batch in enr_batches:
message = AnyOutboundMessage(
FoundNodesMessage(
request_id,
num_batches,
batch,
),
endpoint,
node_id,
)
await self.dispatcher.send_message(message)
return num_batches
async def send_talk_request(
self,
endpoint: Endpoint,
node_id: NodeID,
*,
protocol: bytes,
request: bytes,
request_id: Optional[int] = None,
) -> int:
with self._get_request_id(node_id, request_id) as message_request_id:
message = AnyOutboundMessage(
TalkRequestMessage(message_request_id, protocol, request),
endpoint,
node_id,
)
await self.dispatcher.send_message(message)
return message_request_id
async def send_talk_response(
self,
endpoint: Endpoint,
node_id: NodeID,
*,
response: bytes,
request_id: int,
) -> None:
message = AnyOutboundMessage(
TalkResponseMessage(
request_id,
response,
),
endpoint,
node_id,
)
await self.dispatcher.send_message(message)
async def send_register_topic(
self,
endpoint: Endpoint,
node_id: NodeID,
*,
topic: bytes,
enr: ENRAPI,
ticket: bytes = b"",
request_id: Optional[int] = None,
) -> int:
with self._get_request_id(node_id, request_id) as message_request_id:
message = AnyOutboundMessage(
RegisterTopicMessage(message_request_id, topic, enr, ticket),
endpoint,
node_id,
)
await self.dispatcher.send_message(message)
return message_request_id
async def send_ticket(
self,
endpoint: Endpoint,
node_id: NodeID,
*,
ticket: bytes,
wait_time: int,
request_id: int,
) -> None:
message = AnyOutboundMessage(
TicketMessage(request_id, ticket, wait_time),
endpoint,
node_id,
)
await self.dispatcher.send_message(message)
async def send_registration_confirmation(
self,
endpoint: Endpoint,
node_id: NodeID,
*,
topic: bytes,
request_id: int,
) -> None:
message = AnyOutboundMessage(
RegistrationConfirmationMessage(request_id, topic),
endpoint,
node_id,
)
await self.dispatcher.send_message(message)
async def send_topic_query(
self,
endpoint: Endpoint,
node_id: NodeID,
*,
topic: bytes,
request_id: Optional[int] = None,
) -> int:
with self._get_request_id(node_id, request_id) as message_request_id:
message = AnyOutboundMessage(
TopicQueryMessage(message_request_id, topic),
endpoint,
node_id,
)
await self.dispatcher.send_message(message)
return message_request_id
#
# Request/Response API
#
async def ping(self, endpoint: Endpoint,
node_id: NodeID) -> InboundMessage[PongMessage]:
with self._get_request_id(node_id) as request_id:
request = AnyOutboundMessage(
PingMessage(request_id, self.enr_manager.enr.sequence_number),
endpoint,
node_id,
)
async with self.dispatcher.subscribe_request(
request, PongMessage) as subscription:
with trio.fail_after(REQUEST_RESPONSE_TIMEOUT):
return await subscription.receive()
async def find_nodes(
self,
endpoint: Endpoint,
node_id: NodeID,
distances: Collection[int],
) -> Tuple[InboundMessage[FoundNodesMessage], ...]:
with self._get_request_id(node_id) as request_id:
request = AnyOutboundMessage(
FindNodeMessage(request_id, tuple(distances)),
endpoint,
node_id,
)
async with self.dispatcher.subscribe_request(
request, FoundNodesMessage) as subscription:
with trio.fail_after(REQUEST_RESPONSE_TIMEOUT):
head_response = await subscription.receive()
total = head_response.message.total
responses: Tuple[InboundMessage[FoundNodesMessage], ...]
if total == 1:
responses = (head_response, )
elif total > 1:
tail_responses: List[
InboundMessage[FoundNodesMessage]] = []
for _ in range(total - 1):
tail_responses.append(await subscription.receive())
responses = (head_response, ) + tuple(tail_responses)
else:
# TODO: this code path needs to be excercised and
# probably replaced with some sort of
# `SessionTerminated` exception.
raise Exception("Invalid `total` counter in response")
# Validate that all responses are indeed at one of the
# specified distances.
for response in responses:
for enr in response.message.enrs:
if enr.node_id == node_id:
if 0 not in distances:
raise ValidationError(
f"Invalid response: distance=0 expected={distances}"
)
else:
distance = compute_log_distance(
enr.node_id, node_id)
if distance not in distances:
raise ValidationError(
f"Invalid response: distance={distance} expected={distances}"
)
return responses
async def talk(self, endpoint: Endpoint, node_id: NodeID, protocol: bytes,
request: bytes) -> InboundMessage[TalkResponseMessage]:
raise NotImplementedError
async def register_topic(
self,
endpoint: Endpoint,
node_id: NodeID,
topic: bytes,
ticket: Optional[bytes] = None,
) -> Tuple[InboundMessage[TicketMessage],
Optional[InboundMessage[RegistrationConfirmationMessage]], ]:
raise NotImplementedError
async def topic_query(self, endpoint: Endpoint, node_id: NodeID,
topic: bytes) -> InboundMessage[FoundNodesMessage]:
raise NotImplementedError
class Client(Service, ClientAPI):
logger = logging.getLogger("ddht.Client")
def __init__(
self,
local_private_key: keys.PrivateKey,
listen_on: Endpoint,
enr_db: QueryableENRDatabaseAPI,
session_cache_size: int,
events: EventsAPI = None,
message_type_registry: MessageTypeRegistry = v51_registry,
) -> None:
self.local_private_key = local_private_key
self.listen_on = listen_on
self._listening = trio.Event()
self.enr_manager = ENRManager(
private_key=local_private_key,
enr_db=enr_db,
)
self.enr_db = enr_db
self._registry = message_type_registry
self.request_tracker = RequestTracker()
# Datagrams
(
self._outbound_datagram_send_channel,
self._outbound_datagram_receive_channel,
) = trio.open_memory_channel[OutboundDatagram](256)
(
self._inbound_datagram_send_channel,
self._inbound_datagram_receive_channel,
) = trio.open_memory_channel[InboundDatagram](256)
# EnvelopePair
(
self._outbound_envelope_send_channel,
self._outbound_envelope_receive_channel,
) = trio.open_memory_channel[OutboundEnvelope](256)
(
self._inbound_envelope_send_channel,
self._inbound_envelope_receive_channel,
) = trio.open_memory_channel[InboundEnvelope](256)
# Messages
(
self._outbound_message_send_channel,
self._outbound_message_receive_channel,
) = trio.open_memory_channel[AnyOutboundMessage](256)
(
self._inbound_message_send_channel,
self._inbound_message_receive_channel,
) = trio.open_memory_channel[AnyInboundMessage](256)
if events is None:
events = Events()
self.events = events
self.pool = Pool(
local_private_key=self.local_private_key,
local_node_id=self.enr_manager.enr.node_id,
enr_db=self.enr_db,
outbound_envelope_send_channel=self.
_outbound_envelope_send_channel,
inbound_message_send_channel=self._inbound_message_send_channel,
session_cache_size=session_cache_size,
message_type_registry=self._registry,
events=self.events,
)
self.dispatcher = Dispatcher(
self._inbound_envelope_receive_channel,
self._inbound_message_receive_channel,
self.pool,
self.enr_db,
self.events,
)
self.envelope_decoder = EnvelopeEncoder(
self._outbound_envelope_receive_channel,
self._outbound_datagram_send_channel,
)
self.envelope_encoder = EnvelopeDecoder(
self._inbound_datagram_receive_channel,
self._inbound_envelope_send_channel,
self.enr_manager.enr.node_id,
)
self._ready = trio.Event()
@property
def local_node_id(self) -> NodeID:
return self.pool.local_node_id
async def run(self) -> None:
self.manager.run_daemon_task(
self._run_envelope_and_dispatcher_services)
self.manager.run_daemon_task(self._do_listen, self.listen_on)
await self.manager.wait_finished()
async def _run_envelope_and_dispatcher_services(self) -> None:
"""
Ensure that in the task hierarchy the envelope encode will be shut down
*after* the dispatcher.
run()
|
---EnvelopeEncoder
|
---EnvelopeDecoder
|
---Dispatcher
"""
async with background_trio_service(self.envelope_encoder):
async with background_trio_service(self.envelope_decoder):
async with background_trio_service(self.dispatcher):
await self.manager.wait_finished()
async def wait_listening(self) -> None:
await self._listening.wait()
async def _do_listen(self, listen_on: Endpoint) -> None:
sock = trio.socket.socket(
family=trio.socket.AF_INET,
type=trio.socket.SOCK_DGRAM,
)
ip_address, port = listen_on
await sock.bind((socket.inet_ntoa(ip_address), port))
self._listening.set()
await self.events.listening.trigger(listen_on)
self.logger.debug("Network connection listening on %s", listen_on)
# TODO: the datagram services need to use the `EventsAPI`
datagram_sender = DatagramSender(
self._outbound_datagram_receive_channel,
sock) # type: ignore # noqa: E501
self.manager.run_daemon_child_service(datagram_sender)
datagram_receiver = DatagramReceiver(
sock,
self._inbound_datagram_send_channel) # type: ignore # noqa: E501
self.manager.run_daemon_child_service(datagram_receiver)
await self.manager.wait_finished()
#
# Message API
#
async def send_ping(
self,
node_id: NodeID,
endpoint: Endpoint,
*,
enr_seq: Optional[int] = None,
request_id: Optional[bytes] = None,
) -> bytes:
if enr_seq is None:
enr_seq = self.enr_manager.enr.sequence_number
with self.request_tracker.reserve_request_id(
node_id, request_id) as message_request_id:
message = AnyOutboundMessage(
PingMessage(message_request_id, enr_seq),
endpoint,
node_id,
)
await self.dispatcher.send_message(message)
return message_request_id
async def send_pong(
self,
node_id: NodeID,
endpoint: Endpoint,
*,
enr_seq: Optional[int] = None,
request_id: bytes,
) -> None:
if enr_seq is None:
enr_seq = self.enr_manager.enr.sequence_number
message = AnyOutboundMessage(
PongMessage(
request_id,
enr_seq,
endpoint.ip_address,
endpoint.port,
),
endpoint,
node_id,
)
await self.dispatcher.send_message(message)
async def send_find_nodes(
self,
node_id: NodeID,
endpoint: Endpoint,
*,
distances: Collection[int],
request_id: Optional[bytes] = None,
) -> bytes:
with self.request_tracker.reserve_request_id(
node_id, request_id) as message_request_id:
message = AnyOutboundMessage(
FindNodeMessage(message_request_id, tuple(distances)),
endpoint,
node_id,
)
await self.dispatcher.send_message(message)
return message_request_id
async def send_found_nodes(
self,
node_id: NodeID,
endpoint: Endpoint,
*,
enrs: Sequence[ENRAPI],
request_id: bytes,
) -> int:
enr_batches = partition_enrs(
enrs, max_payload_size=FOUND_NODES_MAX_PAYLOAD_SIZE)
num_batches = len(enr_batches)
for batch in enr_batches:
message = AnyOutboundMessage(
FoundNodesMessage(
request_id,
num_batches,
batch,
),
endpoint,
node_id,
)
await self.dispatcher.send_message(message)
return num_batches
async def send_talk_request(
self,
node_id: NodeID,
endpoint: Endpoint,
*,
protocol: bytes,
payload: bytes,
request_id: Optional[bytes] = None,
) -> bytes:
with self.request_tracker.reserve_request_id(
node_id, request_id) as message_request_id:
message = AnyOutboundMessage(
TalkRequestMessage(message_request_id, protocol, payload),
endpoint,
node_id,
)
await self.dispatcher.send_message(message)
return message_request_id
async def send_talk_response(
self,
node_id: NodeID,
endpoint: Endpoint,
*,
payload: bytes,
request_id: bytes,
) -> None:
message = AnyOutboundMessage(
TalkResponseMessage(request_id, payload),
endpoint,
node_id,
)
await self.dispatcher.send_message(message)
async def send_register_topic(
self,
node_id: NodeID,
endpoint: Endpoint,
*,
topic: bytes,
enr: ENRAPI,
ticket: bytes = b"",
request_id: Optional[bytes] = None,
) -> bytes:
with self.request_tracker.reserve_request_id(
node_id, request_id) as message_request_id:
message = AnyOutboundMessage(
RegisterTopicMessage(message_request_id, topic, enr, ticket),
endpoint,
node_id,
)
await self.dispatcher.send_message(message)
return message_request_id
async def send_ticket(
self,
node_id: NodeID,
endpoint: Endpoint,
*,
ticket: bytes,
wait_time: int,
request_id: bytes,
) -> None:
message = AnyOutboundMessage(
TicketMessage(request_id, ticket, wait_time),
endpoint,
node_id,
)
await self.dispatcher.send_message(message)
async def send_registration_confirmation(
self,
node_id: NodeID,
endpoint: Endpoint,
*,
topic: bytes,
request_id: bytes,
) -> None:
message = AnyOutboundMessage(
RegistrationConfirmationMessage(request_id, topic),
endpoint,
node_id,
)
await self.dispatcher.send_message(message)
async def send_topic_query(
self,
node_id: NodeID,
endpoint: Endpoint,
*,
topic: bytes,
request_id: Optional[bytes] = None,
) -> bytes:
with self.request_tracker.reserve_request_id(
node_id, request_id) as message_request_id:
message = AnyOutboundMessage(
TopicQueryMessage(message_request_id, topic),
endpoint,
node_id,
)
await self.dispatcher.send_message(message)
return message_request_id
#
# Request/Response API
#
async def ping(
self,
node_id: NodeID,
endpoint: Endpoint,
*,
request_id: Optional[bytes] = None,
) -> InboundMessage[PongMessage]:
with self.request_tracker.reserve_request_id(node_id,
request_id) as request_id:
request = AnyOutboundMessage(
PingMessage(request_id, self.enr_manager.enr.sequence_number),
endpoint,
node_id,
)
async with self.dispatcher.subscribe_request(
request, PongMessage) as subscription:
return await subscription.receive()
async def find_nodes(
self,
node_id: NodeID,
endpoint: Endpoint,
distances: Collection[int],
*,
request_id: Optional[bytes] = None,
) -> Tuple[InboundMessage[FoundNodesMessage], ...]:
with self.request_tracker.reserve_request_id(node_id,
request_id) as request_id:
request = AnyOutboundMessage(
FindNodeMessage(request_id, tuple(distances)),
endpoint,
node_id,
)
async with self.dispatcher.subscribe_request(
request, FoundNodesMessage) as subscription:
head_response = await subscription.receive()
total = head_response.message.total
responses: Tuple[InboundMessage[FoundNodesMessage], ...]
if total == 1:
responses = (head_response, )
elif total > 1:
tail_responses: List[
InboundMessage[FoundNodesMessage]] = []
for _ in range(total - 1):
tail_responses.append(await subscription.receive())
responses = (head_response, ) + tuple(tail_responses)
else:
raise ValidationError(
f"Invalid `total` counter in response: total={total}")
return responses
def stream_find_nodes(
self,
node_id: NodeID,
endpoint: Endpoint,
distances: Collection[int],
*,
request_id: Optional[bytes] = None,
) -> AsyncContextManager[trio.abc.ReceiveChannel[
InboundMessage[FoundNodesMessage]]]:
return common_client_stream_find_nodes(self,
node_id,
endpoint,
distances,
request_id=request_id)
async def talk(
self,
node_id: NodeID,
endpoint: Endpoint,
protocol: bytes,
payload: bytes,
*,
request_id: Optional[bytes] = None,
) -> InboundMessage[TalkResponseMessage]:
with self.request_tracker.reserve_request_id(node_id,
request_id) as request_id:
request = AnyOutboundMessage(
TalkRequestMessage(request_id, protocol, payload),
endpoint,
node_id,
)
async with self.dispatcher.subscribe_request(
request, TalkResponseMessage) as subscription:
return await subscription.receive()
async def register_topic(
self,
node_id: NodeID,
endpoint: Endpoint,
topic: bytes,
ticket: Optional[bytes] = None,
*,
request_id: Optional[bytes] = None,
) -> Tuple[InboundMessage[TicketMessage],
Optional[InboundMessage[RegistrationConfirmationMessage]], ]:
raise NotImplementedError
async def topic_query(
self,
node_id: NodeID,
endpoint: Endpoint,
topic: bytes,
*,
request_id: Optional[bytes] = None,
) -> InboundMessage[FoundNodesMessage]:
raise NotImplementedError