mirror of
https://github.com/varun-r-mallya/py-libp2p.git
synced 2026-02-10 23:20:55 +00:00
feat: Matching go-libp2p PeerStore implementation
This commit is contained in:
@ -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:
|
||||
|
||||
@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user