peer_manager set shard info and extend rest test to validate it

Author: Ivansete-status

tests/wakunode_rest/test_rest_admin.nim

@@ -65,7 +65,7 @@ suite "Waku v2 Rest API - Admin":
     ): Future[void] {.async, gcsafe.} =
       await sleepAsync(0.milliseconds)
 
-    let shard = RelayShard(clusterId: clusterId, shardId: 0)
+    let shard = RelayShard(clusterId: clusterId, shardId: 5)
     node1.subscribe((kind: PubsubSub, topic: $shard), simpleHandler).isOkOr:
       assert false, "Failed to subscribe to topic: " & $error
     node2.subscribe((kind: PubsubSub, topic: $shard), simpleHandler).isOkOr:
@@ -212,6 +212,18 @@ suite "Waku v2 Rest API - Admin":
     let conn2 = await node1.peerManager.connectPeer(peerInfo2)
     let conn3 = await node1.peerManager.connectPeer(peerInfo3)
 
+    var count = 0
+    while count < 20:
+      ## Wait ~1s at most for the peer store to update shard info
+      let getRes = await client.getPeers()
+      if getRes.data.allIt(it.shards == @[5.uint16]):
+        break
+
+      count.inc()
+      await sleepAsync(50.milliseconds)
+
+    assert count < 20, "Timeout waiting for shards to be updated in peer store"
+
     # Check successful connections
     check:
       conn2 == true

waku/node/peer_manager/peer_manager.nim

@@ -660,11 +660,8 @@ proc onPeerMetadata(pm: PeerManager, peerId: PeerId) {.async.} =
 
     # Store the shard information from metadata in the peer store
     if pm.switch.peerStore.peerExists(peerId):
-      var peerInfo = pm.switch.peerStore.getPeer(peerId)
-      peerInfo.shards = metadata.shards.mapIt(it.uint16)
-      # Note: We don't need to call updatePeerInfo since we modified the reference directly
-      debug "Updated peer shards from metadata",
-        peerId = peerId, shards = peerInfo.shards
+      let shards = metadata.shards.mapIt(it.uint16)
+      pm.switch.peerStore.setShardInfo(peerId, shards)
 
     return
 

waku/node/peer_manager/waku_peer_store.nim

@@ -39,6 +39,9 @@ type
   # Keeps track of the ENR (Ethereum Node Record) of a peer
   ENRBook* = ref object of PeerBook[enr.Record]
 
+  # Keeps track of peer shards
+  ShardBook* = ref object of PeerBook[seq[uint16]]
+
 proc getPeer*(peerStore: PeerStore, peerId: PeerId): RemotePeerInfo =
   let addresses =
     if peerStore[LastSeenBook][peerId].isSome():
@@ -55,6 +58,7 @@ proc getPeer*(peerStore: PeerStore, peerId: PeerId): RemotePeerInfo =
       else:
         none(enr.Record),
     protocols: peerStore[ProtoBook][peerId],
+    shards: peerStore[ShardBook][peerId],
     agent: peerStore[AgentBook][peerId],
     protoVersion: peerStore[ProtoVersionBook][peerId],
     publicKey: peerStore[KeyBook][peerId],
@@ -76,6 +80,7 @@ proc peers*(peerStore: PeerStore): seq[RemotePeerInfo] =
       toSeq(peerStore[AddressBook].book.keys()),
       toSeq(peerStore[ProtoBook].book.keys()),
       toSeq(peerStore[KeyBook].book.keys()),
+      toSeq(peerStore[ShardBook].book.keys()),
     )
     .toHashSet()
 
@@ -127,6 +132,9 @@ proc addPeer*(peerStore: PeerStore, peer: RemotePeerInfo, origin = UnknownOrigin
   if peer.enr.isSome():
     peerStore[ENRBook][peer.peerId] = peer.enr.get()
 
+proc setShardInfo*(peerStore: PeerStore, peerId: PeerID, shards: seq[uint16]) =
+  peerStore[ShardBook][peerId] = shards
+
 proc peers*(peerStore: PeerStore, proto: string): seq[RemotePeerInfo] =
   peerStore.peers().filterIt(it.protocols.contains(proto))