chore: use AddrInfo instead of ID to match nwaku's libwaku ping function

This commit is contained in:
Richard Ramos 2024-10-28 10:55:21 -04:00
parent c78b09d4ca
commit fc540c780e
No known key found for this signature in database
GPG Key ID: 1CE87DB518195760
3 changed files with 23 additions and 16 deletions

View File

@ -6,11 +6,12 @@ import (
"github.com/libp2p/go-libp2p/core/host" "github.com/libp2p/go-libp2p/core/host"
"github.com/libp2p/go-libp2p/core/peer" "github.com/libp2p/go-libp2p/core/peer"
"github.com/libp2p/go-libp2p/core/peerstore"
"github.com/libp2p/go-libp2p/p2p/protocol/ping" "github.com/libp2p/go-libp2p/p2p/protocol/ping"
) )
type Pinger interface { type Pinger interface {
PingPeer(ctx context.Context, peerID peer.ID) (time.Duration, error) PingPeer(ctx context.Context, peerInfo peer.AddrInfo) (time.Duration, error)
} }
type defaultPingImpl struct { type defaultPingImpl struct {
@ -23,8 +24,9 @@ func NewDefaultPinger(host host.Host) Pinger {
} }
} }
func (d *defaultPingImpl) PingPeer(ctx context.Context, peerID peer.ID) (time.Duration, error) { func (d *defaultPingImpl) PingPeer(ctx context.Context, peerInfo peer.AddrInfo) (time.Duration, error) {
pingResultCh := ping.Ping(ctx, d.host, peerID) d.host.Peerstore().AddAddrs(peerInfo.ID, peerInfo.Addrs, peerstore.AddressTTL)
pingResultCh := ping.Ping(ctx, d.host, peerInfo.ID)
select { select {
case <-ctx.Done(): case <-ctx.Done():
return 0, ctx.Err() return 0, ctx.Err()

View File

@ -45,8 +45,8 @@ type peerStatus struct {
type StorenodeConfigProvider interface { type StorenodeConfigProvider interface {
UseStorenodes() (bool, error) UseStorenodes() (bool, error)
GetPinnedStorenode() (peer.ID, error) GetPinnedStorenode() (peer.AddrInfo, error)
Storenodes() ([]peer.ID, error) Storenodes() ([]peer.AddrInfo, error)
} }
type StorenodeCycle struct { type StorenodeCycle struct {
@ -104,7 +104,7 @@ func (m *StorenodeCycle) connectToNewStorenodeAndWait(ctx context.Context) error
} }
// If no pinned storenode, no need to disconnect and wait for it to be available // If no pinned storenode, no need to disconnect and wait for it to be available
if pinnedStorenode == "" { if pinnedStorenode.ID == "" {
m.disconnectActiveStorenode(graylistBackoff) m.disconnectActiveStorenode(graylistBackoff)
} }
@ -180,21 +180,26 @@ func poolSize(fleetSize int) int {
return int(math.Ceil(float64(fleetSize) / 4)) return int(math.Ceil(float64(fleetSize) / 4))
} }
func (m *StorenodeCycle) getAvailableStorenodesSortedByRTT(ctx context.Context, allStorenodes []peer.ID) []peer.ID { func (m *StorenodeCycle) getAvailableStorenodesSortedByRTT(ctx context.Context, allStorenodes []peer.AddrInfo) []peer.AddrInfo {
peerIDToInfo := make(map[peer.ID]peer.AddrInfo)
for _, p := range allStorenodes {
peerIDToInfo[p.ID] = p
}
availableStorenodes := make(map[peer.ID]time.Duration) availableStorenodes := make(map[peer.ID]time.Duration)
availableStorenodesMutex := sync.Mutex{} availableStorenodesMutex := sync.Mutex{}
availableStorenodesWg := sync.WaitGroup{} availableStorenodesWg := sync.WaitGroup{}
for _, storenode := range allStorenodes { for _, storenode := range allStorenodes {
availableStorenodesWg.Add(1) availableStorenodesWg.Add(1)
go func(peerID peer.ID) { go func(peerInfo peer.AddrInfo) {
defer availableStorenodesWg.Done() defer availableStorenodesWg.Done()
ctx, cancel := context.WithTimeout(ctx, 4*time.Second) ctx, cancel := context.WithTimeout(ctx, 4*time.Second)
defer cancel() defer cancel()
rtt, err := m.pinger.PingPeer(ctx, peerID) rtt, err := m.pinger.PingPeer(ctx, peerInfo)
if err == nil { // pinging storenodes might fail, but we don't care if err == nil { // pinging storenodes might fail, but we don't care
availableStorenodesMutex.Lock() availableStorenodesMutex.Lock()
availableStorenodes[peerID] = rtt availableStorenodes[peerInfo.ID] = rtt
availableStorenodesMutex.Unlock() availableStorenodesMutex.Unlock()
} }
}(storenode) }(storenode)
@ -209,7 +214,7 @@ func (m *StorenodeCycle) getAvailableStorenodesSortedByRTT(ctx context.Context,
var sortedStorenodes []SortedStorenode var sortedStorenodes []SortedStorenode
for storenodeID, rtt := range availableStorenodes { for storenodeID, rtt := range availableStorenodes {
sortedStorenode := SortedStorenode{ sortedStorenode := SortedStorenode{
Storenode: storenodeID, Storenode: peerIDToInfo[storenodeID],
RTT: rtt, RTT: rtt,
} }
m.peersMutex.Lock() m.peersMutex.Lock()
@ -222,7 +227,7 @@ func (m *StorenodeCycle) getAvailableStorenodesSortedByRTT(ctx context.Context,
} }
sort.Sort(byRTTMsAndCanConnectBefore(sortedStorenodes)) sort.Sort(byRTTMsAndCanConnectBefore(sortedStorenodes))
result := make([]peer.ID, len(sortedStorenodes)) result := make([]peer.AddrInfo, len(sortedStorenodes))
for i, s := range sortedStorenodes { for i, s := range sortedStorenodes {
result[i] = s.Storenode result[i] = s.Storenode
} }
@ -252,8 +257,8 @@ func (m *StorenodeCycle) findNewStorenode(ctx context.Context) error {
return err return err
} }
if pinnedStorenode != "" { if pinnedStorenode.ID != "" {
return m.setActiveStorenode(pinnedStorenode) return m.setActiveStorenode(pinnedStorenode.ID)
} }
m.logger.Info("Finding a new storenode..") m.logger.Info("Finding a new storenode..")
@ -287,7 +292,7 @@ func (m *StorenodeCycle) findNewStorenode(ctx context.Context) error {
} }
ms := allStorenodes[r.Int64()] ms := allStorenodes[r.Int64()]
return m.setActiveStorenode(ms) return m.setActiveStorenode(ms.ID)
} }
func (m *StorenodeCycle) storenodeStatus(peerID peer.ID) connStatus { func (m *StorenodeCycle) storenodeStatus(peerID peer.ID) connStatus {

View File

@ -7,7 +7,7 @@ import (
) )
type SortedStorenode struct { type SortedStorenode struct {
Storenode peer.ID Storenode peer.AddrInfo
RTT time.Duration RTT time.Duration
CanConnectAfter time.Time CanConnectAfter time.Time
} }