Fluffy: Update the state bridge to send each content offer to the closest connected portal client (#3278)

* Refactor state bridge to support sending each content to any of the connected portal clients sorted by distance from the content key.

Author: bhartnett

fluffy/tools/portal_bridge/portal_bridge_state.nim

@@ -37,26 +37,25 @@ type
     stateRoot: Hash32
     stateDiffs: seq[TransactionDiff]
 
-  BlockOffersRef = ref object
+  BlockOffers = ref object
     blockNumber: uint64
     accountTrieOffers: seq[AccountTrieOfferWithKey]
     contractTrieOffers: seq[ContractTrieOfferWithKey]
     contractCodeOffers: seq[ContractCodeOfferWithKey]
 
   PortalStateGossipWorker = ref object
     id: int
-    portalClient: RpcClient
-    portalUrl: JsonRpcUrl
-    nodeId: NodeId
-    blockOffersQueue: AsyncQueue[BlockOffersRef]
+    portalClients: OrderedTable[NodeId, RpcClient]
+    portalEndpoints: seq[(JsonRpcUrl, NodeId)]
+    blockOffersQueue: AsyncQueue[BlockOffers]
     gossipBlockOffersLoop: Future[void]
 
   PortalStateBridge* = ref object
     web3Client: RpcClient
     web3Url: JsonRpcUrl
     db: DatabaseRef
     blockDataQueue: AsyncQueue[BlockData]
-    blockOffersQueue: AsyncQueue[BlockOffersRef]
+    blockOffersQueue: AsyncQueue[BlockOffers]
     gossipWorkers: seq[PortalStateGossipWorker]
     collectBlockDataLoop: Future[void]
     buildBlockOffersLoop: Future[void]
@@ -92,27 +91,6 @@ proc putLastPersistedBlockNumber(db: DatabaseRef, blockNumber: uint64) {.inline.
   if blockNumber > db.getLastPersistedBlockNumber().valueOr(0):
     db.put(rlp.encode("lastPersistedBlockNumber"), rlp.encode(blockNumber))
 
-proc collectOffer(
-    offersMap: OrderedTableRef[seq[byte], seq[byte]],
-    offerWithKey:
-      AccountTrieOfferWithKey | ContractTrieOfferWithKey | ContractCodeOfferWithKey,
-) {.inline.} =
-  let keyBytes = offerWithKey.key.toContentKey().encode().asSeq()
-  offersMap[keyBytes] = offerWithKey.offer.encode()
-
-proc recursiveCollectOffer(
-    offersMap: OrderedTableRef[seq[byte], seq[byte]],
-    offerWithKey: AccountTrieOfferWithKey | ContractTrieOfferWithKey,
-) =
-  offersMap.collectOffer(offerWithKey)
-
-  # root node, recursive collect is finished
-  if offerWithKey.key.path.unpackNibbles().len() == 0:
-    return
-
-  # continue the recursive collect
-  offersMap.recursiveCollectOffer(offerWithKey.getParent())
-
 proc runCollectBlockDataLoop(
     bridge: PortalStateBridge, startBlockNumber: uint64
 ) {.async: (raises: []).} =
@@ -237,7 +215,7 @@ proc runBuildBlockOffersLoop(
           builder.buildBlockOffers()
 
           await bridge.blockOffersQueue.addLast(
-            BlockOffersRef(
+            BlockOffers(
               blockNumber: 0.uint64,
               accountTrieOffers: builder.getAccountTrieOffers(),
               contractTrieOffers: builder.getContractTrieOffers(),
@@ -284,7 +262,7 @@ proc runBuildBlockOffersLoop(
           builder.buildBlockOffers()
 
           await bridge.blockOffersQueue.addLast(
-            BlockOffersRef(
+            BlockOffers(
               blockNumber: blockData.blockNumber,
               accountTrieOffers: builder.getAccountTrieOffers(),
               contractTrieOffers: builder.getContractTrieOffers(),
@@ -299,7 +277,107 @@ proc runBuildBlockOffersLoop(
   except CancelledError:
     trace "buildBlockOffersLoop canceled"
 
-proc runGossipBlockOffersLoop(
+proc collectOffer(
+    offersMap: OrderedTableRef[seq[byte], seq[byte]],
+    offerWithKey:
+      AccountTrieOfferWithKey | ContractTrieOfferWithKey | ContractCodeOfferWithKey,
+) {.inline.} =
+  let keyBytes = offerWithKey.key.toContentKey().encode().asSeq()
+  offersMap[keyBytes] = offerWithKey.offer.encode()
+
+proc recursiveCollectOffer(
+    offersMap: OrderedTableRef[seq[byte], seq[byte]],
+    offerWithKey: AccountTrieOfferWithKey | ContractTrieOfferWithKey,
+) =
+  offersMap.collectOffer(offerWithKey)
+
+  # root node, recursive collect is finished
+  if offerWithKey.key.path.unpackNibbles().len() == 0:
+    return
+
+  # continue the recursive collect
+  offersMap.recursiveCollectOffer(offerWithKey.getParent())
+
+func buildOffersMap(blockOffers: BlockOffers): auto =
+  let offersMap = newOrderedTable[seq[byte], seq[byte]]()
+
+  for offerWithKey in blockOffers.accountTrieOffers:
+    offersMap.recursiveCollectOffer(offerWithKey)
+  for offerWithKey in blockOffers.contractTrieOffers:
+    offersMap.recursiveCollectOffer(offerWithKey)
+  for offerWithKey in blockOffers.contractCodeOffers:
+    offersMap.collectOffer(offerWithKey)
+
+  offersMap
+
+proc orderPortalClientsByDistanceFromContent(
+    worker: PortalStateGossipWorker, contentKey: seq[byte]
+) =
+  let contentId = ContentKeyByteList.init(contentKey).toContentId()
+
+  # Closure to sort the portal clients using their nodeIds
+  # and comparing them to the contentId to be gossipped
+  proc portalClientsCmp(x, y: (NodeId, RpcClient)): int =
+    let
+      xDistance = contentId xor x[0]
+      yDistance = contentId xor y[0]
+
+    if xDistance == yDistance:
+      0
+    elif xDistance > yDistance:
+      1
+    else:
+      -1
+
+  # Sort the portalClients based on distance from the content so that
+  # we gossip each piece of content to the closest node first
+  worker.portalClients.sort(portalClientsCmp)
+
+proc contentFoundInNetwork(
+    worker: PortalStateGossipWorker, contentKey: seq[byte]
+): Future[bool] {.async: (raises: [CancelledError]).} =
+  for nodeId, client in worker.portalClients:
+    try:
+      let contentInfo = await client.portal_stateGetContent(contentKey.to0xHex())
+      if contentInfo.content.len() > 0:
+        trace "Found existing content in network",
+          contentKey = contentKey.to0xHex(), nodeId, workerId = worker.id
+        return true
+    except CancelledError as e:
+      raise e
+    except CatchableError as e:
+      debug "Unable to find existing content in network",
+        contentKey = contentKey.to0xHex(), nodeId, error = e.msg, workerId = worker.id
+  return false
+
+proc gossipContentIntoNetwork(
+    worker: PortalStateGossipWorker,
+    minGossipPeers: int,
+    contentKey: seq[byte],
+    contentOffer: seq[byte],
+): Future[bool] {.async: (raises: [CancelledError]).} =
+  for nodeId, client in worker.portalClients:
+    try:
+      let
+        putContentResult = await client.portal_statePutContent(
+          contentKey.to0xHex(), contentOffer.to0xHex()
+        )
+        numPeers = putContentResult.peerCount
+      if numPeers >= minGossipPeers:
+        trace "Offer successfully gossipped to peers",
+          contentKey = contentKey.to0xHex(), nodeId, numPeers, workerId = worker.id
+        return true
+      else:
+        warn "Offer not gossiped to enough peers",
+          contentKey = contentKey.to0xHex(), nodeId, numPeers, workerId = worker.id
+    except CancelledError as e:
+      raise e
+    except CatchableError as e:
+      error "Failed to gossip offer to peers",
+        contentKey = contentKey.to0xHex(), nodeId, error = e.msg, workerId = worker.id
+  return false
+
+proc runGossipLoop(
     worker: PortalStateGossipWorker,
     verifyGossip: bool,
     skipGossipForExisting: bool,
@@ -308,106 +386,52 @@ proc runGossipBlockOffersLoop(
   debug "Starting gossip block offers loop", workerId = worker.id
 
   try:
-    # Create one client per worker in order to improve performance.
+    # Create separate clients in each worker in order to improve performance.
     # WebSocket connections don't perform well when shared by many
     # concurrent workers.
-    worker.portalClient = newRpcClientConnect(worker.portalUrl)
+    for (rpcUrl, nodeId) in worker.portalEndpoints:
+      worker.portalClients[nodeId] = newRpcClientConnect(rpcUrl)
 
-    var blockOffers = await worker.blockOffersQueue.popFirst()
-
-    while true:
+    var
+      blockOffers = await worker.blockOffersQueue.popFirst()
       # A table of offer key, value pairs is used to filter out duplicates so
       # that we don't gossip the same offer multiple times.
-      let offersMap = newOrderedTable[seq[byte], seq[byte]]()
-
-      for offerWithKey in blockOffers.accountTrieOffers:
-        offersMap.recursiveCollectOffer(offerWithKey)
-      for offerWithKey in blockOffers.contractTrieOffers:
-        offersMap.recursiveCollectOffer(offerWithKey)
-      for offerWithKey in blockOffers.contractCodeOffers:
-        offersMap.collectOffer(offerWithKey)
-
-      # We need to use a closure here because nodeId is required to calculate the
-      # distance of each content id from the node
-      proc offersMapCmp(x, y: (seq[byte], seq[byte])): int =
-        let
-          xId = ContentKeyByteList.init(x[0]).toContentId()
-          yId = ContentKeyByteList.init(y[0]).toContentId()
-          xDistance = worker.nodeId xor xId
-          yDistance = worker.nodeId xor yId
-
-        if xDistance == yDistance:
-          0
-        elif xDistance > yDistance:
-          1
-        else:
-          -1
-
-      # Sort the offers based on the distance from the node so that we will gossip
-      # content that is closest to the node first
-      offersMap.sort(offersMapCmp)
+      offersMap = buildOffersMap(blockOffers)
 
+    while true:
       var retryGossip = false
-      for k, v in offersMap:
+
+      for contentKey, contentOffer in offersMap:
+        worker.orderPortalClientsByDistanceFromContent(contentKey)
+
         # Check if we need to gossip the content
-        var gossipContent = true
-
-        if skipGossipForExisting:
-          try:
-            let contentInfo =
-              await worker.portalClient.portal_stateGetContent(k.to0xHex())
-            if contentInfo.content.len() > 0:
-              gossipContent = false
-          except CancelledError as e:
-            raise e
-          except CatchableError as e:
-            debug "Unable to find existing content. Will attempt to gossip content: ",
-              contentKey = k.to0xHex(), error = e.msg, workerId = worker.id
+        if skipGossipForExisting and (await worker.contentFoundInNetwork(contentKey)):
+          continue # move on to the next content key
 
         # Gossip the content into the network
-        if gossipContent:
-          try:
-            let
-              putContentResult = await worker.portalClient.portal_statePutContent(
-                k.to0xHex(), v.to0xHex()
-              )
-              numPeers = putContentResult.peerCount
-            if numPeers >= minGossipPeers:
-              debug "Offer successfully gossipped to peers",
-                numPeers, workerId = worker.id
-            else:
-              warn "Offer not gossiped to enough peers", numPeers, workerId = worker.id
-              retryGossip = true
-              break
-          except CancelledError as e:
-            raise e
-          except CatchableError as e:
-            error "Failed to gossip offer to peers", error = e.msg, workerId = worker.id
-            retryGossip = true
-            break
+        let gossipCompleted = await worker.gossipContentIntoNetwork(
+          minGossipPeers, contentKey, contentOffer
+        )
+        if not gossipCompleted:
+          # Retry gossip of this block
+          retryGossip = true
+          break
 
       # Check if the content can be found in the network
       var foundContentKeys = newSeq[seq[byte]]()
       if verifyGossip and not retryGossip:
-        # wait for the peers to be updated
+        # Wait for the peers to be updated.
+        # Wait time is proportional to the number of offers
         let waitTimeMs = 200 + (offersMap.len() * 20)
         await sleepAsync(waitTimeMs.milliseconds)
-          # wait time is proportional to the number of offers
-
-        for k, _ in offersMap:
-          try:
-            let contentInfo =
-              await worker.portalClient.portal_stateGetContent(k.to0xHex())
-            if contentInfo.content.len() == 0:
-              error "Found empty contentValue", workerId = worker.id
-              retryGossip = true
-              break
-            foundContentKeys.add(k)
-          except CancelledError as e:
-            raise e
-          except CatchableError as e:
-            warn "Unable to find content with key. Will retry gossipping content:",
-              contentKey = k.to0xHex(), error = e.msg, workerId = worker.id
+
+        for contentKey, _ in offersMap:
+          worker.orderPortalClientsByDistanceFromContent(contentKey)
+
+          if await worker.contentFoundInNetwork(contentKey):
+            foundContentKeys.add(contentKey)
+          else:
+            # Retry gossip of this block
             retryGossip = true
             break
 
@@ -418,13 +442,17 @@ proc runGossipBlockOffersLoop(
         # Don't retry gossip for content that was found in the network
         for key in foundContentKeys:
           offersMap.del(key)
+
         warn "Retrying state gossip for block: ",
           blockNumber = blockOffers.blockNumber,
           remainingOffers = offersMap.len(),
           workerId = worker.id
 
-        # We might need to reconnect if using a WebSocket client
-        await worker.portalClient.tryReconnect(worker.portalUrl)
+        # We might need to reconnect if using WebSocket clients
+        for (rpcUrl, nodeId) in worker.portalEndpoints:
+          await worker.portalClients.getOrDefault(nodeId).tryReconnect(rpcUrl)
+
+        # Jump back to the top of while loop to retry processing the current block
         continue
 
       if blockOffers.blockNumber mod 1000 == 0:
@@ -439,6 +467,7 @@ proc runGossipBlockOffersLoop(
           workerId = worker.id
 
       blockOffers = await worker.blockOffersQueue.popFirst()
+      offersMap = buildOffersMap(blockOffers)
   except CancelledError:
     trace "gossipBlockOffersLoop canceled"
 
@@ -528,7 +557,7 @@ proc start*(bridge: PortalStateBridge, config: PortalBridgeConf) =
   info "Starting concurrent gossip workers", workerCount = bridge.gossipWorkers.len()
 
   for worker in bridge.gossipWorkers:
-    worker.gossipBlockOffersLoop = worker.runGossipBlockOffersLoop(
+    worker.gossipBlockOffersLoop = worker.runGossipLoop(
       config.verifyGossip, config.skipGossipForExisting, config.minGossipPeers.int
     )
 
@@ -580,19 +609,16 @@ proc runState*(
     web3Url: config.web3RpcUrl,
     db: db,
     blockDataQueue: newAsyncQueue[BlockData](queueSize),
-    blockOffersQueue: newAsyncQueue[BlockOffersRef](queueSize),
+    blockOffersQueue: newAsyncQueue[BlockOffers](queueSize),
     gossipWorkers: newSeq[PortalStateGossipWorker](),
   )
 
   for i in 0 ..< config.gossipWorkers.int:
-    let
-      (rpcUrl, nodeId) = portalEndpoints[i mod config.portalRpcEndpoints.int]
-      worker = PortalStateGossipWorker(
-        id: i + 1,
-        portalUrl: rpcUrl,
-        nodeId: nodeId,
-        blockOffersQueue: bridge.blockOffersQueue,
-      )
+    let worker = PortalStateGossipWorker(
+      id: i + 1,
+      portalEndpoints: portalEndpoints,
+      blockOffersQueue: bridge.blockOffersQueue,
+    )
     bridge.gossipWorkers.add(worker)
 
   bridge.start(config)