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

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