From f3d8cbf9689f233d269918c225e84cebfcba1494 Mon Sep 17 00:00:00 2001 From: lla-dane Date: Sun, 1 Jun 2025 23:02:41 +0530 Subject: [PATCH] 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()