feat: Matching go-libp2p PeerStore implementation

This commit is contained in:
lla-dane
2025-06-01 23:02:41 +05:30
parent e6f96d32e2
commit f3d8cbf968
5 changed files with 535 additions and 225 deletions

View File

@ -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
"""

View File

@ -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:

View File

@ -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."""

View File

@ -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:

View File

@ -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()