From e397ce25a61777bc9d27732ccddd8e85e92b8f5b Mon Sep 17 00:00:00 2001 From: kaneki003 Date: Sat, 31 May 2025 19:02:18 +0530 Subject: [PATCH 01/23] Updated Yamux impl.,added tests for yamux and mplex --- libp2p/stream_muxer/mplex/mplex_stream.py | 101 +++++---- libp2p/stream_muxer/yamux/yamux.py | 92 ++++++-- newsfragments/639.feature.rst | 1 + .../test_mplex_read_write_lock.py | 124 +++++++++++ .../test_yamux_read_write_lock.py | 207 ++++++++++++++++++ 5 files changed, 458 insertions(+), 67 deletions(-) create mode 100644 newsfragments/639.feature.rst create mode 100644 tests/core/stream_muxer/test_mplex_read_write_lock.py create mode 100644 tests/core/stream_muxer/test_yamux_read_write_lock.py diff --git a/libp2p/stream_muxer/mplex/mplex_stream.py b/libp2p/stream_muxer/mplex/mplex_stream.py index 3b640df1..dc65ac5f 100644 --- a/libp2p/stream_muxer/mplex/mplex_stream.py +++ b/libp2p/stream_muxer/mplex/mplex_stream.py @@ -46,8 +46,9 @@ class MplexStream(IMuxedStream): read_deadline: int | None write_deadline: int | None - # TODO: Add lock for read/write to avoid interleaving receiving messages? close_lock: trio.Lock + read_lock: trio.Lock + write_lock: trio.Lock # NOTE: `dataIn` is size of 8 in Go implementation. incoming_data_channel: "trio.MemoryReceiveChannel[bytes]" @@ -80,6 +81,8 @@ class MplexStream(IMuxedStream): self.event_remote_closed = trio.Event() self.event_reset = trio.Event() self.close_lock = trio.Lock() + self.read_lock = trio.Lock() + self.write_lock = trio.Lock() self.incoming_data_channel = incoming_data_channel self._buf = bytearray() @@ -113,48 +116,49 @@ class MplexStream(IMuxedStream): :param n: number of bytes to read :return: bytes actually read """ - if n is not None and n < 0: - raise ValueError( - "the number of bytes to read `n` must be non-negative or " - f"`None` to indicate read until EOF, got n={n}" - ) - if self.event_reset.is_set(): - raise MplexStreamReset - if n is None: - return await self._read_until_eof() - if len(self._buf) == 0: - data: bytes - # Peek whether there is data available. If yes, we just read until there is - # no data, then return. - try: - data = self.incoming_data_channel.receive_nowait() - self._buf.extend(data) - except trio.EndOfChannel: - raise MplexStreamEOF - except trio.WouldBlock: - # We know `receive` will be blocked here. Wait for data here with - # `receive` and catch all kinds of errors here. + async with self.read_lock: + if n is not None and n < 0: + raise ValueError( + "the number of bytes to read `n` must be non-negative or " + f"`None` to indicate read until EOF, got n={n}" + ) + if self.event_reset.is_set(): + raise MplexStreamReset + if n is None: + return await self._read_until_eof() + if len(self._buf) == 0: + data: bytes + # Peek whether there is data available. If yes, we just read until + # there is no data, then return. try: - data = await self.incoming_data_channel.receive() + data = self.incoming_data_channel.receive_nowait() self._buf.extend(data) except trio.EndOfChannel: - if self.event_reset.is_set(): - raise MplexStreamReset - if self.event_remote_closed.is_set(): - raise MplexStreamEOF - except trio.ClosedResourceError as error: - # Probably `incoming_data_channel` is closed in `reset` when we are - # waiting for `receive`. - if self.event_reset.is_set(): - raise MplexStreamReset - raise Exception( - "`incoming_data_channel` is closed but stream is not reset. " - "This should never happen." - ) from error - self._buf.extend(self._read_return_when_blocked()) - payload = self._buf[:n] - self._buf = self._buf[len(payload) :] - return bytes(payload) + raise MplexStreamEOF + except trio.WouldBlock: + # We know `receive` will be blocked here. Wait for data here with + # `receive` and catch all kinds of errors here. + try: + data = await self.incoming_data_channel.receive() + self._buf.extend(data) + except trio.EndOfChannel: + if self.event_reset.is_set(): + raise MplexStreamReset + if self.event_remote_closed.is_set(): + raise MplexStreamEOF + except trio.ClosedResourceError as error: + # Probably `incoming_data_channel` is closed in `reset` when + # we are waiting for `receive`. + if self.event_reset.is_set(): + raise MplexStreamReset + raise Exception( + "`incoming_data_channel` is closed but stream is not reset." + "This should never happen." + ) from error + self._buf.extend(self._read_return_when_blocked()) + payload = self._buf[:n] + self._buf = self._buf[len(payload) :] + return bytes(payload) async def write(self, data: bytes) -> None: """ @@ -162,14 +166,15 @@ class MplexStream(IMuxedStream): :return: number of bytes written """ - if self.event_local_closed.is_set(): - raise MplexStreamClosed(f"cannot write to closed stream: data={data!r}") - flag = ( - HeaderTags.MessageInitiator - if self.is_initiator - else HeaderTags.MessageReceiver - ) - await self.muxed_conn.send_message(flag, data, self.stream_id) + async with self.write_lock: + if self.event_local_closed.is_set(): + raise MplexStreamClosed(f"cannot write to closed stream: data={data!r}") + flag = ( + HeaderTags.MessageInitiator + if self.is_initiator + else HeaderTags.MessageReceiver + ) + await self.muxed_conn.send_message(flag, data, self.stream_id) async def close(self) -> None: """ diff --git a/libp2p/stream_muxer/yamux/yamux.py b/libp2p/stream_muxer/yamux/yamux.py index 92123465..faf24b29 100644 --- a/libp2p/stream_muxer/yamux/yamux.py +++ b/libp2p/stream_muxer/yamux/yamux.py @@ -77,6 +77,8 @@ class YamuxStream(IMuxedStream): self.send_window = DEFAULT_WINDOW_SIZE self.recv_window = DEFAULT_WINDOW_SIZE self.window_lock = trio.Lock() + self.read_lock = trio.Lock() + self.write_lock = trio.Lock() async def __aenter__(self) -> "YamuxStream": """Enter the async context manager.""" @@ -98,16 +100,32 @@ class YamuxStream(IMuxedStream): # Flow control: Check if we have enough send window total_len = len(data) sent = 0 - + logging.debug(f"Stream {self.stream_id}: Starts writing {total_len} bytes ") while sent < total_len: + # Wait for available window with timeout + timeout = False async with self.window_lock: - # Wait for available window - while self.send_window == 0 and not self.closed: - # Release lock while waiting + if self.send_window == 0: + logging.debug( + f"Stream {self.stream_id}: Window is zero, waiting for update" + ) + # Release lock and wait with timeout self.window_lock.release() - await trio.sleep(0.01) + # To avoid re-acquiring the lock immediately, + with trio.move_on_after(5.0) as cancel_scope: + while self.send_window == 0 and not self.closed: + await trio.sleep(0.01) + # If we timed out, cancel the scope + timeout = cancel_scope.cancelled_caught + # Re-acquire lock await self.window_lock.acquire() + # If we timed out waiting for window update, raise an error + if timeout: + raise MuxedStreamError( + "Timed out waiting for window update after 5 seconds." + ) + if self.closed: raise MuxedStreamError("Stream is closed") @@ -123,25 +141,53 @@ class YamuxStream(IMuxedStream): await self.conn.secured_conn.write(header + chunk) sent += to_send - # If window is getting low, consider updating - if self.send_window < DEFAULT_WINDOW_SIZE // 2: - await self.send_window_update() + async def send_window_update( + self, increment: int | None, skip_lock: bool = False + ) -> None: + """ + Send a window update to peer. - async def send_window_update(self, increment: int | None = None) -> None: - """Send a window update to peer.""" + param:increment: The amount to increment the window size by. + If None, uses the difference between DEFAULT_WINDOW_SIZE + and current receive window. + param:skip_lock (bool): If True, skips acquiring window_lock. + This should only be used when calling from a context + that already holds the lock. + """ + increment_value = 0 if increment is None: - increment = DEFAULT_WINDOW_SIZE - self.recv_window - - if increment <= 0: + increment_value = DEFAULT_WINDOW_SIZE - self.recv_window + else: + increment_value = increment + if increment_value <= 0: + # If increment is zero or negative, skip sending update + logging.debug( + f"Stream {self.stream_id}: Skipping window update" + f"(increment={increment})" + ) return + logging.debug( + f"Stream {self.stream_id}: Sending window update with increment={increment}" + ) - async with self.window_lock: - self.recv_window += increment + async def _do_window_update() -> None: + self.recv_window += increment_value header = struct.pack( - YAMUX_HEADER_FORMAT, 0, TYPE_WINDOW_UPDATE, 0, self.stream_id, increment + YAMUX_HEADER_FORMAT, + 0, + TYPE_WINDOW_UPDATE, + 0, + self.stream_id, + increment_value, ) await self.conn.secured_conn.write(header) + if skip_lock: + await _do_window_update() + else: + async with self.window_lock: + await _do_window_update() + async def read(self, n: int | None = -1) -> bytes: # Handle None value for n by converting it to -1 if n is None: @@ -198,11 +244,19 @@ class YamuxStream(IMuxedStream): # Return all buffered data data = bytes(buffer) buffer.clear() - logging.debug(f"Stream {self.stream_id}: Returning {len(data)} bytes") return data - # For specific size read (n > 0), return available data immediately - return await self.conn.read_stream(self.stream_id, n) + data = await self.conn.read_stream(self.stream_id, n) + async with self.window_lock: + self.recv_window -= len(data) + # Automatically send a window update if recv_window is low + if self.recv_window <= DEFAULT_WINDOW_SIZE // 2: + logging.debug( + f"Stream {self.stream_id}: " + f"Low recv_window ({self.recv_window}), sending update" + ) + await self.send_window_update(None, skip_lock=True) + return data async def close(self) -> None: if not self.send_closed: diff --git a/newsfragments/639.feature.rst b/newsfragments/639.feature.rst new file mode 100644 index 00000000..3f3d7510 --- /dev/null +++ b/newsfragments/639.feature.rst @@ -0,0 +1 @@ +Added separate read and write locks to the `MplexStream` & `YamuxStream` class.This ensures thread-safe access and data integrity when multiple coroutines interact with the same MplexStream instance. diff --git a/tests/core/stream_muxer/test_mplex_read_write_lock.py b/tests/core/stream_muxer/test_mplex_read_write_lock.py new file mode 100644 index 00000000..afc197ac --- /dev/null +++ b/tests/core/stream_muxer/test_mplex_read_write_lock.py @@ -0,0 +1,124 @@ +import pytest +import trio + +from libp2p.abc import ISecureConn +from libp2p.crypto.keys import PrivateKey, PublicKey +from libp2p.peer.id import ID +from libp2p.stream_muxer.mplex.constants import ( + HeaderTags, +) +from libp2p.stream_muxer.mplex.datastructures import ( + StreamID, +) +from libp2p.stream_muxer.mplex.mplex import ( + Mplex, +) +from libp2p.stream_muxer.mplex.mplex_stream import ( + MplexStream, +) + + +class DummySecureConn(ISecureConn): + """A minimal implementation of ISecureConn for testing.""" + + async def write(self, data: bytes) -> None: + pass + + async def read(self, n: int | None = -1) -> bytes: + return b"" + + async def close(self) -> None: + pass + + def get_remote_address(self) -> tuple[str, int] | None: + return None + + def get_local_peer(self) -> ID: + return ID(b"local") + + def get_local_private_key(self) -> PrivateKey: + return PrivateKey() # Dummy key for testing + + def get_remote_peer(self) -> ID: + return ID(b"remote") + + def get_remote_public_key(self) -> PublicKey: + return PublicKey() # Dummy key for testing + + +class DummyMuxedConn(Mplex): + """A minimal mock of Mplex for testing read/write locks.""" + + def __init__(self) -> None: + self.secured_conn = DummySecureConn() + self.peer_id = ID(b"dummy") + self.streams = {} + self.streams_lock = trio.Lock() + self.event_shutting_down = trio.Event() + self.event_closed = trio.Event() + self.event_started = trio.Event() + self.stream_backlog_limit = 256 + self.stream_backlog_semaphore = trio.Semaphore(256) + channels = trio.open_memory_channel[MplexStream](0) + self.new_stream_send_channel, self.new_stream_receive_channel = channels + + async def send_message( + self, flag: HeaderTags, data: bytes, stream_id: StreamID + ) -> None: + await trio.sleep(0.01) + + +@pytest.mark.trio +async def test_concurrent_writes_are_serialized(): + stream_id = StreamID(1, True) + send_log = [] + + class LoggingMuxedConn(DummyMuxedConn): + async def send_message( + self, flag: HeaderTags, data: bytes, stream_id: StreamID + ) -> None: + send_log.append(data) + await trio.sleep(0.01) + + memory_send, memory_recv = trio.open_memory_channel(8) + stream = MplexStream( + name="test", + stream_id=stream_id, + muxed_conn=LoggingMuxedConn(), + incoming_data_channel=memory_recv, + ) + + async def writer(data): + await stream.write(data) + + async with trio.open_nursery() as nursery: + for i in range(5): + nursery.start_soon(writer, f"msg-{i}".encode()) + # Order doesn't matter due to concurrent execution + assert sorted(send_log) == sorted([f"msg-{i}".encode() for i in range(5)]) + + +@pytest.mark.trio +async def test_concurrent_reads_are_serialized(): + stream_id = StreamID(2, True) + muxed_conn = DummyMuxedConn() + memory_send, memory_recv = trio.open_memory_channel(8) + results = [] + stream = MplexStream( + name="test", + stream_id=stream_id, + muxed_conn=muxed_conn, + incoming_data_channel=memory_recv, + ) + for i in range(5): + await memory_send.send(f"data-{i}".encode()) + await memory_send.aclose() + + async def reader(): + data = await stream.read(6) + results.append(data) + + async with trio.open_nursery() as nursery: + for _ in range(5): + nursery.start_soon(reader) + assert sorted(results) == [f"data-{i}".encode() for i in range(5)] diff --git a/tests/core/stream_muxer/test_yamux_read_write_lock.py b/tests/core/stream_muxer/test_yamux_read_write_lock.py new file mode 100644 index 00000000..74a4c2f1 --- /dev/null +++ b/tests/core/stream_muxer/test_yamux_read_write_lock.py @@ -0,0 +1,207 @@ +import logging + +import pytest +import trio +from trio.testing import ( + memory_stream_pair, +) + +from libp2p.abc import IRawConnection +from libp2p.crypto.ed25519 import ( + create_new_key_pair, +) +from libp2p.peer.id import ( + ID, +) +from libp2p.security.insecure.transport import ( + InsecureTransport, +) +from libp2p.stream_muxer.yamux.yamux import ( + Yamux, + YamuxStream, +) + + +class TrioStreamAdapter(IRawConnection): + """Adapter to make trio memory streams work with libp2p.""" + + def __init__(self, send_stream, receive_stream, is_initiator=False): + self.send_stream = send_stream + self.receive_stream = receive_stream + self.is_initiator = is_initiator + + async def write(self, data: bytes) -> None: + logging.debug(f"Attempting to write {len(data)} bytes") + with trio.move_on_after(2): + await self.send_stream.send_all(data) + + async def read(self, n: int | None = None) -> bytes: + if n is None or n <= 0: + raise ValueError("Reading unbounded or zero bytes not supported") + logging.debug(f"Attempting to read {n} bytes") + with trio.move_on_after(2): + data = await self.receive_stream.receive_some(n) + logging.debug(f"Read {len(data)} bytes") + return data + + async def close(self) -> None: + logging.debug("Closing stream") + await self.send_stream.aclose() + await self.receive_stream.aclose() + + def get_remote_address(self) -> tuple[str, int] | None: + """Return None since this is a test adapter without real network info.""" + return None + + +@pytest.fixture +def key_pair(): + return create_new_key_pair() + + +@pytest.fixture +def peer_id(key_pair): + return ID.from_pubkey(key_pair.public_key) + + +@pytest.fixture +async def secure_conn_pair(key_pair, peer_id): + """Create a pair of secure connections for testing.""" + logging.debug("Setting up secure_conn_pair") + client_send, server_receive = memory_stream_pair() + server_send, client_receive = memory_stream_pair() + + client_rw = TrioStreamAdapter(client_send, client_receive) + server_rw = TrioStreamAdapter(server_send, server_receive) + + insecure_transport = InsecureTransport(key_pair) + + async def run_outbound(nursery_results): + with trio.move_on_after(5): + client_conn = await insecure_transport.secure_outbound(client_rw, peer_id) + logging.debug("Outbound handshake complete") + nursery_results["client"] = client_conn + + async def run_inbound(nursery_results): + with trio.move_on_after(5): + server_conn = await insecure_transport.secure_inbound(server_rw) + logging.debug("Inbound handshake complete") + nursery_results["server"] = server_conn + + nursery_results = {} + async with trio.open_nursery() as nursery: + nursery.start_soon(run_outbound, nursery_results) + nursery.start_soon(run_inbound, nursery_results) + await trio.sleep(0.1) # Give tasks a chance to finish + + client_conn = nursery_results.get("client") + server_conn = nursery_results.get("server") + + if client_conn is None or server_conn is None: + raise RuntimeError("Handshake failed: client_conn or server_conn is None") + + logging.debug("secure_conn_pair setup complete") + return client_conn, server_conn + + +@pytest.fixture +async def yamux_pair(secure_conn_pair, peer_id): + """Create a pair of Yamux multiplexers for testing.""" + logging.debug("Setting up yamux_pair") + client_conn, server_conn = secure_conn_pair + client_yamux = Yamux(client_conn, peer_id, is_initiator=True) + server_yamux = Yamux(server_conn, peer_id, is_initiator=False) + async with trio.open_nursery() as nursery: + with trio.move_on_after(5): + nursery.start_soon(client_yamux.start) + nursery.start_soon(server_yamux.start) + await trio.sleep(0.1) + logging.debug("yamux_pair started") + yield client_yamux, server_yamux + logging.debug("yamux_pair cleanup") + + +@pytest.mark.trio +async def test_yamux_race_condition_without_locks(yamux_pair): + """ + Test for race-around/interleaving in Yamux streams when read/write + locks are disabled. + This launches concurrent writers/readers on both sides of a stream. + If there is no proper locking, the received data may be interleaved + or corrupted. + + The test creates structured messages and verifies they are received + intact and in order. + Without proper locking, concurrent read/write operations could cause + data corruption + or message interleaving, which this test will catch. + """ + client_yamux, server_yamux = yamux_pair + client_stream: YamuxStream = await client_yamux.open_stream() + server_stream: YamuxStream = await server_yamux.accept_stream() + MSG_COUNT = 10 + MSG_SIZE = 256 * 1024 + client_msgs = [ + f"CLIENT-MSG-{i:03d}-".encode().ljust(MSG_SIZE, b"C") for i in range(MSG_COUNT) + ] + server_msgs = [ + f"SERVER-MSG-{i:03d}-".encode().ljust(MSG_SIZE, b"S") for i in range(MSG_COUNT) + ] + client_received = [] + server_received = [] + + async def writer(stream, msgs, name): + """Write messages with minimal delays to encourage race conditions.""" + for i, msg in enumerate(msgs): + await stream.write(msg) + # Yield control frequently to encourage interleaving + if i % 5 == 0: + await trio.sleep(0.005) + + async def reader(stream, received, name): + """Read messages and store them for verification.""" + for i in range(MSG_COUNT): + data = await stream.read(MSG_SIZE) + received.append(data) + if i % 3 == 0: + await trio.sleep(0.001) + + # Running all operations concurrently + async with trio.open_nursery() as nursery: + nursery.start_soon(writer, client_stream, client_msgs, "client") + nursery.start_soon(writer, server_stream, server_msgs, "server") + nursery.start_soon(reader, client_stream, client_received, "client") + nursery.start_soon(reader, server_stream, server_received, "server") + + assert len(client_received) == MSG_COUNT, ( + f"Client received {len(client_received)} messages, expected {MSG_COUNT}" + ) + assert len(server_received) == MSG_COUNT, ( + f"Server received {len(server_received)} messages, expected {MSG_COUNT}" + ) + assert client_received == server_msgs, ( + "Client did not receive server messages in order or intact!" + ) + assert server_received == client_msgs, ( + "Server did not receive client messages in order or intact!" + ) + for i, msg in enumerate(client_received): + # logging.debug(f"datatype of msg: {type(msg)}, length: {len(msg)}") + # logging.debug(f"datatype of msg: {type(b"SERVER-MSG-")}") + assert len(msg) == MSG_SIZE, ( + f"Client message {i} has wrong size: {len(msg)} != {MSG_SIZE}" + ) + assert msg.startswith(b"SERVER-MSG-"), ( + f"Client message {i} doesn't start with expected prefix" + ) + + for i, msg in enumerate(server_received): + assert len(msg) == MSG_SIZE, ( + f"Server message {i} has wrong size: {len(msg)} != {MSG_SIZE}" + ) + assert msg.startswith(b"CLIENT-MSG-"), ( + f"Server message {i} doesn't start with expected prefix" + ) + + await client_stream.close() + await server_stream.close() From 209deffc8aaf438caea7f18fd0f4615d90044cbf Mon Sep 17 00:00:00 2001 From: kaneki003 Date: Sat, 21 Jun 2025 13:39:03 +0530 Subject: [PATCH 02/23] resolved recv_window updates,added support for read_EOF --- libp2p/stream_muxer/yamux/yamux.py | 105 +++++++----------- .../test_mplex_read_write_lock.py | 15 ++- 2 files changed, 52 insertions(+), 68 deletions(-) diff --git a/libp2p/stream_muxer/yamux/yamux.py b/libp2p/stream_muxer/yamux/yamux.py index faf24b29..f58e98c4 100644 --- a/libp2p/stream_muxer/yamux/yamux.py +++ b/libp2p/stream_muxer/yamux/yamux.py @@ -141,9 +141,7 @@ class YamuxStream(IMuxedStream): await self.conn.secured_conn.write(header + chunk) sent += to_send - async def send_window_update( - self, increment: int | None, skip_lock: bool = False - ) -> None: + async def send_window_update(self, increment: int, skip_lock: bool = False) -> None: """ Send a window update to peer. @@ -154,12 +152,7 @@ class YamuxStream(IMuxedStream): This should only be used when calling from a context that already holds the lock. """ - increment_value = 0 - if increment is None: - increment_value = DEFAULT_WINDOW_SIZE - self.recv_window - else: - increment_value = increment - if increment_value <= 0: + if increment <= 0: # If increment is zero or negative, skip sending update logging.debug( f"Stream {self.stream_id}: Skipping window update" @@ -171,14 +164,13 @@ class YamuxStream(IMuxedStream): ) async def _do_window_update() -> None: - self.recv_window += increment_value header = struct.pack( YAMUX_HEADER_FORMAT, 0, TYPE_WINDOW_UPDATE, 0, self.stream_id, - increment_value, + increment, ) await self.conn.secured_conn.write(header) @@ -188,6 +180,22 @@ class YamuxStream(IMuxedStream): async with self.window_lock: await _do_window_update() + async def read_EOF(self) -> bytes: + """ + To read data from stream until it is closed. + """ + data = b"" + try: + while True: + recv = await self.read() + if recv: + data += recv + except MuxedStreamEOF: + logging.debug( + f"Stream {self.stream_id}:EOF reached,total data read:{len(data)} bytes" + ) + return data + async def read(self, n: int | None = -1) -> bytes: # Handle None value for n by converting it to -1 if n is None: @@ -202,61 +210,34 @@ class YamuxStream(IMuxedStream): # If reading until EOF (n == -1), block until stream is closed if n == -1: - while not self.recv_closed and not self.conn.event_shutting_down.is_set(): - # Check if there's data in the buffer - buffer = self.conn.stream_buffers.get(self.stream_id) - if buffer and len(buffer) > 0: - # Wait for closure even if data is available - logging.debug( - f"Stream {self.stream_id}:Waiting for FIN before returning data" - ) - await self.conn.stream_events[self.stream_id].wait() - self.conn.stream_events[self.stream_id] = trio.Event() - else: - # No data, wait for data or closure - logging.debug(f"Stream {self.stream_id}: Waiting for data or FIN") - await self.conn.stream_events[self.stream_id].wait() - self.conn.stream_events[self.stream_id] = trio.Event() - - # After loop, check if stream is closed or shutting down - async with self.conn.streams_lock: - if self.conn.event_shutting_down.is_set(): - logging.debug(f"Stream {self.stream_id}: Connection shutting down") - raise MuxedStreamEOF("Connection shut down") - if self.closed: - if self.reset_received: - logging.debug(f"Stream {self.stream_id}: Stream was reset") - raise MuxedStreamReset("Stream was reset") - else: - logging.debug( - f"Stream {self.stream_id}: Stream closed cleanly (EOF)" - ) - raise MuxedStreamEOF("Stream closed cleanly (EOF)") - buffer = self.conn.stream_buffers.get(self.stream_id) - if buffer is None: - logging.debug( - f"Stream {self.stream_id}: Buffer gone, assuming closed" - ) - raise MuxedStreamEOF("Stream buffer closed") - if self.recv_closed and len(buffer) == 0: - logging.debug(f"Stream {self.stream_id}: EOF reached") - raise MuxedStreamEOF("Stream is closed for receiving") - # Return all buffered data + # Check if there's data in the buffer + buffer = self.conn.stream_buffers.get(self.stream_id) + size = len(buffer) if buffer else 0 + if size > 0: + # If any data is available,return it immediately + assert buffer is not None data = bytes(buffer) buffer.clear() - return data - - data = await self.conn.read_stream(self.stream_id, n) - async with self.window_lock: - self.recv_window -= len(data) - # Automatically send a window update if recv_window is low - if self.recv_window <= DEFAULT_WINDOW_SIZE // 2: + async with self.window_lock: + self.recv_window += len(data) + await self.send_window_update(len(data), skip_lock=True) + return data + # Otherwise,wait for data or FIN + if self.recv_closed: + raise MuxedStreamEOF("Stream is closed for receiving") + await self.conn.stream_events[self.stream_id].wait() + self.conn.stream_events[self.stream_id] = trio.Event() + return b"" + else: + data = await self.conn.read_stream(self.stream_id, n) + async with self.window_lock: + self.recv_window += len(data) logging.debug( - f"Stream {self.stream_id}: " - f"Low recv_window ({self.recv_window}), sending update" + f"Stream {self.stream_id}: Sending window update after read, " + f"increment={len(data)}" ) - await self.send_window_update(None, skip_lock=True) - return data + await self.send_window_update(len(data), skip_lock=True) + return data async def close(self) -> None: if not self.send_closed: diff --git a/tests/core/stream_muxer/test_mplex_read_write_lock.py b/tests/core/stream_muxer/test_mplex_read_write_lock.py index afc197ac..d00d5b8e 100644 --- a/tests/core/stream_muxer/test_mplex_read_write_lock.py +++ b/tests/core/stream_muxer/test_mplex_read_write_lock.py @@ -1,7 +1,7 @@ import pytest import trio -from libp2p.abc import ISecureConn +from libp2p.abc import IMuxedStream, ISecureConn from libp2p.crypto.keys import PrivateKey, PublicKey from libp2p.peer.id import ID from libp2p.stream_muxer.mplex.constants import ( @@ -59,13 +59,15 @@ class DummyMuxedConn(Mplex): self.event_started = trio.Event() self.stream_backlog_limit = 256 self.stream_backlog_semaphore = trio.Semaphore(256) - channels = trio.open_memory_channel[MplexStream](0) + # Use IMuxedStream for type consistency with Mplex + channels = trio.open_memory_channel[IMuxedStream](0) self.new_stream_send_channel, self.new_stream_receive_channel = channels async def send_message( - self, flag: HeaderTags, data: bytes, stream_id: StreamID - ) -> None: + self, flag: HeaderTags, data: bytes | None, stream_id: StreamID + ) -> int: await trio.sleep(0.01) + return 0 @pytest.mark.trio @@ -75,10 +77,11 @@ async def test_concurrent_writes_are_serialized(): class LoggingMuxedConn(DummyMuxedConn): async def send_message( - self, flag: HeaderTags, data: bytes, stream_id: StreamID - ) -> None: + self, flag: HeaderTags, data: bytes | None, stream_id: StreamID + ) -> int: send_log.append(data) await trio.sleep(0.01) + return 0 memory_send, memory_recv = trio.open_memory_channel(8) stream = MplexStream( From df17788ec3d9ed8b4f4fc32920ba0fabc1020510 Mon Sep 17 00:00:00 2001 From: kaneki003 Date: Sat, 21 Jun 2025 14:10:09 +0530 Subject: [PATCH 03/23] resolving build-fails --- tests/core/stream_muxer/test_yamux.py | 6 ++++-- tests/core/stream_muxer/test_yamux_read_write_lock.py | 6 ------ 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/tests/core/stream_muxer/test_yamux.py b/tests/core/stream_muxer/test_yamux.py index 81d05676..bfd8eb5a 100644 --- a/tests/core/stream_muxer/test_yamux.py +++ b/tests/core/stream_muxer/test_yamux.py @@ -224,14 +224,16 @@ async def test_yamux_stream_reset(yamux_pair): await client_stream.reset() # After reset, reading should raise MuxedStreamReset or MuxedStreamEOF try: - await server_stream.read() + while True: + await server_stream.read() except (MuxedStreamEOF, MuxedStreamError): pass else: pytest.fail("Expected MuxedStreamEOF or MuxedStreamError") # Verify subsequent operations fail with StreamReset or EOF with pytest.raises(MuxedStreamError): - await server_stream.read() + while True: + await server_stream.read() with pytest.raises(MuxedStreamError): await server_stream.write(b"test") logging.debug("test_yamux_stream_reset complete") diff --git a/tests/core/stream_muxer/test_yamux_read_write_lock.py b/tests/core/stream_muxer/test_yamux_read_write_lock.py index 74a4c2f1..b73284e8 100644 --- a/tests/core/stream_muxer/test_yamux_read_write_lock.py +++ b/tests/core/stream_muxer/test_yamux_read_write_lock.py @@ -191,17 +191,11 @@ async def test_yamux_race_condition_without_locks(yamux_pair): assert len(msg) == MSG_SIZE, ( f"Client message {i} has wrong size: {len(msg)} != {MSG_SIZE}" ) - assert msg.startswith(b"SERVER-MSG-"), ( - f"Client message {i} doesn't start with expected prefix" - ) for i, msg in enumerate(server_received): assert len(msg) == MSG_SIZE, ( f"Server message {i} has wrong size: {len(msg)} != {MSG_SIZE}" ) - assert msg.startswith(b"CLIENT-MSG-"), ( - f"Server message {i} doesn't start with expected prefix" - ) await client_stream.close() await server_stream.close() From d7cdae8a0fc1a58e874af1847e40c466069435f1 Mon Sep 17 00:00:00 2001 From: kaneki003 Date: Sat, 21 Jun 2025 17:51:27 +0530 Subject: [PATCH 04/23] intgrated n==-1 case in read() --- libp2p/stream_muxer/yamux/yamux.py | 84 +++++++++++-------- .../test_yamux_read_write_lock.py | 33 +++++--- 2 files changed, 70 insertions(+), 47 deletions(-) diff --git a/libp2p/stream_muxer/yamux/yamux.py b/libp2p/stream_muxer/yamux/yamux.py index f58e98c4..023251ed 100644 --- a/libp2p/stream_muxer/yamux/yamux.py +++ b/libp2p/stream_muxer/yamux/yamux.py @@ -180,22 +180,6 @@ class YamuxStream(IMuxedStream): async with self.window_lock: await _do_window_update() - async def read_EOF(self) -> bytes: - """ - To read data from stream until it is closed. - """ - data = b"" - try: - while True: - recv = await self.read() - if recv: - data += recv - except MuxedStreamEOF: - logging.debug( - f"Stream {self.stream_id}:EOF reached,total data read:{len(data)} bytes" - ) - return data - async def read(self, n: int | None = -1) -> bytes: # Handle None value for n by converting it to -1 if n is None: @@ -208,25 +192,57 @@ class YamuxStream(IMuxedStream): ) raise MuxedStreamEOF("Stream is closed for receiving") - # If reading until EOF (n == -1), block until stream is closed if n == -1: - # Check if there's data in the buffer - buffer = self.conn.stream_buffers.get(self.stream_id) - size = len(buffer) if buffer else 0 - if size > 0: - # If any data is available,return it immediately - assert buffer is not None - data = bytes(buffer) - buffer.clear() - async with self.window_lock: - self.recv_window += len(data) - await self.send_window_update(len(data), skip_lock=True) - return data - # Otherwise,wait for data or FIN - if self.recv_closed: - raise MuxedStreamEOF("Stream is closed for receiving") - await self.conn.stream_events[self.stream_id].wait() - self.conn.stream_events[self.stream_id] = trio.Event() + data = b"" + while not self.conn.event_shutting_down.is_set(): + # Check if there's data in the buffer + buffer = self.conn.stream_buffers.get(self.stream_id) + + # If buffer is not available, check if stream is closed + if buffer is None: + logging.debug(f"Stream {self.stream_id}: No buffer available") + raise MuxedStreamEOF("Stream buffer closed") + + # If we have data in buffer, process it + if len(buffer) > 0: + chunk = bytes(buffer) + buffer.clear() + data += chunk + + # Send window update for the chunk we just read + async with self.window_lock: + self.recv_window += len(chunk) + logging.debug(f"Stream {self.stream_id}: Update {len(chunk)}") + await self.send_window_update(len(chunk), skip_lock=True) + + # If stream is closed (FIN received) and buffer is empty, break + if self.recv_closed and len(buffer) == 0: + logging.debug(f"Stream {self.stream_id}: Closed with empty buffer") + break + + # If stream was reset, raise reset error + if self.reset_received: + logging.debug(f"Stream {self.stream_id}: Stream was reset") + raise MuxedStreamReset("Stream was reset") + + # Wait for more data or stream closure + logging.debug(f"Stream {self.stream_id}: Waiting for data or FIN") + await self.conn.stream_events[self.stream_id].wait() + self.conn.stream_events[self.stream_id] = trio.Event() + + # After loop exit, first check if we have data to return + if data: + logging.debug( + f"Stream {self.stream_id}: Returning {len(data)} bytes after loop" + ) + return data + + # No data accumulated, now check why we exited the loop + if self.conn.event_shutting_down.is_set(): + logging.debug(f"Stream {self.stream_id}: Connection shutting down") + raise MuxedStreamEOF("Connection shut down") + + # Return empty data return b"" else: data = await self.conn.read_stream(self.stream_id, n) diff --git a/tests/core/stream_muxer/test_yamux_read_write_lock.py b/tests/core/stream_muxer/test_yamux_read_write_lock.py index b73284e8..6981f9d3 100644 --- a/tests/core/stream_muxer/test_yamux_read_write_lock.py +++ b/tests/core/stream_muxer/test_yamux_read_write_lock.py @@ -16,6 +16,7 @@ from libp2p.peer.id import ( from libp2p.security.insecure.transport import ( InsecureTransport, ) +from libp2p.stream_muxer.exceptions import MuxedStreamEOF from libp2p.stream_muxer.yamux.yamux import ( Yamux, YamuxStream, @@ -139,8 +140,8 @@ async def test_yamux_race_condition_without_locks(yamux_pair): client_yamux, server_yamux = yamux_pair client_stream: YamuxStream = await client_yamux.open_stream() server_stream: YamuxStream = await server_yamux.accept_stream() - MSG_COUNT = 10 - MSG_SIZE = 256 * 1024 + MSG_COUNT = 1 + MSG_SIZE = 512 * 1024 client_msgs = [ f"CLIENT-MSG-{i:03d}-".encode().ljust(MSG_SIZE, b"C") for i in range(MSG_COUNT) ] @@ -160,11 +161,17 @@ async def test_yamux_race_condition_without_locks(yamux_pair): async def reader(stream, received, name): """Read messages and store them for verification.""" - for i in range(MSG_COUNT): - data = await stream.read(MSG_SIZE) - received.append(data) - if i % 3 == 0: - await trio.sleep(0.001) + try: + data = await stream.read() + if data: + received.append(data) + except MuxedStreamEOF: + pass + # for i in range(MSG_COUNT): + # data = await stream.read() + # received.append(data) + # if i % 3 == 0: + # await trio.sleep(0.001) # Running all operations concurrently async with trio.open_nursery() as nursery: @@ -173,12 +180,12 @@ async def test_yamux_race_condition_without_locks(yamux_pair): nursery.start_soon(reader, client_stream, client_received, "client") nursery.start_soon(reader, server_stream, server_received, "server") - assert len(client_received) == MSG_COUNT, ( - f"Client received {len(client_received)} messages, expected {MSG_COUNT}" - ) - assert len(server_received) == MSG_COUNT, ( - f"Server received {len(server_received)} messages, expected {MSG_COUNT}" - ) + # assert len(client_received) == MSG_COUNT, ( + # f"Client received {len(client_received)} messages, expected {MSG_COUNT}" + # ) + # assert len(server_received) == MSG_COUNT, ( + # f"Server received {len(server_received)} messages, expected {MSG_COUNT}" + # ) assert client_received == server_msgs, ( "Client did not receive server messages in order or intact!" ) From 724375e1fabac07751242af9beafe20c304ec65c Mon Sep 17 00:00:00 2001 From: kaneki003 Date: Tue, 24 Jun 2025 18:05:15 +0530 Subject: [PATCH 05/23] updated doc-string and reverted mplex-changes --- libp2p/stream_muxer/mplex/mplex_stream.py | 101 +++++++------- libp2p/stream_muxer/yamux/yamux.py | 2 - newsfragments/639.feature.rst | 7 +- .../test_mplex_read_write_lock.py | 127 ------------------ tests/core/stream_muxer/test_yamux.py | 6 +- ...write_lock.py => test_yamux_read_write.py} | 0 6 files changed, 56 insertions(+), 187 deletions(-) delete mode 100644 tests/core/stream_muxer/test_mplex_read_write_lock.py rename tests/core/stream_muxer/{test_yamux_read_write_lock.py => test_yamux_read_write.py} (100%) diff --git a/libp2p/stream_muxer/mplex/mplex_stream.py b/libp2p/stream_muxer/mplex/mplex_stream.py index dc65ac5f..3b640df1 100644 --- a/libp2p/stream_muxer/mplex/mplex_stream.py +++ b/libp2p/stream_muxer/mplex/mplex_stream.py @@ -46,9 +46,8 @@ class MplexStream(IMuxedStream): read_deadline: int | None write_deadline: int | None + # TODO: Add lock for read/write to avoid interleaving receiving messages? close_lock: trio.Lock - read_lock: trio.Lock - write_lock: trio.Lock # NOTE: `dataIn` is size of 8 in Go implementation. incoming_data_channel: "trio.MemoryReceiveChannel[bytes]" @@ -81,8 +80,6 @@ class MplexStream(IMuxedStream): self.event_remote_closed = trio.Event() self.event_reset = trio.Event() self.close_lock = trio.Lock() - self.read_lock = trio.Lock() - self.write_lock = trio.Lock() self.incoming_data_channel = incoming_data_channel self._buf = bytearray() @@ -116,49 +113,48 @@ class MplexStream(IMuxedStream): :param n: number of bytes to read :return: bytes actually read """ - async with self.read_lock: - if n is not None and n < 0: - raise ValueError( - "the number of bytes to read `n` must be non-negative or " - f"`None` to indicate read until EOF, got n={n}" - ) - if self.event_reset.is_set(): - raise MplexStreamReset - if n is None: - return await self._read_until_eof() - if len(self._buf) == 0: - data: bytes - # Peek whether there is data available. If yes, we just read until - # there is no data, then return. + if n is not None and n < 0: + raise ValueError( + "the number of bytes to read `n` must be non-negative or " + f"`None` to indicate read until EOF, got n={n}" + ) + if self.event_reset.is_set(): + raise MplexStreamReset + if n is None: + return await self._read_until_eof() + if len(self._buf) == 0: + data: bytes + # Peek whether there is data available. If yes, we just read until there is + # no data, then return. + try: + data = self.incoming_data_channel.receive_nowait() + self._buf.extend(data) + except trio.EndOfChannel: + raise MplexStreamEOF + except trio.WouldBlock: + # We know `receive` will be blocked here. Wait for data here with + # `receive` and catch all kinds of errors here. try: - data = self.incoming_data_channel.receive_nowait() + data = await self.incoming_data_channel.receive() self._buf.extend(data) except trio.EndOfChannel: - raise MplexStreamEOF - except trio.WouldBlock: - # We know `receive` will be blocked here. Wait for data here with - # `receive` and catch all kinds of errors here. - try: - data = await self.incoming_data_channel.receive() - self._buf.extend(data) - except trio.EndOfChannel: - if self.event_reset.is_set(): - raise MplexStreamReset - if self.event_remote_closed.is_set(): - raise MplexStreamEOF - except trio.ClosedResourceError as error: - # Probably `incoming_data_channel` is closed in `reset` when - # we are waiting for `receive`. - if self.event_reset.is_set(): - raise MplexStreamReset - raise Exception( - "`incoming_data_channel` is closed but stream is not reset." - "This should never happen." - ) from error - self._buf.extend(self._read_return_when_blocked()) - payload = self._buf[:n] - self._buf = self._buf[len(payload) :] - return bytes(payload) + if self.event_reset.is_set(): + raise MplexStreamReset + if self.event_remote_closed.is_set(): + raise MplexStreamEOF + except trio.ClosedResourceError as error: + # Probably `incoming_data_channel` is closed in `reset` when we are + # waiting for `receive`. + if self.event_reset.is_set(): + raise MplexStreamReset + raise Exception( + "`incoming_data_channel` is closed but stream is not reset. " + "This should never happen." + ) from error + self._buf.extend(self._read_return_when_blocked()) + payload = self._buf[:n] + self._buf = self._buf[len(payload) :] + return bytes(payload) async def write(self, data: bytes) -> None: """ @@ -166,15 +162,14 @@ class MplexStream(IMuxedStream): :return: number of bytes written """ - async with self.write_lock: - if self.event_local_closed.is_set(): - raise MplexStreamClosed(f"cannot write to closed stream: data={data!r}") - flag = ( - HeaderTags.MessageInitiator - if self.is_initiator - else HeaderTags.MessageReceiver - ) - await self.muxed_conn.send_message(flag, data, self.stream_id) + if self.event_local_closed.is_set(): + raise MplexStreamClosed(f"cannot write to closed stream: data={data!r}") + flag = ( + HeaderTags.MessageInitiator + if self.is_initiator + else HeaderTags.MessageReceiver + ) + await self.muxed_conn.send_message(flag, data, self.stream_id) async def close(self) -> None: """ diff --git a/libp2p/stream_muxer/yamux/yamux.py b/libp2p/stream_muxer/yamux/yamux.py index 023251ed..f90ba9a1 100644 --- a/libp2p/stream_muxer/yamux/yamux.py +++ b/libp2p/stream_muxer/yamux/yamux.py @@ -77,8 +77,6 @@ class YamuxStream(IMuxedStream): self.send_window = DEFAULT_WINDOW_SIZE self.recv_window = DEFAULT_WINDOW_SIZE self.window_lock = trio.Lock() - self.read_lock = trio.Lock() - self.write_lock = trio.Lock() async def __aenter__(self) -> "YamuxStream": """Enter the async context manager.""" diff --git a/newsfragments/639.feature.rst b/newsfragments/639.feature.rst index 3f3d7510..93476b68 100644 --- a/newsfragments/639.feature.rst +++ b/newsfragments/639.feature.rst @@ -1 +1,6 @@ -Added separate read and write locks to the `MplexStream` & `YamuxStream` class.This ensures thread-safe access and data integrity when multiple coroutines interact with the same MplexStream instance. +Fixed several flow-control and concurrency issues in the `YamuxStream` class. Previously, stress-testing revealed that transferring data over `DEFAULT_WINDOW_SIZE` would break the stream due to inconsistent window update handling and lock management. The fixes include: + +- Removed sending of window updates during writes to maintain correct flow-control. +- Added proper timeout handling when releasing and acquiring locks to prevent concurrency errors. +- Corrected the `read` function to properly handle window updates for both `read_until_EOF` and `read_n_bytes`. +- Added event logging at `send_window_updates` and `waiting_for_window_updates` for better observability. diff --git a/tests/core/stream_muxer/test_mplex_read_write_lock.py b/tests/core/stream_muxer/test_mplex_read_write_lock.py deleted file mode 100644 index d00d5b8e..00000000 --- a/tests/core/stream_muxer/test_mplex_read_write_lock.py +++ /dev/null @@ -1,127 +0,0 @@ -import pytest -import trio - -from libp2p.abc import IMuxedStream, ISecureConn -from libp2p.crypto.keys import PrivateKey, PublicKey -from libp2p.peer.id import ID -from libp2p.stream_muxer.mplex.constants import ( - HeaderTags, -) -from libp2p.stream_muxer.mplex.datastructures import ( - StreamID, -) -from libp2p.stream_muxer.mplex.mplex import ( - Mplex, -) -from libp2p.stream_muxer.mplex.mplex_stream import ( - MplexStream, -) - - -class DummySecureConn(ISecureConn): - """A minimal implementation of ISecureConn for testing.""" - - async def write(self, data: bytes) -> None: - pass - - async def read(self, n: int | None = -1) -> bytes: - return b"" - - async def close(self) -> None: - pass - - def get_remote_address(self) -> tuple[str, int] | None: - return None - - def get_local_peer(self) -> ID: - return ID(b"local") - - def get_local_private_key(self) -> PrivateKey: - return PrivateKey() # Dummy key for testing - - def get_remote_peer(self) -> ID: - return ID(b"remote") - - def get_remote_public_key(self) -> PublicKey: - return PublicKey() # Dummy key for testing - - -class DummyMuxedConn(Mplex): - """A minimal mock of Mplex for testing read/write locks.""" - - def __init__(self) -> None: - self.secured_conn = DummySecureConn() - self.peer_id = ID(b"dummy") - self.streams = {} - self.streams_lock = trio.Lock() - self.event_shutting_down = trio.Event() - self.event_closed = trio.Event() - self.event_started = trio.Event() - self.stream_backlog_limit = 256 - self.stream_backlog_semaphore = trio.Semaphore(256) - # Use IMuxedStream for type consistency with Mplex - channels = trio.open_memory_channel[IMuxedStream](0) - self.new_stream_send_channel, self.new_stream_receive_channel = channels - - async def send_message( - self, flag: HeaderTags, data: bytes | None, stream_id: StreamID - ) -> int: - await trio.sleep(0.01) - return 0 - - -@pytest.mark.trio -async def test_concurrent_writes_are_serialized(): - stream_id = StreamID(1, True) - send_log = [] - - class LoggingMuxedConn(DummyMuxedConn): - async def send_message( - self, flag: HeaderTags, data: bytes | None, stream_id: StreamID - ) -> int: - send_log.append(data) - await trio.sleep(0.01) - return 0 - - memory_send, memory_recv = trio.open_memory_channel(8) - stream = MplexStream( - name="test", - stream_id=stream_id, - muxed_conn=LoggingMuxedConn(), - incoming_data_channel=memory_recv, - ) - - async def writer(data): - await stream.write(data) - - async with trio.open_nursery() as nursery: - for i in range(5): - nursery.start_soon(writer, f"msg-{i}".encode()) - # Order doesn't matter due to concurrent execution - assert sorted(send_log) == sorted([f"msg-{i}".encode() for i in range(5)]) - - -@pytest.mark.trio -async def test_concurrent_reads_are_serialized(): - stream_id = StreamID(2, True) - muxed_conn = DummyMuxedConn() - memory_send, memory_recv = trio.open_memory_channel(8) - results = [] - stream = MplexStream( - name="test", - stream_id=stream_id, - muxed_conn=muxed_conn, - incoming_data_channel=memory_recv, - ) - for i in range(5): - await memory_send.send(f"data-{i}".encode()) - await memory_send.aclose() - - async def reader(): - data = await stream.read(6) - results.append(data) - - async with trio.open_nursery() as nursery: - for _ in range(5): - nursery.start_soon(reader) - assert sorted(results) == [f"data-{i}".encode() for i in range(5)] diff --git a/tests/core/stream_muxer/test_yamux.py b/tests/core/stream_muxer/test_yamux.py index bfd8eb5a..81d05676 100644 --- a/tests/core/stream_muxer/test_yamux.py +++ b/tests/core/stream_muxer/test_yamux.py @@ -224,16 +224,14 @@ async def test_yamux_stream_reset(yamux_pair): await client_stream.reset() # After reset, reading should raise MuxedStreamReset or MuxedStreamEOF try: - while True: - await server_stream.read() + await server_stream.read() except (MuxedStreamEOF, MuxedStreamError): pass else: pytest.fail("Expected MuxedStreamEOF or MuxedStreamError") # Verify subsequent operations fail with StreamReset or EOF with pytest.raises(MuxedStreamError): - while True: - await server_stream.read() + await server_stream.read() with pytest.raises(MuxedStreamError): await server_stream.write(b"test") logging.debug("test_yamux_stream_reset complete") diff --git a/tests/core/stream_muxer/test_yamux_read_write_lock.py b/tests/core/stream_muxer/test_yamux_read_write.py similarity index 100% rename from tests/core/stream_muxer/test_yamux_read_write_lock.py rename to tests/core/stream_muxer/test_yamux_read_write.py From 5ac4fc1abaa0fc5d6d55350cb711550916b9e95e Mon Sep 17 00:00:00 2001 From: kaneki003 Date: Thu, 3 Jul 2025 22:20:35 +0530 Subject: [PATCH 06/23] seperated tests for better understanding --- .../stream_muxer/test_yamux_interleaving.py | 199 ++++++++++++++++++ ...rite.py => test_yamux_interleaving_EOF.py} | 17 +- 2 files changed, 201 insertions(+), 15 deletions(-) create mode 100644 tests/core/stream_muxer/test_yamux_interleaving.py rename tests/core/stream_muxer/{test_yamux_read_write.py => test_yamux_interleaving_EOF.py} (90%) diff --git a/tests/core/stream_muxer/test_yamux_interleaving.py b/tests/core/stream_muxer/test_yamux_interleaving.py new file mode 100644 index 00000000..1ce62952 --- /dev/null +++ b/tests/core/stream_muxer/test_yamux_interleaving.py @@ -0,0 +1,199 @@ +import logging + +import pytest +import trio +from trio.testing import ( + memory_stream_pair, +) + +from libp2p.abc import IRawConnection +from libp2p.crypto.ed25519 import ( + create_new_key_pair, +) +from libp2p.peer.id import ( + ID, +) +from libp2p.security.insecure.transport import ( + InsecureTransport, +) +from libp2p.stream_muxer.yamux.yamux import ( + Yamux, + YamuxStream, +) + + +class TrioStreamAdapter(IRawConnection): + """Adapter to make trio memory streams work with libp2p.""" + + def __init__(self, send_stream, receive_stream, is_initiator=False): + self.send_stream = send_stream + self.receive_stream = receive_stream + self.is_initiator = is_initiator + + async def write(self, data: bytes) -> None: + logging.debug(f"Attempting to write {len(data)} bytes") + with trio.move_on_after(2): + await self.send_stream.send_all(data) + + async def read(self, n: int | None = None) -> bytes: + if n is None or n <= 0: + raise ValueError("Reading unbounded or zero bytes not supported") + logging.debug(f"Attempting to read {n} bytes") + with trio.move_on_after(2): + data = await self.receive_stream.receive_some(n) + logging.debug(f"Read {len(data)} bytes") + return data + + async def close(self) -> None: + logging.debug("Closing stream") + await self.send_stream.aclose() + await self.receive_stream.aclose() + + def get_remote_address(self) -> tuple[str, int] | None: + """Return None since this is a test adapter without real network info.""" + return None + + +@pytest.fixture +def key_pair(): + return create_new_key_pair() + + +@pytest.fixture +def peer_id(key_pair): + return ID.from_pubkey(key_pair.public_key) + + +@pytest.fixture +async def secure_conn_pair(key_pair, peer_id): + """Create a pair of secure connections for testing.""" + logging.debug("Setting up secure_conn_pair") + client_send, server_receive = memory_stream_pair() + server_send, client_receive = memory_stream_pair() + + client_rw = TrioStreamAdapter(client_send, client_receive) + server_rw = TrioStreamAdapter(server_send, server_receive) + + insecure_transport = InsecureTransport(key_pair) + + async def run_outbound(nursery_results): + with trio.move_on_after(5): + client_conn = await insecure_transport.secure_outbound(client_rw, peer_id) + logging.debug("Outbound handshake complete") + nursery_results["client"] = client_conn + + async def run_inbound(nursery_results): + with trio.move_on_after(5): + server_conn = await insecure_transport.secure_inbound(server_rw) + logging.debug("Inbound handshake complete") + nursery_results["server"] = server_conn + + nursery_results = {} + async with trio.open_nursery() as nursery: + nursery.start_soon(run_outbound, nursery_results) + nursery.start_soon(run_inbound, nursery_results) + await trio.sleep(0.1) # Give tasks a chance to finish + + client_conn = nursery_results.get("client") + server_conn = nursery_results.get("server") + + if client_conn is None or server_conn is None: + raise RuntimeError("Handshake failed: client_conn or server_conn is None") + + logging.debug("secure_conn_pair setup complete") + return client_conn, server_conn + + +@pytest.fixture +async def yamux_pair(secure_conn_pair, peer_id): + """Create a pair of Yamux multiplexers for testing.""" + logging.debug("Setting up yamux_pair") + client_conn, server_conn = secure_conn_pair + client_yamux = Yamux(client_conn, peer_id, is_initiator=True) + server_yamux = Yamux(server_conn, peer_id, is_initiator=False) + async with trio.open_nursery() as nursery: + with trio.move_on_after(5): + nursery.start_soon(client_yamux.start) + nursery.start_soon(server_yamux.start) + await trio.sleep(0.1) + logging.debug("yamux_pair started") + yield client_yamux, server_yamux + logging.debug("yamux_pair cleanup") + + +@pytest.mark.trio +async def test_yamux_race_condition_without_locks(yamux_pair): + """ + Test for race-around/interleaving in Yamux streams,when reading in + segments of data. + This launches concurrent writers/readers on both sides of a stream. + If there is no proper locking, the received data may be interleaved + or corrupted. + + The test creates structured messages and verifies they are received + intact and in order. + Without proper locking, concurrent read/write operations could cause + data corruption + or message interleaving, which this test will catch. + """ + client_yamux, server_yamux = yamux_pair + client_stream: YamuxStream = await client_yamux.open_stream() + server_stream: YamuxStream = await server_yamux.accept_stream() + MSG_COUNT = 10 + MSG_SIZE = 256 * 1024 # At max,only DEFAULT_WINDOW_SIZE bytes can be read + client_msgs = [ + f"CLIENT-MSG-{i:03d}-".encode().ljust(MSG_SIZE, b"C") for i in range(MSG_COUNT) + ] + server_msgs = [ + f"SERVER-MSG-{i:03d}-".encode().ljust(MSG_SIZE, b"S") for i in range(MSG_COUNT) + ] + client_received = [] + server_received = [] + + async def writer(stream, msgs, name): + """Write messages with minimal delays to encourage race conditions.""" + for i, msg in enumerate(msgs): + await stream.write(msg) + # Yield control frequently to encourage interleaving + if i % 5 == 0: + await trio.sleep(0.005) + + async def reader(stream, received, name): + """Read messages and store them for verification.""" + for i in range(MSG_COUNT): + data = await stream.read(MSG_SIZE) + received.append(data) + if i % 3 == 0: + await trio.sleep(0.001) + + # Running all operations concurrently + async with trio.open_nursery() as nursery: + nursery.start_soon(writer, client_stream, client_msgs, "client") + nursery.start_soon(writer, server_stream, server_msgs, "server") + nursery.start_soon(reader, client_stream, client_received, "client") + nursery.start_soon(reader, server_stream, server_received, "server") + + assert len(client_received) == MSG_COUNT, ( + f"Client received {len(client_received)} messages, expected {MSG_COUNT}" + ) + assert len(server_received) == MSG_COUNT, ( + f"Server received {len(server_received)} messages, expected {MSG_COUNT}" + ) + assert client_received == server_msgs, ( + "Client did not receive server messages in order or intact!" + ) + assert server_received == client_msgs, ( + "Server did not receive client messages in order or intact!" + ) + for i, msg in enumerate(client_received): + assert len(msg) == MSG_SIZE, ( + f"Client message {i} has wrong size: {len(msg)} != {MSG_SIZE}" + ) + + for i, msg in enumerate(server_received): + assert len(msg) == MSG_SIZE, ( + f"Server message {i} has wrong size: {len(msg)} != {MSG_SIZE}" + ) + + await client_stream.close() + await server_stream.close() diff --git a/tests/core/stream_muxer/test_yamux_read_write.py b/tests/core/stream_muxer/test_yamux_interleaving_EOF.py similarity index 90% rename from tests/core/stream_muxer/test_yamux_read_write.py rename to tests/core/stream_muxer/test_yamux_interleaving_EOF.py index 6981f9d3..23d2c2b4 100644 --- a/tests/core/stream_muxer/test_yamux_read_write.py +++ b/tests/core/stream_muxer/test_yamux_interleaving_EOF.py @@ -125,8 +125,8 @@ async def yamux_pair(secure_conn_pair, peer_id): @pytest.mark.trio async def test_yamux_race_condition_without_locks(yamux_pair): """ - Test for race-around/interleaving in Yamux streams when read/write - locks are disabled. + Test for race-around/interleaving in Yamux streams,when reading till + EOF is being used. This launches concurrent writers/readers on both sides of a stream. If there is no proper locking, the received data may be interleaved or corrupted. @@ -167,11 +167,6 @@ async def test_yamux_race_condition_without_locks(yamux_pair): received.append(data) except MuxedStreamEOF: pass - # for i in range(MSG_COUNT): - # data = await stream.read() - # received.append(data) - # if i % 3 == 0: - # await trio.sleep(0.001) # Running all operations concurrently async with trio.open_nursery() as nursery: @@ -180,12 +175,6 @@ async def test_yamux_race_condition_without_locks(yamux_pair): nursery.start_soon(reader, client_stream, client_received, "client") nursery.start_soon(reader, server_stream, server_received, "server") - # assert len(client_received) == MSG_COUNT, ( - # f"Client received {len(client_received)} messages, expected {MSG_COUNT}" - # ) - # assert len(server_received) == MSG_COUNT, ( - # f"Server received {len(server_received)} messages, expected {MSG_COUNT}" - # ) assert client_received == server_msgs, ( "Client did not receive server messages in order or intact!" ) @@ -193,8 +182,6 @@ async def test_yamux_race_condition_without_locks(yamux_pair): "Server did not receive client messages in order or intact!" ) for i, msg in enumerate(client_received): - # logging.debug(f"datatype of msg: {type(msg)}, length: {len(msg)}") - # logging.debug(f"datatype of msg: {type(b"SERVER-MSG-")}") assert len(msg) == MSG_SIZE, ( f"Client message {i} has wrong size: {len(msg)} != {MSG_SIZE}" ) From 31b6a6f237f5dfe7f7ee73ca134965ccb23dec37 Mon Sep 17 00:00:00 2001 From: lla-dane Date: Fri, 27 Jun 2025 14:33:49 +0530 Subject: [PATCH 07/23] todo/bounded nursery in identify-push --- libp2p/identity/identify_push/identify_push.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libp2p/identity/identify_push/identify_push.py b/libp2p/identity/identify_push/identify_push.py index 914264ed..77a65187 100644 --- a/libp2p/identity/identify_push/identify_push.py +++ b/libp2p/identity/identify_push/identify_push.py @@ -190,4 +190,4 @@ async def push_identify_to_peers( # limiting concurrent connections to 10 async with trio.open_nursery() as nursery: for peer_id in peer_ids: - nursery.start_soon(push_identify_to_peer, host, peer_id, observed_multiaddr) + nursery.start_soon(limited_push, peer_id) From a89ba8ef81e98be2a7c2c9ab8def271c6acf2bb5 Mon Sep 17 00:00:00 2001 From: lla-dane Date: Fri, 27 Jun 2025 14:50:10 +0530 Subject: [PATCH 08/23] added newsfragment --- newsfragments/708.performance.rst | 2 ++ 1 file changed, 2 insertions(+) create mode 100644 newsfragments/708.performance.rst diff --git a/newsfragments/708.performance.rst b/newsfragments/708.performance.rst new file mode 100644 index 00000000..551a70e4 --- /dev/null +++ b/newsfragments/708.performance.rst @@ -0,0 +1,2 @@ +Limit concurrency in `push_identify_to_peers` to prevent resource exhaustion under high peer counts. +This makes peer communication more stable and efficient, especially at scale. From 383d7cb72211ed3cde415f8c5d7b4cecb7171d5e Mon Sep 17 00:00:00 2001 From: lla-dane Date: Wed, 2 Jul 2025 19:43:31 +0530 Subject: [PATCH 09/23] added tests --- libp2p/identity/identify_push/identify_push.py | 6 +++++- tests/core/identity/identify_push/test_identify_push.py | 1 + 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/libp2p/identity/identify_push/identify_push.py b/libp2p/identity/identify_push/identify_push.py index 77a65187..cf39fd76 100644 --- a/libp2p/identity/identify_push/identify_push.py +++ b/libp2p/identity/identify_push/identify_push.py @@ -176,7 +176,9 @@ async def push_identify_to_peers( host: IHost, peer_ids: set[ID] | None = None, observed_multiaddr: Multiaddr | None = None, -) -> None: + counter: dict[str, int] | None = None, + lock: trio.Lock | None = None, +) -> int: # <-- return the max concurrency """ Push an identify message to multiple peers in parallel. @@ -191,3 +193,5 @@ async def push_identify_to_peers( async with trio.open_nursery() as nursery: for peer_id in peer_ids: nursery.start_soon(limited_push, peer_id) + + return counter["max"] if counter else 0 diff --git a/tests/core/identity/identify_push/test_identify_push.py b/tests/core/identity/identify_push/test_identify_push.py index b0ffb677..292d1948 100644 --- a/tests/core/identity/identify_push/test_identify_push.py +++ b/tests/core/identity/identify_push/test_identify_push.py @@ -10,6 +10,7 @@ import trio from libp2p import ( new_host, ) +from libp2p.abc import IHost from libp2p.crypto.secp256k1 import ( create_new_key_pair, ) From 8bfd4bde948edc6ad04f47cc028faa8f863ab2a3 Mon Sep 17 00:00:00 2001 From: lla-dane Date: Thu, 3 Jul 2025 11:53:58 +0530 Subject: [PATCH 10/23] created concurrency limit configurable --- libp2p/identity/identify_push/identify_push.py | 1 + 1 file changed, 1 insertion(+) diff --git a/libp2p/identity/identify_push/identify_push.py b/libp2p/identity/identify_push/identify_push.py index cf39fd76..e9e979af 100644 --- a/libp2p/identity/identify_push/identify_push.py +++ b/libp2p/identity/identify_push/identify_push.py @@ -178,6 +178,7 @@ async def push_identify_to_peers( observed_multiaddr: Multiaddr | None = None, counter: dict[str, int] | None = None, lock: trio.Lock | None = None, + limit: int = CONCURRENCY_LIMIT, ) -> int: # <-- return the max concurrency """ Push an identify message to multiple peers in parallel. From a7d122a0f944c15fc377b28911a15cbc249ac5b8 Mon Sep 17 00:00:00 2001 From: lla-dane Date: Fri, 4 Jul 2025 17:28:44 +0530 Subject: [PATCH 11/23] added extra tests for identifu push for concurrency cap --- .../identity/identify_push/identify_push.py | 9 +- .../identify_push/test_identify_push.py | 91 ++++++++++++++++++- tests/utils/utils.py | 21 ++++- 3 files changed, 111 insertions(+), 10 deletions(-) diff --git a/libp2p/identity/identify_push/identify_push.py b/libp2p/identity/identify_push/identify_push.py index e9e979af..914264ed 100644 --- a/libp2p/identity/identify_push/identify_push.py +++ b/libp2p/identity/identify_push/identify_push.py @@ -176,10 +176,7 @@ async def push_identify_to_peers( host: IHost, peer_ids: set[ID] | None = None, observed_multiaddr: Multiaddr | None = None, - counter: dict[str, int] | None = None, - lock: trio.Lock | None = None, - limit: int = CONCURRENCY_LIMIT, -) -> int: # <-- return the max concurrency +) -> None: """ Push an identify message to multiple peers in parallel. @@ -193,6 +190,4 @@ async def push_identify_to_peers( # limiting concurrent connections to 10 async with trio.open_nursery() as nursery: for peer_id in peer_ids: - nursery.start_soon(limited_push, peer_id) - - return counter["max"] if counter else 0 + nursery.start_soon(push_identify_to_peer, host, peer_id, observed_multiaddr) diff --git a/tests/core/identity/identify_push/test_identify_push.py b/tests/core/identity/identify_push/test_identify_push.py index 292d1948..935fb2c0 100644 --- a/tests/core/identity/identify_push/test_identify_push.py +++ b/tests/core/identity/identify_push/test_identify_push.py @@ -10,7 +10,6 @@ import trio from libp2p import ( new_host, ) -from libp2p.abc import IHost from libp2p.crypto.secp256k1 import ( create_new_key_pair, ) @@ -36,6 +35,8 @@ from tests.utils.factories import ( ) from tests.utils.utils import ( create_mock_connections, + run_host_forever, + wait_until_listening, ) logger = logging.getLogger("libp2p.identity.identify-push-test") @@ -504,3 +505,91 @@ async def test_push_identify_to_peers_respects_concurrency_limit(): assert state["max_observed"] <= CONCURRENCY_LIMIT, ( f"Max concurrency observed: {state['max_observed']}" ) + + +@pytest.mark.trio +async def test_all_peers_receive_identify_push_with_semaphore(security_protocol): + dummy_peers = [] + + async with host_pair_factory(security_protocol=security_protocol) as (host_a, _): + # Create dummy peers + for _ in range(50): + key_pair = create_new_key_pair() + dummy_host = new_host(key_pair=key_pair) + dummy_host.set_stream_handler( + ID_PUSH, identify_push_handler_for(dummy_host) + ) + listen_addr = multiaddr.Multiaddr("/ip4/127.0.0.1/tcp/0") + dummy_peers.append((dummy_host, listen_addr)) + + async with trio.open_nursery() as nursery: + # Start all dummy hosts + for host, listen_addr in dummy_peers: + nursery.start_soon(run_host_forever, host, listen_addr) + + # Wait for all hosts to finish setting up listeners + for host, _ in dummy_peers: + await wait_until_listening(host) + + # Now connect host_a → dummy peers + for host, _ in dummy_peers: + await host_a.connect(info_from_p2p_addr(host.get_addrs()[0])) + + await push_identify_to_peers( + host_a, + ) + + await trio.sleep(0.5) + + peer_id_a = host_a.get_id() + for host, _ in dummy_peers: + dummy_peerstore = host.get_peerstore() + assert peer_id_a in dummy_peerstore.peer_ids() + + nursery.cancel_scope.cancel() + + +@pytest.mark.trio +async def test_all_peers_receive_identify_push_with_semaphore_under_high_peer_load( + security_protocol, +): + dummy_peers = [] + + async with host_pair_factory(security_protocol=security_protocol) as (host_a, _): + # Create dummy peers + # Breaking with more than 500 peers + # Trio have a async tasks limit of 1000 + for _ in range(499): + key_pair = create_new_key_pair() + dummy_host = new_host(key_pair=key_pair) + dummy_host.set_stream_handler( + ID_PUSH, identify_push_handler_for(dummy_host) + ) + listen_addr = multiaddr.Multiaddr("/ip4/127.0.0.1/tcp/0") + dummy_peers.append((dummy_host, listen_addr)) + + async with trio.open_nursery() as nursery: + # Start all dummy hosts + for host, listen_addr in dummy_peers: + nursery.start_soon(run_host_forever, host, listen_addr) + + # Wait for all hosts to finish setting up listeners + for host, _ in dummy_peers: + await wait_until_listening(host) + + # Now connect host_a → dummy peers + for host, _ in dummy_peers: + await host_a.connect(info_from_p2p_addr(host.get_addrs()[0])) + + await push_identify_to_peers( + host_a, + ) + + await trio.sleep(0.5) + + peer_id_a = host_a.get_id() + for host, _ in dummy_peers: + dummy_peerstore = host.get_peerstore() + assert peer_id_a in dummy_peerstore.peer_ids() + + nursery.cancel_scope.cancel() diff --git a/tests/utils/utils.py b/tests/utils/utils.py index 6e23ecdd..4fe2a6e1 100644 --- a/tests/utils/utils.py +++ b/tests/utils/utils.py @@ -2,13 +2,30 @@ from unittest.mock import ( MagicMock, ) +import trio -def create_mock_connections() -> dict: +from libp2p.abc import IHost + + +def create_mock_connections(count: int = 50) -> dict: connections = {} - for i in range(1, 31): + for i in range(1, count): peer_id = f"peer-{i}" mock_conn = MagicMock(name=f"INetConn-{i}") connections[peer_id] = mock_conn return connections + + +async def run_host_forever(host: IHost, addr): + async with host.run([addr]): + await trio.sleep_forever() + + +async def wait_until_listening(host, timeout=3): + with trio.move_on_after(timeout): + while not host.get_addrs(): + await trio.sleep(0.05) + return + raise RuntimeError("Timed out waiting for host to get an address") From bfe3dee781365f9674700f9f388f7c906c7e4cf2 Mon Sep 17 00:00:00 2001 From: lla-dane Date: Fri, 4 Jul 2025 17:32:48 +0530 Subject: [PATCH 12/23] updated newsfragment --- newsfragments/708.performance.rst | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/newsfragments/708.performance.rst b/newsfragments/708.performance.rst index 551a70e4..286615e9 100644 --- a/newsfragments/708.performance.rst +++ b/newsfragments/708.performance.rst @@ -1,2 +1 @@ -Limit concurrency in `push_identify_to_peers` to prevent resource exhaustion under high peer counts. -This makes peer communication more stable and efficient, especially at scale. +Added extra tests for identify push concurrency cap under high peer load From f3d8cbf9689f233d269918c225e84cebfcba1494 Mon Sep 17 00:00:00 2001 From: lla-dane Date: Sun, 1 Jun 2025 23:02:41 +0530 Subject: [PATCH 13/23] feat: Matching go-libp2p PeerStore implementation --- libp2p/abc.py | 610 +++++++++++++-------- libp2p/peer/peerdata.py | 82 ++- libp2p/peer/peerstore.py | 62 ++- libp2p/tools/async_service/base.py | 3 +- libp2p/tools/async_service/trio_service.py | 3 +- 5 files changed, 535 insertions(+), 225 deletions(-) diff --git a/libp2p/abc.py b/libp2p/abc.py index dc941c43..0b63ed04 100644 --- a/libp2p/abc.py +++ b/libp2p/abc.py @@ -385,6 +385,10 @@ class IPeerMetadata(ABC): :raises Exception: If the operation is unsuccessful. """ + @abstractmethod + def clear_metadata(self, peer_id: ID) -> None: + """Clears the metadata""" + # -------------------------- addrbook interface.py -------------------------- @@ -475,11 +479,114 @@ class IAddrBook(ABC): """ + @abstractmethod + def set_addr(self, peer_id: ID, addr: Multiaddr, ttl: int) -> None: + """Set addr""" + + @abstractmethod + def set_addrs(self, peer_id: ID, addrs: Sequence[Multiaddr], ttl: int) -> None: + """Set addrs""" + + @abstractmethod + def update_addrs(self, peer_id: ID, oldTTL: int, newTTL: int) -> None: + """Update addrs""" + + @abstractmethod + def addr_stream(self, peer_id: ID) -> None: + """Addr stream""" + + +# -------------------------- keybook interface.py -------------------------- + + +class IKeyBook(ABC): + """IKeyBook""" + + @abstractmethod + def pubkey(self, peer_id: ID) -> PublicKey: + """Pubkey""" + + @abstractmethod + def privkey(self, peer_id: ID) -> PrivateKey: + """Privkey""" + + @abstractmethod + def add_pubkey(self, peer_id: ID, pubkey: PublicKey) -> None: + """add_pubkey""" + + @abstractmethod + def add_privkey(self, peer_id: ID, privkey: PrivateKey) -> None: + """add_privkey""" + + @abstractmethod + def add_key_pair(self, peer_id: ID, key_pair: KeyPair) -> None: + """add_key_pair""" + + @abstractmethod + def peer_with_keys(self) -> list[ID]: + """peer_with_keys""" + + @abstractmethod + def clear_keydata(self, peer_id: ID) -> PublicKey: + """clear_keydata""" + + +# -------------------------- metrics interface.py -------------------------- + + +class IMetrics(ABC): + """IMetrics""" + + @abstractmethod + def record_latency(self, peer_id: ID, RTT: float) -> None: + """record_latency""" + + @abstractmethod + def latency_EWMA(self, peer_id: ID) -> float: + """latency_EWMA""" + + @abstractmethod + def clear_metrics(self, peer_id: ID) -> None: + """clear_metrics""" + + +# -------------------------- protobook interface.py -------------------------- + + +class IProtoBook(ABC): + @abstractmethod + def get_protocols(self, peer_id: ID) -> list[str]: + """get_protocols""" + + @abstractmethod + def add_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None: + """add_protocols""" + + @abstractmethod + def set_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None: + """set_protocols""" + + @abstractmethod + def remove_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None: + """remove_protocols""" + + @abstractmethod + def supports_protocols(self, peer_id: ID, protocols: Sequence[str]) -> list[str]: + """supports_protocols""" + + @abstractmethod + def first_supported_protocol(self, peer_id: ID, protocols: Sequence[str]) -> str: + """first_supported_protocol""" + + @abstractmethod + def clear_protocol_data(self, peer_id: ID) -> None: + """clear_protocol_data""" + # -------------------------- peerstore interface.py -------------------------- -class IPeerStore(IAddrBook, IPeerMetadata): +class IPeerStore(IPeerMetadata, IAddrBook, IKeyBook, IMetrics, IProtoBook): """ Interface for a peer store. @@ -488,9 +595,98 @@ class IPeerStore(IAddrBook, IPeerMetadata): """ @abstractmethod - def peer_info(self, peer_id: ID) -> PeerInfo: + def get(self, peer_id: ID, key: str) -> Any: """ - Retrieve the peer information for the specified peer. + Retrieve the value associated with a key for a specified peer. + + Parameters + ---------- + peer_id : ID + The identifier of the peer. + key : str + The key to look up. + + Returns + ------- + Any + The value corresponding to the specified key. + + Raises + ------ + PeerStoreError + If the peer ID or value is not found. + + """ + + @abstractmethod + def put(self, peer_id: ID, key: str, val: Any) -> None: + """ + Store a key-value pair for the specified peer. + + Parameters + ---------- + peer_id : ID + The identifier of the peer. + key : str + The key for the data. + val : Any + The value to store. + + """ + + @abstractmethod + def clear_metadata(self, peer_id: ID) -> None: + """clear_metadata""" + + ## + @abstractmethod + def add_addr(self, peer_id: ID, addr: Multiaddr, ttl: int) -> None: + """ + Add an address for the specified peer. + + Parameters + ---------- + peer_id : ID + The identifier of the peer. + addr : Multiaddr + The multiaddress to add. + ttl : int + The time-to-live for the record. + + """ + + @abstractmethod + def add_addrs(self, peer_id: ID, addrs: Sequence[Multiaddr], ttl: int) -> None: + """ + Add multiple addresses for the specified peer. + + Parameters + ---------- + peer_id : ID + The identifier of the peer. + addrs : Sequence[Multiaddr] + A sequence of multiaddresses to add. + ttl : int + The time-to-live for the record. + + """ + + @abstractmethod + def set_addr(self, peer_id: ID, addr: Multiaddr, ttl: int) -> None: + """set_addr""" + + @abstractmethod + def set_addrs(self, peer_id: ID, addrs: Sequence[Multiaddr], ttl: int) -> None: + """set_addrs""" + + @abstractmethod + def update_addrs(self, peer_id: ID, oldTTL: int, newTTL: int) -> None: + """update_addrs""" + + @abstractmethod + def addrs(self, peer_id: ID) -> list[Multiaddr]: + """ + Retrieve the addresses for the specified peer. Parameters ---------- @@ -499,11 +695,163 @@ class IPeerStore(IAddrBook, IPeerMetadata): Returns ------- - PeerInfo - The peer information object for the given peer. + list[Multiaddr] + A list of multiaddresses. """ + @abstractmethod + def clear_addrs(self, peer_id: ID) -> None: + """ + Clear all addresses for the specified peer. + + Parameters + ---------- + peer_id : ID + The identifier of the peer. + + """ + + @abstractmethod + def peers_with_addrs(self) -> list[ID]: + """ + Retrieve all peer identifiers with stored addresses. + + Returns + ------- + list[ID] + A list of peer IDs. + + """ + + @abstractmethod + def addr_stream(self, peer_id: ID) -> None: + """addr_stream""" + + ## + @abstractmethod + def pubkey(self, peer_id: ID) -> PublicKey: + """ + Retrieve the public key for the specified peer. + + Parameters + ---------- + peer_id : ID + The identifier of the peer. + + Returns + ------- + PublicKey + The public key of the peer. + + Raises + ------ + PeerStoreError + If the peer ID is not found. + + """ + + @abstractmethod + def privkey(self, peer_id: ID) -> PrivateKey: + """ + Retrieve the private key for the specified peer. + + Parameters + ---------- + peer_id : ID + The identifier of the peer. + + Returns + ------- + PrivateKey + The private key of the peer. + + Raises + ------ + PeerStoreError + If the peer ID is not found. + + """ + + @abstractmethod + def add_pubkey(self, peer_id: ID, pubkey: PublicKey) -> None: + """ + Add a public key for the specified peer. + + Parameters + ---------- + peer_id : ID + The identifier of the peer. + pubkey : PublicKey + The public key to add. + + Raises + ------ + PeerStoreError + If the peer already has a public key set. + + """ + + @abstractmethod + def add_privkey(self, peer_id: ID, privkey: PrivateKey) -> None: + """ + Add a private key for the specified peer. + + Parameters + ---------- + peer_id : ID + The identifier of the peer. + privkey : PrivateKey + The private key to add. + + Raises + ------ + PeerStoreError + If the peer already has a private key set. + + """ + + @abstractmethod + def add_key_pair(self, peer_id: ID, key_pair: KeyPair) -> None: + """ + Add a key pair for the specified peer. + + Parameters + ---------- + peer_id : ID + The identifier of the peer. + key_pair : KeyPair + The key pair to add. + + Raises + ------ + PeerStoreError + If the peer already has a public or private key set. + + """ + + @abstractmethod + def peer_with_keys(self) -> list[ID]: + """peer_with_keys""" + + @abstractmethod + def clear_keydata(self, peer_id: ID) -> PublicKey: + """clear_keydata""" + + ## + @abstractmethod + def record_latency(self, peer_id: ID, RTT: float) -> None: + """record_latency""" + + @abstractmethod + def latency_EWMA(self, peer_id: ID) -> float: + """latency_EWMA""" + + @abstractmethod + def clear_metrics(self, peer_id: ID) -> None: + """clear_metrics""" + + ## @abstractmethod def get_protocols(self, peer_id: ID) -> list[str]: """ @@ -554,6 +902,40 @@ class IPeerStore(IAddrBook, IPeerMetadata): """ + @abstractmethod + def remove_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None: + """remove_protocols""" + + @abstractmethod + def supports_protocols(self, peer_id: ID, protocols: Sequence[str]) -> list[str]: + """supports_protocols""" + + @abstractmethod + def first_supported_protocol(self, peer_id: ID, protocols: Sequence[str]) -> str: + """first_supported_protocol""" + + @abstractmethod + def clear_protocol_data(self, peer_id: ID) -> None: + """clear_protocol_data""" + + ## + @abstractmethod + def peer_info(self, peer_id: ID) -> PeerInfo: + """ + Retrieve the peer information for the specified peer. + + Parameters + ---------- + peer_id : ID + The identifier of the peer. + + Returns + ------- + PeerInfo + The peer information object for the given peer. + + """ + @abstractmethod def peer_ids(self) -> list[ID]: """ @@ -567,218 +949,8 @@ class IPeerStore(IAddrBook, IPeerMetadata): """ @abstractmethod - def get(self, peer_id: ID, key: str) -> Any: - """ - Retrieve the value associated with a key for a specified peer. - - Parameters - ---------- - peer_id : ID - The identifier of the peer. - key : str - The key to look up. - - Returns - ------- - Any - The value corresponding to the specified key. - - Raises - ------ - PeerStoreError - If the peer ID or value is not found. - - """ - - @abstractmethod - def put(self, peer_id: ID, key: str, val: Any) -> None: - """ - Store a key-value pair for the specified peer. - - Parameters - ---------- - peer_id : ID - The identifier of the peer. - key : str - The key for the data. - val : Any - The value to store. - - """ - - @abstractmethod - def add_addr(self, peer_id: ID, addr: Multiaddr, ttl: int) -> None: - """ - Add an address for the specified peer. - - Parameters - ---------- - peer_id : ID - The identifier of the peer. - addr : Multiaddr - The multiaddress to add. - ttl : int - The time-to-live for the record. - - """ - - @abstractmethod - def add_addrs(self, peer_id: ID, addrs: Sequence[Multiaddr], ttl: int) -> None: - """ - Add multiple addresses for the specified peer. - - Parameters - ---------- - peer_id : ID - The identifier of the peer. - addrs : Sequence[Multiaddr] - A sequence of multiaddresses to add. - ttl : int - The time-to-live for the record. - - """ - - @abstractmethod - def addrs(self, peer_id: ID) -> list[Multiaddr]: - """ - Retrieve the addresses for the specified peer. - - Parameters - ---------- - peer_id : ID - The identifier of the peer. - - Returns - ------- - list[Multiaddr] - A list of multiaddresses. - - """ - - @abstractmethod - def clear_addrs(self, peer_id: ID) -> None: - """ - Clear all addresses for the specified peer. - - Parameters - ---------- - peer_id : ID - The identifier of the peer. - - """ - - @abstractmethod - def peers_with_addrs(self) -> list[ID]: - """ - Retrieve all peer identifiers with stored addresses. - - Returns - ------- - list[ID] - A list of peer IDs. - - """ - - @abstractmethod - def add_pubkey(self, peer_id: ID, pubkey: PublicKey) -> None: - """ - Add a public key for the specified peer. - - Parameters - ---------- - peer_id : ID - The identifier of the peer. - pubkey : PublicKey - The public key to add. - - Raises - ------ - PeerStoreError - If the peer already has a public key set. - - """ - - @abstractmethod - def pubkey(self, peer_id: ID) -> PublicKey: - """ - Retrieve the public key for the specified peer. - - Parameters - ---------- - peer_id : ID - The identifier of the peer. - - Returns - ------- - PublicKey - The public key of the peer. - - Raises - ------ - PeerStoreError - If the peer ID is not found. - - """ - - @abstractmethod - def add_privkey(self, peer_id: ID, privkey: PrivateKey) -> None: - """ - Add a private key for the specified peer. - - Parameters - ---------- - peer_id : ID - The identifier of the peer. - privkey : PrivateKey - The private key to add. - - Raises - ------ - PeerStoreError - If the peer already has a private key set. - - """ - - @abstractmethod - def privkey(self, peer_id: ID) -> PrivateKey: - """ - Retrieve the private key for the specified peer. - - Parameters - ---------- - peer_id : ID - The identifier of the peer. - - Returns - ------- - PrivateKey - The private key of the peer. - - Raises - ------ - PeerStoreError - If the peer ID is not found. - - """ - - @abstractmethod - def add_key_pair(self, peer_id: ID, key_pair: KeyPair) -> None: - """ - Add a key pair for the specified peer. - - Parameters - ---------- - peer_id : ID - The identifier of the peer. - key_pair : KeyPair - The key pair to add. - - Raises - ------ - PeerStoreError - If the peer already has a public or private key set. - - """ + def clear_peerdata(self, peer_id: ID) -> None: + """clear_peerdata""" # -------------------------- listener interface.py -------------------------- @@ -1316,7 +1488,7 @@ class IPeerData(ABC): """ @abstractmethod - def add_addrs(self, addrs: Sequence[Multiaddr]) -> None: + def add_addrs(self, addrs: Sequence[Multiaddr], ttl: int) -> None: """ Add multiple multiaddresses to the peer's data. @@ -1324,6 +1496,8 @@ class IPeerData(ABC): ---------- addrs : Sequence[Multiaddr] A sequence of multiaddresses to add. + ttl: inr + Time to live for the peer record """ diff --git a/libp2p/peer/peerdata.py b/libp2p/peer/peerdata.py index 386e31ef..bd0c4d0b 100644 --- a/libp2p/peer/peerdata.py +++ b/libp2p/peer/peerdata.py @@ -18,6 +18,13 @@ from libp2p.crypto.keys import ( PublicKey, ) +""" +Latency EWMA Smoothing governs the deacy of the EWMA (the speed at which +is changes). This must be a normalized (0-1) value. +1 is 100% change, 0 is no change. +""" +LATENCY_EWMA_SMOOTHING = 0.1 + class PeerData(IPeerData): pubkey: PublicKey | None @@ -55,13 +62,82 @@ class PeerData(IPeerData): """ self.protocols = list(protocols) - def add_addrs(self, addrs: Sequence[Multiaddr]) -> None: + def remove_protocols(self, protocols: Sequence[str]) -> None: + """ + :param protocols: protocols to remove + """ + for protocol in protocols: + if protocol in self.protocols: + self.protocols.remove(protocol) + + def supports_protocols(self, protocols: Sequence[str]) -> list[str]: + """ + :param protocols: protocols to check from + :return: all supported protocols in the given list + """ + return [proto for proto in protocols if proto in self.protocols] + + def first_supported_protocol(self, protocols: Sequence[str]) -> str: + """ + :param protocols: protocols to check from + :return: first supported protocol in the given list + """ + for protocol in protocols: + if protocol in self.protocols: + return protocol + + return "None supported" + + def clear_protocol_data(self) -> None: + """Clear all protocols""" + self.protocols = [] + + def add_addrs(self, addrs: Sequence[Multiaddr], ttl: int) -> None: """ :param addrs: multiaddresses to add """ + expiry = time.time() + ttl if ttl is not None else float("inf") for addr in addrs: if addr not in self.addrs: self.addrs.append(addr) + current_expiry = self.addrs_ttl.get(addr, 0) + if expiry > current_expiry: + self.addrs_ttl[addr] = expiry + + def set_addrs(self, addrs: Sequence[Multiaddr], ttl: int) -> None: + """ + :param addrs: multiaddresses to update + :param ttl: new ttl + """ + now = time.time() + + if ttl <= 0: + # Put the TTL value to -1 + for addr in addrs: + # TODO! if addr in self.addrs, remove them? + if addr in self.addrs_ttl: + del self.addrs_ttl[addr] + return + + expiry = now + ttl + for addr in addrs: + # TODO! if addr not in self.addrs, add them? + self.addrs_ttl[addr] = expiry + + def update_addrs(self, oldTTL: int, newTTL: int) -> None: + """ + :param oldTTL: old ttl + :param newTTL: new ttl + """ + now = time.time() + + new_expiry = now + newTTL + old_expiry = now + oldTTL + + for addr, expiry in list(self.addrs_ttl.items()): + # Approximate match by expiry time + if abs(expiry - old_expiry) < 1: + self.addrs_ttl[addr] = new_expiry def get_addrs(self) -> list[Multiaddr]: """ @@ -90,6 +166,10 @@ class PeerData(IPeerData): return self.metadata[key] raise PeerDataError("key not found") + def clear_metadata(self) -> None: + """Clears metadata.""" + self.metadata = {} + def add_pubkey(self, pubkey: PublicKey) -> None: """ :param pubkey: diff --git a/libp2p/peer/peerstore.py b/libp2p/peer/peerstore.py index 3bb729d2..ada56f47 100644 --- a/libp2p/peer/peerstore.py +++ b/libp2p/peer/peerstore.py @@ -53,6 +53,15 @@ class PeerStore(IPeerStore): return PeerInfo(peer_id, peer_data.get_addrs()) raise PeerStoreError("peer ID not found") + def peer_ids(self) -> list[ID]: + """ + :return: all of the peer IDs stored in peer store + """ + return list(self.peer_data_map.keys()) + + def clear_peerdata(self, peer_id: ID) -> None: + """Clears the peer data of the peer""" + def get_protocols(self, peer_id: ID) -> list[str]: """ :param peer_id: peer ID to get protocols for @@ -79,7 +88,15 @@ class PeerStore(IPeerStore): peer_data = self.peer_data_map[peer_id] peer_data.set_protocols(list(protocols)) - def peer_ids(self) -> list[ID]: + def remove_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None: + """ + :param peer_id: peer ID to get info for + :param protocols: unsupported protocols to remove + """ + peer_data = self.peer_data_map[peer_id] + peer_data.remove_protocols(protocols) + + def supports_protocols(self, peer_id: ID, protocols: Sequence[str]) -> list[str]: """ :return: all of the peer IDs stored in peer store """ @@ -165,7 +182,7 @@ class PeerStore(IPeerStore): def peers_with_addrs(self) -> list[ID]: """ - :return: all of the peer IDs which has addrs stored in peer store + :return: all of the peer IDs which has addrsfloat stored in peer store """ # Add all peers with addrs at least 1 to output output: list[ID] = [] @@ -179,6 +196,10 @@ class PeerStore(IPeerStore): peer_data.clear_addrs() return output + def addr_stream(self, peer_id: ID) -> None: + """addr_stream""" + # TODO! + def add_pubkey(self, peer_id: ID, pubkey: PublicKey) -> None: """ :param peer_id: peer ID to add public key for @@ -239,6 +260,43 @@ class PeerStore(IPeerStore): self.add_pubkey(peer_id, key_pair.public_key) self.add_privkey(peer_id, key_pair.private_key) + def peer_with_keys(self) -> list[ID]: + """Returns the peer_ids for which keys are stored""" + return [ + peer_id + for peer_id, pdata in self.peer_data_map.items() + if pdata.pubkey is not None + ] + + def clear_keydata(self, peer_id: ID) -> None: + """Clears all the keys of the peer""" + peer_data = self.peer_data_map[peer_id] + peer_data.clear_keydata() + + def record_latency(self, peer_id: ID, RTT: float) -> None: + """ + Records a new latency measurement for the given peer + using Exponentially Weighted Moving Average (EWMA) + + :param peer_id: peer ID to get private key for + :param RTT: the new latency value (round trip time) + """ + peer_data = self.peer_data_map[peer_id] + peer_data.record_latency(RTT) + + def latency_EWMA(self, peer_id: ID) -> float: + """ + :param peer_id: peer ID to get private key for + :return: The latency EWMA value for that peer + """ + peer_data = self.peer_data_map[peer_id] + return peer_data.latency_EWMA() + + def clear_metrics(self, peer_id: ID) -> None: + """Clear the latency metrics""" + peer_data = self.peer_data_map[peer_id] + peer_data.clear_metrics() + class PeerStoreError(KeyError): """Raised when peer ID is not found in peer store.""" diff --git a/libp2p/tools/async_service/base.py b/libp2p/tools/async_service/base.py index a23f0e75..bd6c3ef0 100644 --- a/libp2p/tools/async_service/base.py +++ b/libp2p/tools/async_service/base.py @@ -18,7 +18,6 @@ import sys from typing import ( Any, TypeVar, - cast, ) import uuid @@ -361,7 +360,7 @@ class BaseManager(InternalManagerAPI): # Only show stacktrace if this is **not** a DaemonTaskExit error exc_info=not isinstance(err, DaemonTaskExit), ) - self._errors.append(cast(EXC_INFO, sys.exc_info())) + self._errors.append(sys.exc_info()) self.cancel() else: if task.parent is None: diff --git a/libp2p/tools/async_service/trio_service.py b/libp2p/tools/async_service/trio_service.py index 3fdddb81..61b5cb7a 100644 --- a/libp2p/tools/async_service/trio_service.py +++ b/libp2p/tools/async_service/trio_service.py @@ -52,7 +52,6 @@ from .exceptions import ( LifecycleError, ) from .typing import ( - EXC_INFO, AsyncFn, ) @@ -232,7 +231,7 @@ class TrioManager(BaseManager): # Exceptions from any tasks spawned by our service will be # caught by trio and raised here, so we store them to report # together with any others we have already captured. - self._errors.append(cast(EXC_INFO, sys.exc_info())) + self._errors.append(sys.exc_info()) finally: system_nursery.cancel_scope.cancel() From 5de458482c2676830d93c0dbbd43a928baa2af25 Mon Sep 17 00:00:00 2001 From: lla-dane Date: Wed, 18 Jun 2025 14:22:21 +0530 Subject: [PATCH 14/23] refactor after rebase --- libp2p/abc.py | 30 +-------- libp2p/peer/peerdata.py | 73 ++++++++++------------ libp2p/peer/peerstore.py | 17 ++++- libp2p/tools/async_service/base.py | 3 +- libp2p/tools/async_service/trio_service.py | 3 +- 5 files changed, 56 insertions(+), 70 deletions(-) diff --git a/libp2p/abc.py b/libp2p/abc.py index 0b63ed04..343ae0a7 100644 --- a/libp2p/abc.py +++ b/libp2p/abc.py @@ -479,18 +479,6 @@ class IAddrBook(ABC): """ - @abstractmethod - def set_addr(self, peer_id: ID, addr: Multiaddr, ttl: int) -> None: - """Set addr""" - - @abstractmethod - def set_addrs(self, peer_id: ID, addrs: Sequence[Multiaddr], ttl: int) -> None: - """Set addrs""" - - @abstractmethod - def update_addrs(self, peer_id: ID, oldTTL: int, newTTL: int) -> None: - """Update addrs""" - @abstractmethod def addr_stream(self, peer_id: ID) -> None: """Addr stream""" @@ -527,7 +515,7 @@ class IKeyBook(ABC): """peer_with_keys""" @abstractmethod - def clear_keydata(self, peer_id: ID) -> PublicKey: + def clear_keydata(self, peer_id: ID) -> None: """clear_keydata""" @@ -671,18 +659,6 @@ class IPeerStore(IPeerMetadata, IAddrBook, IKeyBook, IMetrics, IProtoBook): """ - @abstractmethod - def set_addr(self, peer_id: ID, addr: Multiaddr, ttl: int) -> None: - """set_addr""" - - @abstractmethod - def set_addrs(self, peer_id: ID, addrs: Sequence[Multiaddr], ttl: int) -> None: - """set_addrs""" - - @abstractmethod - def update_addrs(self, peer_id: ID, oldTTL: int, newTTL: int) -> None: - """update_addrs""" - @abstractmethod def addrs(self, peer_id: ID) -> list[Multiaddr]: """ @@ -835,7 +811,7 @@ class IPeerStore(IPeerMetadata, IAddrBook, IKeyBook, IMetrics, IProtoBook): """peer_with_keys""" @abstractmethod - def clear_keydata(self, peer_id: ID) -> PublicKey: + def clear_keydata(self, peer_id: ID) -> None: """clear_keydata""" ## @@ -1488,7 +1464,7 @@ class IPeerData(ABC): """ @abstractmethod - def add_addrs(self, addrs: Sequence[Multiaddr], ttl: int) -> None: + def add_addrs(self, addrs: Sequence[Multiaddr]) -> None: """ Add multiple multiaddresses to the peer's data. diff --git a/libp2p/peer/peerdata.py b/libp2p/peer/peerdata.py index bd0c4d0b..bf54a494 100644 --- a/libp2p/peer/peerdata.py +++ b/libp2p/peer/peerdata.py @@ -34,6 +34,7 @@ class PeerData(IPeerData): addrs: list[Multiaddr] last_identified: int ttl: int # Keep ttl=0 by default for always valid + latmap: float def __init__(self) -> None: self.pubkey = None @@ -43,6 +44,7 @@ class PeerData(IPeerData): self.addrs = [] self.last_identified = int(time.time()) self.ttl = 0 + self.latmap = 0 def get_protocols(self) -> list[str]: """ @@ -92,52 +94,13 @@ class PeerData(IPeerData): """Clear all protocols""" self.protocols = [] - def add_addrs(self, addrs: Sequence[Multiaddr], ttl: int) -> None: + def add_addrs(self, addrs: Sequence[Multiaddr]) -> None: """ :param addrs: multiaddresses to add """ - expiry = time.time() + ttl if ttl is not None else float("inf") for addr in addrs: if addr not in self.addrs: self.addrs.append(addr) - current_expiry = self.addrs_ttl.get(addr, 0) - if expiry > current_expiry: - self.addrs_ttl[addr] = expiry - - def set_addrs(self, addrs: Sequence[Multiaddr], ttl: int) -> None: - """ - :param addrs: multiaddresses to update - :param ttl: new ttl - """ - now = time.time() - - if ttl <= 0: - # Put the TTL value to -1 - for addr in addrs: - # TODO! if addr in self.addrs, remove them? - if addr in self.addrs_ttl: - del self.addrs_ttl[addr] - return - - expiry = now + ttl - for addr in addrs: - # TODO! if addr not in self.addrs, add them? - self.addrs_ttl[addr] = expiry - - def update_addrs(self, oldTTL: int, newTTL: int) -> None: - """ - :param oldTTL: old ttl - :param newTTL: new ttl - """ - now = time.time() - - new_expiry = now + newTTL - old_expiry = now + oldTTL - - for addr, expiry in list(self.addrs_ttl.items()): - # Approximate match by expiry time - if abs(expiry - old_expiry) < 1: - self.addrs_ttl[addr] = new_expiry def get_addrs(self) -> list[Multiaddr]: """ @@ -200,6 +163,36 @@ class PeerData(IPeerData): raise PeerDataError("private key not found") return self.privkey + def clear_keydata(self) -> None: + """Clears keydata""" + self.pubkey = None + self.privkey = None + + def record_latency(self, new_latency: float) -> None: + """ + Records a new latency measurement for the given peer + using Exponentially Weighted Moving Average (EWMA) + :param new_latency: the new latency value + """ + s = LATENCY_EWMA_SMOOTHING + if s > 1 or s < 0: + s = 0.1 + + if self.latmap is None: + self.latmap = new_latency + else: + prev = self.latmap + updated = ((1.0 - s) * prev) + (s * new_latency) + self.latmap = updated + + def latency_EWMA(self) -> float: + """Returns the latency EWMA value""" + return self.latmap + + def clear_metrics(self) -> None: + """Clear the latency metrics""" + self.latmap = 0 + def update_last_identified(self) -> None: self.last_identified = int(time.time()) diff --git a/libp2p/peer/peerstore.py b/libp2p/peer/peerstore.py index ada56f47..4539fc87 100644 --- a/libp2p/peer/peerstore.py +++ b/libp2p/peer/peerstore.py @@ -100,7 +100,17 @@ class PeerStore(IPeerStore): """ :return: all of the peer IDs stored in peer store """ - return list(self.peer_data_map.keys()) + peer_data = self.peer_data_map[peer_id] + return peer_data.supports_protocols(protocols) + + def first_supported_protocol(self, peer_id: ID, protocols: Sequence[str]) -> str: + peer_data = self.peer_data_map[peer_id] + return peer_data.first_supported_protocol(protocols) + + def clear_protocol_data(self, peer_id: ID) -> None: + """Clears prtocoldata""" + peer_data = self.peer_data_map[peer_id] + peer_data.clear_protocol_data() def valid_peer_ids(self) -> list[ID]: """ @@ -138,6 +148,11 @@ class PeerStore(IPeerStore): peer_data = self.peer_data_map[peer_id] peer_data.put_metadata(key, val) + def clear_metadata(self, peer_id: ID) -> None: + """Clears metadata""" + peer_data = self.peer_data_map[peer_id] + peer_data.clear_metadata() + def add_addr(self, peer_id: ID, addr: Multiaddr, ttl: int = 0) -> None: """ :param peer_id: peer ID to add address for diff --git a/libp2p/tools/async_service/base.py b/libp2p/tools/async_service/base.py index bd6c3ef0..a23f0e75 100644 --- a/libp2p/tools/async_service/base.py +++ b/libp2p/tools/async_service/base.py @@ -18,6 +18,7 @@ import sys from typing import ( Any, TypeVar, + cast, ) import uuid @@ -360,7 +361,7 @@ class BaseManager(InternalManagerAPI): # Only show stacktrace if this is **not** a DaemonTaskExit error exc_info=not isinstance(err, DaemonTaskExit), ) - self._errors.append(sys.exc_info()) + self._errors.append(cast(EXC_INFO, sys.exc_info())) self.cancel() else: if task.parent is None: diff --git a/libp2p/tools/async_service/trio_service.py b/libp2p/tools/async_service/trio_service.py index 61b5cb7a..3fdddb81 100644 --- a/libp2p/tools/async_service/trio_service.py +++ b/libp2p/tools/async_service/trio_service.py @@ -52,6 +52,7 @@ from .exceptions import ( LifecycleError, ) from .typing import ( + EXC_INFO, AsyncFn, ) @@ -231,7 +232,7 @@ class TrioManager(BaseManager): # Exceptions from any tasks spawned by our service will be # caught by trio and raised here, so we store them to report # together with any others we have already captured. - self._errors.append(sys.exc_info()) + self._errors.append(cast(EXC_INFO, sys.exc_info())) finally: system_nursery.cancel_scope.cancel() From 3d369bc142ac3a0117dfd8426163a47ecdd2515e Mon Sep 17 00:00:00 2001 From: lla-dane Date: Wed, 18 Jun 2025 16:47:31 +0530 Subject: [PATCH 15/23] Proto-Book: added tests --- libp2p/peer/peerdata.py | 9 ++++++ libp2p/peer/peerstore.py | 32 +++++++++++++------- tests/core/peer/test_peerdata.py | 50 ++++++++++++++++++++++++++++++++ 3 files changed, 80 insertions(+), 11 deletions(-) diff --git a/libp2p/peer/peerdata.py b/libp2p/peer/peerdata.py index bf54a494..725ca5a3 100644 --- a/libp2p/peer/peerdata.py +++ b/libp2p/peer/peerdata.py @@ -46,6 +46,8 @@ class PeerData(IPeerData): self.ttl = 0 self.latmap = 0 + # --------PROTO-BOOK-------- + def get_protocols(self) -> list[str]: """ :return: all protocols associated with given peer @@ -94,6 +96,7 @@ class PeerData(IPeerData): """Clear all protocols""" self.protocols = [] + # -------ADDR-BOOK--------- def add_addrs(self, addrs: Sequence[Multiaddr]) -> None: """ :param addrs: multiaddresses to add @@ -112,6 +115,9 @@ class PeerData(IPeerData): """Clear all addresses.""" self.addrs = [] + # TODO! ADDRS_STREAM + + # -------METADATA----------- def put_metadata(self, key: str, val: Any) -> None: """ :param key: key in KV pair @@ -133,6 +139,7 @@ class PeerData(IPeerData): """Clears metadata.""" self.metadata = {} + # -------KEY-BOOK--------------- def add_pubkey(self, pubkey: PublicKey) -> None: """ :param pubkey: @@ -168,6 +175,7 @@ class PeerData(IPeerData): self.pubkey = None self.privkey = None + # ----------METRICS-------------- def record_latency(self, new_latency: float) -> None: """ Records a new latency measurement for the given peer @@ -196,6 +204,7 @@ class PeerData(IPeerData): def update_last_identified(self) -> None: self.last_identified = int(time.time()) + # ----------TTL------------------ def get_last_identified(self) -> int: """ :return: last identified timestamp diff --git a/libp2p/peer/peerstore.py b/libp2p/peer/peerstore.py index 4539fc87..da18a5f0 100644 --- a/libp2p/peer/peerstore.py +++ b/libp2p/peer/peerstore.py @@ -62,6 +62,20 @@ class PeerStore(IPeerStore): def clear_peerdata(self, peer_id: ID) -> None: """Clears the peer data of the peer""" + def valid_peer_ids(self) -> list[ID]: + """ + :return: all of the valid peer IDs stored in peer store + """ + valid_peer_ids: list[ID] = [] + for peer_id, peer_data in self.peer_data_map.items(): + if not peer_data.is_expired(): + valid_peer_ids.append(peer_id) + else: + peer_data.clear_addrs() + return valid_peer_ids + + # --------PROTO-BOOK-------- + def get_protocols(self, peer_id: ID) -> list[str]: """ :param peer_id: peer ID to get protocols for @@ -112,17 +126,7 @@ class PeerStore(IPeerStore): peer_data = self.peer_data_map[peer_id] peer_data.clear_protocol_data() - def valid_peer_ids(self) -> list[ID]: - """ - :return: all of the valid peer IDs stored in peer store - """ - valid_peer_ids: list[ID] = [] - for peer_id, peer_data in self.peer_data_map.items(): - if not peer_data.is_expired(): - valid_peer_ids.append(peer_id) - else: - peer_data.clear_addrs() - return valid_peer_ids + # ------METADATA--------- def get(self, peer_id: ID, key: str) -> Any: """ @@ -153,6 +157,8 @@ class PeerStore(IPeerStore): peer_data = self.peer_data_map[peer_id] peer_data.clear_metadata() + # -------ADDR-BOOK-------- + def add_addr(self, peer_id: ID, addr: Multiaddr, ttl: int = 0) -> None: """ :param peer_id: peer ID to add address for @@ -215,6 +221,8 @@ class PeerStore(IPeerStore): """addr_stream""" # TODO! + # -------KEY-BOOK--------- + def add_pubkey(self, peer_id: ID, pubkey: PublicKey) -> None: """ :param peer_id: peer ID to add public key for @@ -288,6 +296,8 @@ class PeerStore(IPeerStore): peer_data = self.peer_data_map[peer_id] peer_data.clear_keydata() + # --------METRICS-------- + def record_latency(self, peer_id: ID, RTT: float) -> None: """ Records a new latency measurement for the given peer diff --git a/tests/core/peer/test_peerdata.py b/tests/core/peer/test_peerdata.py index 65e98959..500e19d4 100644 --- a/tests/core/peer/test_peerdata.py +++ b/tests/core/peer/test_peerdata.py @@ -39,6 +39,56 @@ def test_set_protocols(): assert peer_data.get_protocols() == protocols +# Test case when removing protocols: +def test_remove_protocols(): + peer_data = PeerData() + protocols: Sequence[str] = ["protocol1", "protocol2"] + peer_data.set_protocols(protocols) + + peer_data.remove_protocols(["protocol1"]) + assert peer_data.get_protocols() == ["protocol2"] + + +# Test case when supports protocols: +def test_supports_protocols(): + peer_data = PeerData() + peer_data.set_protocols(["protocol1", "protocol2", "protocol3"]) + + input_protocols = ["protocol1", "protocol4", "protocol2"] + supported = peer_data.supports_protocols(input_protocols) + + assert supported == ["protocol1", "protocol2"] + + +def test_first_supported_protocol_found(): + peer_data = PeerData() + peer_data.set_protocols(["protocolA", "protocolB"]) + + input_protocols = ["protocolC", "protocolB", "protocolA"] + first = peer_data.first_supported_protocol(input_protocols) + + assert first == "protocolB" + + +def test_first_supported_protocol_none(): + peer_data = PeerData() + peer_data.set_protocols(["protocolX", "protocolY"]) + + input_protocols = ["protocolA", "protocolB"] + first = peer_data.first_supported_protocol(input_protocols) + + assert first == "None supported" + + +def test_clear_protocol_data(): + peer_data = PeerData() + peer_data.set_protocols(["proto1", "proto2"]) + + peer_data.clear_protocol_data() + + assert peer_data.get_protocols() == [] + + # Test case when adding addresses def test_add_addrs(): peer_data = PeerData() From 4e533270793e8dc7b4c3e3120eb27129d7d74b55 Mon Sep 17 00:00:00 2001 From: lla-dane Date: Wed, 18 Jun 2025 17:05:38 +0530 Subject: [PATCH 16/23] Metrics: added tests --- libp2p/peer/peerdata.py | 2 +- tests/core/peer/test_peerdata.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/libp2p/peer/peerdata.py b/libp2p/peer/peerdata.py index 725ca5a3..bf45ca23 100644 --- a/libp2p/peer/peerdata.py +++ b/libp2p/peer/peerdata.py @@ -186,7 +186,7 @@ class PeerData(IPeerData): if s > 1 or s < 0: s = 0.1 - if self.latmap is None: + if self.latmap == 0: self.latmap = new_latency else: prev = self.latmap diff --git a/tests/core/peer/test_peerdata.py b/tests/core/peer/test_peerdata.py index 500e19d4..f2d76643 100644 --- a/tests/core/peer/test_peerdata.py +++ b/tests/core/peer/test_peerdata.py @@ -157,3 +157,35 @@ def test_get_privkey_not_found(): peer_data = PeerData() with pytest.raises(PeerDataError): peer_data.get_privkey() + + +# Test case for recording latency for the first time +def test_record_latency_initial(): + peer_data = PeerData() + assert peer_data.latency_EWMA() == 0 + + peer_data.record_latency(100.0) + assert peer_data.latency_EWMA() == 100.0 + + +# Test case for updating latency +def test_record_latency_updates_ewma(): + peer_data = PeerData() + peer_data.record_latency(100.0) # first measurement + first = peer_data.latency_EWMA() + + peer_data.record_latency(50.0) # second measurement + second = peer_data.latency_EWMA() + + assert second < first # EWMA should have smoothed downward + assert second > 50.0 # Not as low as the new latency + assert second != first + + +def test_clear_metrics(): + peer_data = PeerData() + peer_data.record_latency(200.0) + assert peer_data.latency_EWMA() == 200.0 + + peer_data.clear_metrics() + assert peer_data.latency_EWMA() == 0 From 1b025e552c81187f37971fe59016a72dc15428f9 Mon Sep 17 00:00:00 2001 From: lla-dane Date: Wed, 18 Jun 2025 17:19:07 +0530 Subject: [PATCH 17/23] Key-Book: added tests --- tests/core/peer/test_peerdata.py | 41 ++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/tests/core/peer/test_peerdata.py b/tests/core/peer/test_peerdata.py index f2d76643..24730844 100644 --- a/tests/core/peer/test_peerdata.py +++ b/tests/core/peer/test_peerdata.py @@ -6,10 +6,12 @@ from multiaddr import Multiaddr from libp2p.crypto.secp256k1 import ( create_new_key_pair, ) +from libp2p.peer.id import ID from libp2p.peer.peerdata import ( PeerData, PeerDataError, ) +from libp2p.peer.peerstore import PeerStore MOCK_ADDR = Multiaddr("/ip4/127.0.0.1/tcp/4001") MOCK_KEYPAIR = create_new_key_pair() @@ -60,6 +62,7 @@ def test_supports_protocols(): assert supported == ["protocol1", "protocol2"] +# Test case for first supported protocol is found def test_first_supported_protocol_found(): peer_data = PeerData() peer_data.set_protocols(["protocolA", "protocolB"]) @@ -70,6 +73,7 @@ def test_first_supported_protocol_found(): assert first == "protocolB" +# Test case for first supported protocol not found def test_first_supported_protocol_none(): peer_data = PeerData() peer_data.set_protocols(["protocolX", "protocolY"]) @@ -80,6 +84,7 @@ def test_first_supported_protocol_none(): assert first == "None supported" +# Test case for clearing protocol data def test_clear_protocol_data(): peer_data = PeerData() peer_data.set_protocols(["proto1", "proto2"]) @@ -159,6 +164,42 @@ def test_get_privkey_not_found(): peer_data.get_privkey() +# Test case for returning all the peers with stored keys +def test_peer_with_keys(): + peer_store = PeerStore() + peer_id_1 = ID(b"peer1") + peer_id_2 = ID(b"peer2") + + peer_data_1 = PeerData() + peer_data_2 = PeerData() + + peer_data_1.pubkey = MOCK_PUBKEY + peer_data_2.pubkey = None + + peer_store.peer_data_map = { + peer_id_1: peer_data_1, + peer_id_2: peer_data_2, + } + + assert peer_store.peer_with_keys() == [peer_id_1] + + +# Test case for clearing the key book +def test_clear_keydata(): + peer_store = PeerStore() + peer_id = ID(b"peer123") + peer_data = PeerData() + + peer_data.pubkey = MOCK_PUBKEY + peer_data.privkey = MOCK_PRIVKEY + peer_store.peer_data_map = {peer_id: peer_data} + + peer_store.clear_keydata(peer_id) + + assert peer_data.pubkey is None + assert peer_data.privkey is None + + # Test case for recording latency for the first time def test_record_latency_initial(): peer_data = PeerData() From ff966bbfa027805f4801cb797180213aae3524fb Mon Sep 17 00:00:00 2001 From: lla-dane Date: Wed, 18 Jun 2025 17:24:10 +0530 Subject: [PATCH 18/23] Metadata: added test --- tests/core/peer/test_peerdata.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/core/peer/test_peerdata.py b/tests/core/peer/test_peerdata.py index 24730844..46f9ed5a 100644 --- a/tests/core/peer/test_peerdata.py +++ b/tests/core/peer/test_peerdata.py @@ -136,6 +136,15 @@ def test_get_metadata_key_not_found(): peer_data.get_metadata("nonexistent_key") +# Test case for clearing metadata +def test_clear_metadata(): + peer_data = PeerData() + peer_data.metadata = {"key1": "value1", "key2": "value2"} + + peer_data.clear_metadata() + assert peer_data.metadata == {} + + # Test case for adding public key def test_add_pubkey(): peer_data = PeerData() From 994369705473fe38a8c8bae92d098ff0e0524eac Mon Sep 17 00:00:00 2001 From: lla-dane Date: Wed, 18 Jun 2025 18:52:34 +0530 Subject: [PATCH 19/23] Added docstrings --- libp2p/abc.py | 454 ++++++++++++++++++++++++++++++++++++---- libp2p/peer/peerdata.py | 2 - 2 files changed, 411 insertions(+), 45 deletions(-) diff --git a/libp2p/abc.py b/libp2p/abc.py index 343ae0a7..70c4ab71 100644 --- a/libp2p/abc.py +++ b/libp2p/abc.py @@ -387,7 +387,15 @@ class IPeerMetadata(ABC): @abstractmethod def clear_metadata(self, peer_id: ID) -> None: - """Clears the metadata""" + """ + Remove all stored metadata for the specified peer. + + Parameters + ---------- + peer_id : ID + The peer identifier whose metadata are to be removed. + + """ # -------------------------- addrbook interface.py -------------------------- @@ -479,96 +487,267 @@ class IAddrBook(ABC): """ - @abstractmethod - def addr_stream(self, peer_id: ID) -> None: - """Addr stream""" - # -------------------------- keybook interface.py -------------------------- class IKeyBook(ABC): - """IKeyBook""" + """ + Interface for an key book. + + Provides methods for managing cryptographic keys. + """ @abstractmethod def pubkey(self, peer_id: ID) -> PublicKey: - """Pubkey""" + """ + Returns the public key of the specified peer + + Parameters + ---------- + peer_id : ID + The peer identifier whose public key is to be returned. + + """ @abstractmethod def privkey(self, peer_id: ID) -> PrivateKey: - """Privkey""" + """ + Returns the private key of the specified peer + + Parameters + ---------- + peer_id : ID + The peer identifier whose private key is to be returned. + + """ @abstractmethod def add_pubkey(self, peer_id: ID, pubkey: PublicKey) -> None: - """add_pubkey""" + """ + Adds the public key for a specified peer + + Parameters + ---------- + peer_id : ID + The peer identifier whose public key is to be added + pubkey: PublicKey + The public key of the peer + + """ @abstractmethod def add_privkey(self, peer_id: ID, privkey: PrivateKey) -> None: - """add_privkey""" + """ + Adds the private key for a specified peer + + Parameters + ---------- + peer_id : ID + The peer identifier whose private key is to be added + privkey: PrivateKey + The private key of the peer + + """ @abstractmethod def add_key_pair(self, peer_id: ID, key_pair: KeyPair) -> None: - """add_key_pair""" + """ + Adds the key pair for a specified peer + + Parameters + ---------- + peer_id : ID + The peer identifier whose key pair is to be added + key_pair: KeyPair + The key pair of the peer + + """ @abstractmethod def peer_with_keys(self) -> list[ID]: - """peer_with_keys""" + """Returns all the peer IDs stored in the AddrBook""" @abstractmethod def clear_keydata(self, peer_id: ID) -> None: - """clear_keydata""" + """ + Remove all stored keydata for the specified peer. + + Parameters + ---------- + peer_id : ID + The peer identifier whose keys are to be removed. + + """ # -------------------------- metrics interface.py -------------------------- class IMetrics(ABC): - """IMetrics""" + """ + Interface for metrics of peer interaction. + + Provides methods for managing the metrics. + """ @abstractmethod def record_latency(self, peer_id: ID, RTT: float) -> None: - """record_latency""" + """ + Records a new round-trip time (RTT) latency value for the specified peer + using Exponentially Weighted Moving Average (EWMA). + + Parameters + ---------- + peer_id : ID + The identifier of the peer for which latency is being recorded. + + RTT : float + The round-trip time latency value to record. + + """ @abstractmethod def latency_EWMA(self, peer_id: ID) -> float: - """latency_EWMA""" + """ + Returns the current latency value for the specified peer using + Exponentially Weighted Moving Average (EWMA). + + Parameters + ---------- + peer_id : ID + The identifier of the peer whose latency EWMA is to be returned. + + """ @abstractmethod def clear_metrics(self, peer_id: ID) -> None: - """clear_metrics""" + """ + Clears the stored latency metrics for the specified peer. + + Parameters + ---------- + peer_id : ID + The identifier of the peer whose latency metrics are to be cleared. + + """ # -------------------------- protobook interface.py -------------------------- class IProtoBook(ABC): + """ + Interface for a protocol book. + + Provides methods for managing the list of supported protocols. + """ + @abstractmethod def get_protocols(self, peer_id: ID) -> list[str]: - """get_protocols""" + """ + Returns the list of protocols associated with the specified peer. + + Parameters + ---------- + peer_id : ID + The identifier of the peer whose supported protocols are to be returned. + + """ @abstractmethod def add_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None: - """add_protocols""" + """ + Adds the given protocols to the specified peer's protocol list. + + Parameters + ---------- + peer_id : ID + The identifier of the peer to which protocols will be added. + + protocols : Sequence[str] + A sequence of protocol strings to add. + + """ @abstractmethod def set_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None: - """set_protocols""" + """ + Replaces the existing protocols of the specified peer with the given list. + + Parameters + ---------- + peer_id : ID + The identifier of the peer whose protocols are to be set. + + protocols : Sequence[str] + A sequence of protocol strings to assign. + + """ @abstractmethod def remove_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None: - """remove_protocols""" + """ + Removes the specified protocols from the peer's protocol list. + + Parameters + ---------- + peer_id : ID + The identifier of the peer from which protocols will be removed. + + protocols : Sequence[str] + A sequence of protocol strings to remove. + + """ @abstractmethod def supports_protocols(self, peer_id: ID, protocols: Sequence[str]) -> list[str]: - """supports_protocols""" + """ + Returns the list of protocols from the input sequence that the peer supports. + + Parameters + ---------- + peer_id : ID + The identifier of the peer to check for protocol support. + + protocols : Sequence[str] + A sequence of protocol strings to check against the peer's + supported protocols. + + """ @abstractmethod def first_supported_protocol(self, peer_id: ID, protocols: Sequence[str]) -> str: - """first_supported_protocol""" + """ + Returns the first protocol from the input list that the peer supports. + + Parameters + ---------- + peer_id : ID + The identifier of the peer to check for supported protocols. + + protocols : Sequence[str] + A sequence of protocol strings to check. + + Returns + ------- + str + The first matching protocol string, or an empty string + if none are supported. + + """ @abstractmethod def clear_protocol_data(self, peer_id: ID) -> None: - """clear_protocol_data""" + """ + Clears all protocol data associated with the specified peer. + + Parameters + ---------- + peer_id : ID + The identifier of the peer whose protocol data will be cleared. + + """ # -------------------------- peerstore interface.py -------------------------- @@ -582,6 +761,7 @@ class IPeerStore(IPeerMetadata, IAddrBook, IKeyBook, IMetrics, IProtoBook): management, protocol handling, and key storage. """ + # -------METADATA--------- @abstractmethod def get(self, peer_id: ID, key: str) -> Any: """ @@ -624,9 +804,17 @@ class IPeerStore(IPeerMetadata, IAddrBook, IKeyBook, IMetrics, IProtoBook): @abstractmethod def clear_metadata(self, peer_id: ID) -> None: - """clear_metadata""" + """ + Clears the stored latency metrics for the specified peer. - ## + Parameters + ---------- + peer_id : ID + The identifier of the peer whose latency metrics are to be cleared. + + """ + + # --------ADDR-BOOK--------- @abstractmethod def add_addr(self, peer_id: ID, addr: Multiaddr, ttl: int) -> None: """ @@ -700,11 +888,7 @@ class IPeerStore(IPeerMetadata, IAddrBook, IKeyBook, IMetrics, IProtoBook): """ - @abstractmethod - def addr_stream(self, peer_id: ID) -> None: - """addr_stream""" - - ## + # --------KEY-BOOK---------- @abstractmethod def pubkey(self, peer_id: ID) -> PublicKey: """ @@ -808,26 +992,63 @@ class IPeerStore(IPeerMetadata, IAddrBook, IKeyBook, IMetrics, IProtoBook): @abstractmethod def peer_with_keys(self) -> list[ID]: - """peer_with_keys""" + """Returns all the peer IDs stored in the AddrBook""" @abstractmethod def clear_keydata(self, peer_id: ID) -> None: - """clear_keydata""" + """ + Remove all stored keydata for the specified peer. - ## + Parameters + ---------- + peer_id : ID + The peer identifier whose keys are to be removed. + + """ + + # -------METRICS--------- @abstractmethod def record_latency(self, peer_id: ID, RTT: float) -> None: - """record_latency""" + """ + Records a new round-trip time (RTT) latency value for the specified peer + using Exponentially Weighted Moving Average (EWMA). + + Parameters + ---------- + peer_id : ID + The identifier of the peer for which latency is being recorded. + + RTT : float + The round-trip time latency value to record. + + """ @abstractmethod def latency_EWMA(self, peer_id: ID) -> float: - """latency_EWMA""" + """ + Returns the current latency value for the specified peer using + Exponentially Weighted Moving Average (EWMA). + + Parameters + ---------- + peer_id : ID + The identifier of the peer whose latency EWMA is to be returned. + + """ @abstractmethod def clear_metrics(self, peer_id: ID) -> None: - """clear_metrics""" + """ + Clears the stored latency metrics for the specified peer. - ## + Parameters + ---------- + peer_id : ID + The identifier of the peer whose latency metrics are to be cleared. + + """ + + # --------PROTO-BOOK---------- @abstractmethod def get_protocols(self, peer_id: ID) -> list[str]: """ @@ -880,21 +1101,69 @@ class IPeerStore(IPeerMetadata, IAddrBook, IKeyBook, IMetrics, IProtoBook): @abstractmethod def remove_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None: - """remove_protocols""" + """ + Removes the specified protocols from the peer's protocol list. + + Parameters + ---------- + peer_id : ID + The identifier of the peer from which protocols will be removed. + + protocols : Sequence[str] + A sequence of protocol strings to remove. + + """ @abstractmethod def supports_protocols(self, peer_id: ID, protocols: Sequence[str]) -> list[str]: - """supports_protocols""" + """ + Returns the list of protocols from the input sequence that the peer supports. + + Parameters + ---------- + peer_id : ID + The identifier of the peer to check for protocol support. + + protocols : Sequence[str] + A sequence of protocol strings to check against the peer's + supported protocols. + + """ @abstractmethod def first_supported_protocol(self, peer_id: ID, protocols: Sequence[str]) -> str: - """first_supported_protocol""" + """ + Returns the first protocol from the input list that the peer supports. + + Parameters + ---------- + peer_id : ID + The identifier of the peer to check for supported protocols. + + protocols : Sequence[str] + A sequence of protocol strings to check. + + Returns + ------- + str + The first matching protocol string, or an empty string + if none are supported. + + """ @abstractmethod def clear_protocol_data(self, peer_id: ID) -> None: - """clear_protocol_data""" + """ + Clears all protocol data associated with the specified peer. - ## + Parameters + ---------- + peer_id : ID + The identifier of the peer whose protocol data will be cleared. + + """ + + # --------PEER-STORE-------- @abstractmethod def peer_info(self, peer_id: ID) -> PeerInfo: """ @@ -1463,6 +1732,60 @@ class IPeerData(ABC): """ + @abstractmethod + def remove_protocols(self, protocols: Sequence[str]) -> None: + """ + Removes the specified protocols from this peer's list of supported protocols. + + Parameters + ---------- + protocols : Sequence[str] + A sequence of protocol strings to be removed. + + """ + + @abstractmethod + def supports_protocols(self, protocols: Sequence[str]) -> list[str]: + """ + Returns the list of protocols from the input sequence that are supported + by this peer. + + Parameters + ---------- + protocols : Sequence[str] + A sequence of protocol strings to check against this peer's supported + protocols. + + Returns + ------- + list[str] + A list of protocol strings that are supported. + + """ + + @abstractmethod + def first_supported_protocol(self, protocols: Sequence[str]) -> str: + """ + Returns the first protocol from the input list that this peer supports. + + Parameters + ---------- + protocols : Sequence[str] + A sequence of protocol strings to check for support. + + Returns + ------- + str + The first matching protocol, or an empty string if none are supported. + + """ + + @abstractmethod + def clear_protocol_data(self) -> None: + """ + Clears all protocol data associated with this peer. + """ + @abstractmethod def add_addrs(self, addrs: Sequence[Multiaddr]) -> None: """ @@ -1532,6 +1855,12 @@ class IPeerData(ABC): """ + @abstractmethod + def clear_metadata(self) -> None: + """ + Clears all metadata entries associated with this peer. + """ + @abstractmethod def add_pubkey(self, pubkey: PublicKey) -> None: """ @@ -1590,6 +1919,45 @@ class IPeerData(ABC): """ + @abstractmethod + def clear_keydata(self) -> None: + """ + Clears all cryptographic key data associated with this peer, + including both public and private keys. + """ + + @abstractmethod + def record_latency(self, new_latency: float) -> None: + """ + Records a new latency measurement using + Exponentially Weighted Moving Average (EWMA). + + Parameters + ---------- + new_latency : float + The new round-trip time (RTT) latency value to incorporate + into the EWMA calculation. + + """ + + @abstractmethod + def latency_EWMA(self) -> float: + """ + Returns the current EWMA value of the recorded latency. + + Returns + ------- + float + The current latency estimate based on EWMA. + + """ + + @abstractmethod + def clear_metrics(self) -> None: + """ + Clears all latency-related metrics and resets the internal state. + """ + @abstractmethod def update_last_identified(self) -> None: """ diff --git a/libp2p/peer/peerdata.py b/libp2p/peer/peerdata.py index bf45ca23..0d1a2f35 100644 --- a/libp2p/peer/peerdata.py +++ b/libp2p/peer/peerdata.py @@ -115,8 +115,6 @@ class PeerData(IPeerData): """Clear all addresses.""" self.addrs = [] - # TODO! ADDRS_STREAM - # -------METADATA----------- def put_metadata(self, key: str, val: Any) -> None: """ From faeacf686a99bd8a40ed838bc281350d632aabb1 Mon Sep 17 00:00:00 2001 From: lla-dane Date: Wed, 18 Jun 2025 19:34:40 +0530 Subject: [PATCH 20/23] fix typos --- libp2p/peer/peerstore.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/libp2p/peer/peerstore.py b/libp2p/peer/peerstore.py index da18a5f0..e23e014a 100644 --- a/libp2p/peer/peerstore.py +++ b/libp2p/peer/peerstore.py @@ -217,10 +217,6 @@ class PeerStore(IPeerStore): peer_data.clear_addrs() return output - def addr_stream(self, peer_id: ID) -> None: - """addr_stream""" - # TODO! - # -------KEY-BOOK--------- def add_pubkey(self, peer_id: ID, pubkey: PublicKey) -> None: @@ -292,7 +288,7 @@ class PeerStore(IPeerStore): ] def clear_keydata(self, peer_id: ID) -> None: - """Clears all the keys of the peer""" + """Clears the keys of the peer""" peer_data = self.peer_data_map[peer_id] peer_data.clear_keydata() From 51c08de1bce4ae2262002bf0c90d8b0bb8e065c7 Mon Sep 17 00:00:00 2001 From: lla-dane Date: Thu, 3 Jul 2025 18:13:26 +0530 Subject: [PATCH 21/23] test added: clear protocol data --- tests/core/peer/test_peerdata.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/core/peer/test_peerdata.py b/tests/core/peer/test_peerdata.py index 46f9ed5a..49825915 100644 --- a/tests/core/peer/test_peerdata.py +++ b/tests/core/peer/test_peerdata.py @@ -51,6 +51,16 @@ def test_remove_protocols(): assert peer_data.get_protocols() == ["protocol2"] +# Test case when clearing the protocol list: +def test_clear_protocol_data(): + peer_data = PeerData() + protocols: Sequence[str] = ["protocol1", "protocol2"] + peer_data.set_protocols(protocols) + + peer_data.clear_protocol_data() + assert peer_data.get_protocols() == [] + + # Test case when supports protocols: def test_supports_protocols(): peer_data = PeerData() @@ -84,16 +94,6 @@ def test_first_supported_protocol_none(): assert first == "None supported" -# Test case for clearing protocol data -def test_clear_protocol_data(): - peer_data = PeerData() - peer_data.set_protocols(["proto1", "proto2"]) - - peer_data.clear_protocol_data() - - assert peer_data.get_protocols() == [] - - # Test case when adding addresses def test_add_addrs(): peer_data = PeerData() From d1c31483bd02445266396e808035be6d8afc4c7a Mon Sep 17 00:00:00 2001 From: lla-dane Date: Fri, 4 Jul 2025 14:53:44 +0530 Subject: [PATCH 22/23] Implemented addr_stream in the peerstore --- libp2p/peer/peerstore.py | 30 +++++++++++++++++++++++++++ tests/core/peer/test_peerstore.py | 34 +++++++++++++++++++++++++++++++ 2 files changed, 64 insertions(+) diff --git a/libp2p/peer/peerstore.py b/libp2p/peer/peerstore.py index e23e014a..40cb7893 100644 --- a/libp2p/peer/peerstore.py +++ b/libp2p/peer/peerstore.py @@ -2,6 +2,7 @@ from collections import ( defaultdict, ) from collections.abc import ( + AsyncIterable, Sequence, ) from typing import ( @@ -11,6 +12,8 @@ from typing import ( from multiaddr import ( Multiaddr, ) +import trio +from trio import MemoryReceiveChannel, MemorySendChannel from libp2p.abc import ( IPeerStore, @@ -40,6 +43,7 @@ class PeerStore(IPeerStore): def __init__(self) -> None: self.peer_data_map = defaultdict(PeerData) + self.addr_update_channels: dict[ID, MemorySendChannel[Multiaddr]] = {} def peer_info(self, peer_id: ID) -> PeerInfo: """ @@ -178,6 +182,13 @@ class PeerStore(IPeerStore): peer_data.set_ttl(ttl) peer_data.update_last_identified() + if peer_id in self.addr_update_channels: + for addr in addrs: + try: + self.addr_update_channels[peer_id].send_nowait(addr) + except trio.WouldBlock: + pass # Or consider logging / dropping / replacing stream + def addrs(self, peer_id: ID) -> list[Multiaddr]: """ :param peer_id: peer ID to get addrs for @@ -217,6 +228,25 @@ class PeerStore(IPeerStore): peer_data.clear_addrs() return output + async def addr_stream(self, peer_id: ID) -> AsyncIterable[Multiaddr]: + """ + Returns an async stream of newly added addresses for the given peer. + + This function allows consumers to subscribe to address updates for a peer + and receive each new address as it is added via `add_addr` or `add_addrs`. + + :param peer_id: The ID of the peer to monitor address updates for. + :return: An async iterator yielding Multiaddr instances as they are added. + """ + send: MemorySendChannel[Multiaddr] + receive: MemoryReceiveChannel[Multiaddr] + + send, receive = trio.open_memory_channel(0) + self.addr_update_channels[peer_id] = send + + async for addr in receive: + yield addr + # -------KEY-BOOK--------- def add_pubkey(self, peer_id: ID, pubkey: PublicKey) -> None: diff --git a/tests/core/peer/test_peerstore.py b/tests/core/peer/test_peerstore.py index b0d8ed81..85fc1863 100644 --- a/tests/core/peer/test_peerstore.py +++ b/tests/core/peer/test_peerstore.py @@ -2,6 +2,7 @@ import time import pytest from multiaddr import Multiaddr +import trio from libp2p.peer.id import ID from libp2p.peer.peerstore import ( @@ -89,3 +90,36 @@ def test_peers(): store.add_addr(ID(b"peer3"), Multiaddr("/ip4/127.0.0.1/tcp/4001"), 10) assert set(store.peer_ids()) == {ID(b"peer1"), ID(b"peer2"), ID(b"peer3")} + + +@pytest.mark.trio +async def test_addr_stream_yields_new_addrs(): + store = PeerStore() + peer_id = ID(b"peer1") + addr1 = Multiaddr("/ip4/127.0.0.1/tcp/4001") + addr2 = Multiaddr("/ip4/127.0.0.1/tcp/4002") + + # 🔧 Pre-initialize peer in peer_data_map + # store.add_addr(peer_id, Multiaddr("/ip4/127.0.0.1/tcp/0"), ttl=1) + + collected = [] + + async def consume_addrs(): + async for addr in store.addr_stream(peer_id): + collected.append(addr) + if len(collected) == 2: + break + + async with trio.open_nursery() as nursery: + nursery.start_soon(consume_addrs) + await trio.sleep(2) # Give time for the stream to start + + store.add_addr(peer_id, addr1, ttl=10) + await trio.sleep(0.2) + store.add_addr(peer_id, addr2, ttl=10) + await trio.sleep(0.2) + + # After collecting expected addresses, cancel the stream + nursery.cancel_scope.cancel() + + assert collected == [addr1, addr2] From b21591f8d591a2563c706139669f176ed9f6b9f5 Mon Sep 17 00:00:00 2001 From: lla-dane Date: Fri, 4 Jul 2025 15:07:12 +0530 Subject: [PATCH 23/23] remove redundants --- tests/core/peer/test_peerstore.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/core/peer/test_peerstore.py b/tests/core/peer/test_peerstore.py index 85fc1863..c5f31767 100644 --- a/tests/core/peer/test_peerstore.py +++ b/tests/core/peer/test_peerstore.py @@ -99,9 +99,6 @@ async def test_addr_stream_yields_new_addrs(): addr1 = Multiaddr("/ip4/127.0.0.1/tcp/4001") addr2 = Multiaddr("/ip4/127.0.0.1/tcp/4002") - # 🔧 Pre-initialize peer in peer_data_map - # store.add_addr(peer_id, Multiaddr("/ip4/127.0.0.1/tcp/0"), ttl=1) - collected = [] async def consume_addrs():