Fix relay peer management & logs

Author: SionoiS

waku/node/peer_manager/peer_manager.nim

@@ -132,15 +132,14 @@ proc addPeer*(pm: PeerManager, remotePeerInfo: RemotePeerInfo, origin = UnknownO
     peerId = remotePeerInfo.peerId,
     addresses = remotePeerInfo.addrs,
     origin = origin,
-    protocols = remotePeerinfo.protocols,
     enr = remotePeerInfo.enr
 
   pm.peerStore[AddressBook][remotePeerInfo.peerId] = remotePeerInfo.addrs
   pm.peerStore[KeyBook][remotePeerInfo.peerId] = publicKey
   pm.peerStore[SourceBook][remotePeerInfo.peerId] = origin
 
-  if (let enr = remotePeerInfo.enr.get(); remotePeerInfo.enr.isSome()):
-    pm.peerStore[ENRBook][remotePeerInfo.peerId] = enr
+  if remotePeerInfo.enr.isSome():
+    pm.peerStore[ENRBook][remotePeerInfo.peerId] = remotePeerInfo.enr.get()
 
   # Add peer to storage. Entry will subsequently be updated with connectedness information
   if not pm.storage.isNil:
@@ -327,24 +326,38 @@ proc onPeerMetadata(pm: PeerManager, peerId: PeerId) {.async.} =
     return
 
   let res = catch: await pm.switch.dial(peerId, WakuMetadataCodec)
+  let conn = res.valueOr:
+    info "disconnecting from peer", peerId=peerId, reason="dial failed: " & error.msg
+    asyncSpawn(pm.switch.disconnect(peerId))
+    pm.peerStore.delete(peerId)
+    return
+  
+  let metadata = (await pm.wakuMetadata.request(conn)).valueOr:
+    info "disconnecting from peer", peerId=peerId, reason="waku metatdata request failed: " & error
+    asyncSpawn(pm.switch.disconnect(peerId))
+    pm.peerStore.delete(peerId)
+    return
 
-  let reason = 
-    if (let conn = res; conn.isOk()):
-      if (let metadata = (await pm.wakuMetadata.request(conn.get())); metadata.isOk()):
-        if (let clusterId = metadata.get().clusterId; clusterId.isSome()):
-          if pm.wakuMetadata.clusterId == clusterId.get():
-            if metadata.get().shards.anyIt(pm.wakuMetadata.shards.contains(it)):
-              return
-            else: "no shards in common"
-          else: "different clusterId reported: " & $pm.wakuMetadata.clusterId & " vs " & $clusterId
-        else: "empty clusterId reported"
-      else: "failed waku metadata codec request"
-    else: "waku metadata codec not supported"
-
-  info "disconnecting from peer", peerId=peerId, reason=reason
-  asyncSpawn(pm.switch.disconnect(peerId))
-  pm.peerStore.delete(peerId)
+  let clusterId = metadata.clusterId.valueOr:
+    info "disconnecting from peer", peerId=peerId, reason="empty clusterId reported"
+    asyncSpawn(pm.switch.disconnect(peerId))
+    pm.peerStore.delete(peerId)
+    return
+
+  if pm.wakuMetadata.clusterId != clusterId:
+    info "disconnecting from peer",
+      peerId=peerId,
+      reason="different clusterId reported: " & $pm.wakuMetadata.clusterId & " vs " & $clusterId
+    asyncSpawn(pm.switch.disconnect(peerId))
+    pm.peerStore.delete(peerId)
+    return
 
+  if not metadata.shards.anyIt(pm.wakuMetadata.shards.contains(it)):
+    info "disconnecting from peer", peerId=peerId, reason="no shard in common"
+    asyncSpawn(pm.switch.disconnect(peerId))
+    pm.peerStore.delete(peerId)
+    return
+   
 # called when a peer i) first connects to us ii) disconnects all connections from us
 proc onPeerEvent(pm: PeerManager, peerId: PeerId, event: PeerEvent) {.async.} =
   if not pm.wakuMetadata.isNil() and event.kind == PeerEventKind.Joined:
@@ -586,9 +599,10 @@ proc connectToNodes*(pm: PeerManager,
   # later.
   await sleepAsync(chronos.seconds(5))
 
-# Returns the peerIds of physical connections (in and out)
-# containing at least one stream with the given protocol.
 proc connectedPeers*(pm: PeerManager, protocol: string): (seq[PeerId], seq[PeerId]) =
+  ## Returns the peerIds of physical connections (in and out)
+  ## containing at least one stream with the given protocol.
+  
   var inPeers: seq[PeerId]
   var outPeers: seq[PeerId]
 
@@ -631,20 +645,22 @@ proc manageRelayPeers*(pm: PeerManager) {.async.} =
   var peersToConnect: seq[RemotePeerInfo]
   var peersToDisconnect: int
 
-  for shard in pm.wakuMetadata.shards.items:
-    var peers = pm.peerStore.getPeersByShard(uint16(pm.wakuMetadata.clusterId), uint16(shard))
+  # Get all connected peers for Waku Relay
+  var (inPeers, outPeers) = pm.connectedPeers(WakuRelayCodec)
 
-    peers.keepItIf(it.protocols.contains(WakuRelayCodec))
+  # Calculate in/out target number of peers for each shards
+  let inTarget = pm.inRelayPeersTarget div pm.wakuMetadata.shards.len
+  let outTarget = pm.outRelayPeersTarget div pm.wakuMetadata.shards.len
 
-    let connectedInPeers = peers.filterIt(
-      pm.peerStore.isConnected(it.peerId) and it.direction == Inbound)
-    
-    let connectedOutPeers = peers.filterIt(
-      pm.peerStore.isConnected(it.peerId) and it.direction == Outbound)
+  for shard in pm.wakuMetadata.shards.items:
+    # Filter out peer not on this shard
+    let connectedInPeers = inPeers.filterIt(
+      pm.peerStore.hasShard(it, uint16(pm.wakuMetadata.clusterId), uint16(shard)))
 
-    let inTarget = pm.inRelayPeersTarget div pm.wakuMetadata.shards.len
-    let outTarget = pm.outRelayPeersTarget div pm.wakuMetadata.shards.len
+    let connectedOutPeers = outPeers.filterIt(
+      pm.peerStore.hasShard(it, uint16(pm.wakuMetadata.clusterId), uint16(shard)))
 
+    # Calculate the difference between current values and targets
     let inPeerDiff = connectedInPeers.len - inTarget
     let outPeerDiff = outTarget - connectedOutPeers.len
 
@@ -654,12 +670,17 @@ proc manageRelayPeers*(pm: PeerManager) {.async.} =
     if outPeerDiff <= 0:
       continue
 
-    let connectablePeers = peers.filterIt(
-      (not pm.peerStore.isConnected(it.peerId)) and pm.canBeConnected(it.peerId))
+    # Get all connectable relay peers for this shard
+    var connectablePeers = pm.peerStore.getPeersByShard(
+      uint16(pm.wakuMetadata.clusterId), uint16(shard))
+  
+    connectablePeers.keepItIf(
+      it.protocols.contains(WakuRelayCodec) and
+      not pm.peerStore.isConnected(it.peerId) and
+      pm.canBeConnected(it.peerId))
 
     debug "Target Peer Count per Shard",
-      Shard = shard,
-      RelayPeers = peers.len,
+      ShardNumber = shard,
       Inbound = $connectedInPeers.len & "/" & $inTarget,
       Outbound = $connectedOutPeers.len & "/" & $outTarget,
       Connectable = connectablePeers.len
@@ -675,6 +696,7 @@ proc manageRelayPeers*(pm: PeerManager) {.async.} =
   # Even with duplicates, after a couple of iteration the target will be reached
   let uniquePeers = peersToConnect.deduplicate()
 
+  # Connect to all nodes
   for i in countup(0, uniquePeers.len, MaxParallelDials):
     let stop = min(i + MaxParallelDials, uniquePeers.len)
     await pm.connectToNodes(uniquePeers[i..<stop])

waku/node/peer_manager/waku_peer_store.nim

@@ -102,6 +102,9 @@ proc connectedness*(peerStore: PeerStore, peerId: PeerID): Connectedness =
   # TODO: richer return than just bool, e.g. add enum "CanConnect", "CannotConnect", etc. based on recent connection attempts
   return peerStore[ConnectionBook].book.getOrDefault(peerId, NotConnected)
 
+proc hasShard*(peerStore: PeerStore, peerId: PeerID, cluster, shard: uint16): bool =
+  peerStore[ENRBook][peerId].containsShard(cluster, shard)
+
 proc isConnected*(peerStore: PeerStore, peerId: PeerID): bool =
   # Returns `true` if the peer is connected
   peerStore.connectedness(peerId) == Connected