diff --git a/libp2p/abc.py b/libp2p/abc.py index dc941c43..70c4ab71 100644 --- a/libp2p/abc.py +++ b/libp2p/abc.py @@ -385,6 +385,18 @@ class IPeerMetadata(ABC): :raises Exception: If the operation is unsuccessful. """ + @abstractmethod + def clear_metadata(self, peer_id: ID) -> None: + """ + Remove all stored metadata for the specified peer. + + Parameters + ---------- + peer_id : ID + The peer identifier whose metadata are to be removed. + + """ + # -------------------------- addrbook interface.py -------------------------- @@ -476,10 +488,272 @@ class IAddrBook(ABC): """ +# -------------------------- keybook interface.py -------------------------- + + +class IKeyBook(ABC): + """ + Interface for an key book. + + Provides methods for managing cryptographic keys. + """ + + @abstractmethod + def pubkey(self, peer_id: ID) -> PublicKey: + """ + 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: + """ + 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: + """ + 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: + """ + 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: + """ + 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]: + """Returns all the peer IDs stored in the AddrBook""" + + @abstractmethod + def clear_keydata(self, peer_id: ID) -> None: + """ + 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): + """ + Interface for metrics of peer interaction. + + Provides methods for managing the metrics. + """ + + @abstractmethod + def record_latency(self, peer_id: ID, RTT: float) -> None: + """ + 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: + """ + 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: + """ + 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]: + """ + 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: + """ + 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: + """ + 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: + """ + 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]: + """ + 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: + """ + 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: + """ + 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 -------------------------- -class IPeerStore(IAddrBook, IPeerMetadata): +class IPeerStore(IPeerMetadata, IAddrBook, IKeyBook, IMetrics, IProtoBook): """ Interface for a peer store. @@ -487,85 +761,7 @@ class IPeerStore(IAddrBook, IPeerMetadata): management, protocol handling, and key storage. """ - @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 get_protocols(self, peer_id: ID) -> list[str]: - """ - Retrieve the protocols associated with the specified peer. - - Parameters - ---------- - peer_id : ID - The identifier of the peer. - - Returns - ------- - list[str] - A list of protocol identifiers. - - Raises - ------ - PeerStoreError - If the peer ID is not found. - - """ - - @abstractmethod - def add_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None: - """ - Add additional protocols for the specified peer. - - Parameters - ---------- - peer_id : ID - The identifier of the peer. - protocols : Sequence[str] - The protocols to add. - - """ - - @abstractmethod - def set_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None: - """ - Set the protocols for the specified peer. - - Parameters - ---------- - peer_id : ID - The identifier of the peer. - protocols : Sequence[str] - The protocols to set. - - """ - - @abstractmethod - def peer_ids(self) -> list[ID]: - """ - Retrieve all peer identifiers stored in the peer store. - - Returns - ------- - list[ID] - A list of all peer IDs in the store. - - """ - + # -------METADATA--------- @abstractmethod def get(self, peer_id: ID, key: str) -> Any: """ @@ -606,6 +802,19 @@ class IPeerStore(IAddrBook, IPeerMetadata): """ + @abstractmethod + def clear_metadata(self, peer_id: ID) -> None: + """ + 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: """ @@ -679,25 +888,7 @@ class IPeerStore(IAddrBook, IPeerMetadata): """ - @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. - - """ - + # --------KEY-BOOK---------- @abstractmethod def pubkey(self, peer_id: ID) -> PublicKey: """ @@ -720,25 +911,6 @@ class IPeerStore(IAddrBook, IPeerMetadata): """ - @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: """ @@ -761,6 +933,44 @@ class IPeerStore(IAddrBook, IPeerMetadata): """ + @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: """ @@ -780,6 +990,213 @@ class IPeerStore(IAddrBook, IPeerMetadata): """ + @abstractmethod + def peer_with_keys(self) -> list[ID]: + """Returns all the peer IDs stored in the AddrBook""" + + @abstractmethod + def clear_keydata(self, peer_id: ID) -> None: + """ + 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: + """ + 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: + """ + 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: + """ + 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]: + """ + Retrieve the protocols associated with the specified peer. + + Parameters + ---------- + peer_id : ID + The identifier of the peer. + + Returns + ------- + list[str] + A list of protocol identifiers. + + Raises + ------ + PeerStoreError + If the peer ID is not found. + + """ + + @abstractmethod + def add_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None: + """ + Add additional protocols for the specified peer. + + Parameters + ---------- + peer_id : ID + The identifier of the peer. + protocols : Sequence[str] + The protocols to add. + + """ + + @abstractmethod + def set_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None: + """ + Set the protocols for the specified peer. + + Parameters + ---------- + peer_id : ID + The identifier of the peer. + protocols : Sequence[str] + The protocols to set. + + """ + + @abstractmethod + def remove_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None: + """ + 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]: + """ + 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: + """ + 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: + """ + 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: + """ + 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]: + """ + Retrieve all peer identifiers stored in the peer store. + + Returns + ------- + list[ID] + A list of all peer IDs in the store. + + """ + + @abstractmethod + def clear_peerdata(self, peer_id: ID) -> None: + """clear_peerdata""" + # -------------------------- listener interface.py -------------------------- @@ -1315,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: """ @@ -1324,6 +1795,8 @@ class IPeerData(ABC): ---------- addrs : Sequence[Multiaddr] A sequence of multiaddresses to add. + ttl: inr + Time to live for the peer record """ @@ -1382,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: """ @@ -1440,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 386e31ef..0d1a2f35 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 @@ -27,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 @@ -36,6 +44,9 @@ class PeerData(IPeerData): self.addrs = [] self.last_identified = int(time.time()) self.ttl = 0 + self.latmap = 0 + + # --------PROTO-BOOK-------- def get_protocols(self) -> list[str]: """ @@ -55,6 +66,37 @@ class PeerData(IPeerData): """ self.protocols = list(protocols) + 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 = [] + + # -------ADDR-BOOK--------- def add_addrs(self, addrs: Sequence[Multiaddr]) -> None: """ :param addrs: multiaddresses to add @@ -73,6 +115,7 @@ class PeerData(IPeerData): """Clear all addresses.""" self.addrs = [] + # -------METADATA----------- def put_metadata(self, key: str, val: Any) -> None: """ :param key: key in KV pair @@ -90,6 +133,11 @@ class PeerData(IPeerData): return self.metadata[key] raise PeerDataError("key not found") + def clear_metadata(self) -> None: + """Clears metadata.""" + self.metadata = {} + + # -------KEY-BOOK--------------- def add_pubkey(self, pubkey: PublicKey) -> None: """ :param pubkey: @@ -120,9 +168,41 @@ 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 + + # ----------METRICS-------------- + 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 == 0: + 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()) + # ----------TTL------------------ def get_last_identified(self) -> int: """ :return: last identified timestamp diff --git a/libp2p/peer/peerstore.py b/libp2p/peer/peerstore.py index 3bb729d2..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: """ @@ -53,6 +57,29 @@ 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 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 @@ -79,23 +106,31 @@ 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 """ - return list(self.peer_data_map.keys()) + peer_data = self.peer_data_map[peer_id] + return peer_data.supports_protocols(protocols) - 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 + 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() + + # ------METADATA--------- def get(self, peer_id: ID, key: str) -> Any: """ @@ -121,6 +156,13 @@ 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() + + # -------ADDR-BOOK-------- + def add_addr(self, peer_id: ID, addr: Multiaddr, ttl: int = 0) -> None: """ :param peer_id: peer ID to add address for @@ -140,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 @@ -165,7 +214,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 +228,27 @@ 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: """ :param peer_id: peer ID to add public key for @@ -239,6 +309,45 @@ 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 the keys of the peer""" + 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 + 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/tests/core/peer/test_peerdata.py b/tests/core/peer/test_peerdata.py index 65e98959..49825915 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() @@ -39,6 +41,59 @@ 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 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() + peer_data.set_protocols(["protocol1", "protocol2", "protocol3"]) + + input_protocols = ["protocol1", "protocol4", "protocol2"] + supported = peer_data.supports_protocols(input_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"]) + + input_protocols = ["protocolC", "protocolB", "protocolA"] + first = peer_data.first_supported_protocol(input_protocols) + + 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"]) + + input_protocols = ["protocolA", "protocolB"] + first = peer_data.first_supported_protocol(input_protocols) + + assert first == "None supported" + + # Test case when adding addresses def test_add_addrs(): peer_data = PeerData() @@ -81,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() @@ -107,3 +171,71 @@ def test_get_privkey_not_found(): peer_data = PeerData() with pytest.raises(PeerDataError): 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() + 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 diff --git a/tests/core/peer/test_peerstore.py b/tests/core/peer/test_peerstore.py index b0d8ed81..c5f31767 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,33 @@ 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") + + 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]