fix(service-disco): maxBuckets scaling

Author: gmelodie

libp2p/protocols/kademlia/routing_table.nim

@@ -40,7 +40,15 @@ proc new*(
     selfId: selfId, localNodeId: localNodeId.get(selfId), buckets: @[], config: config
   )
 
+func bucketCount(maxBuckets: int): int =
+  ## Buckets past the key length can never be reached, and at least one bucket
+  ## is always needed.
+  clamp(maxBuckets, 1, IdLength * 8)
+
 proc bucketIndex*(rtable: RoutingTable, key: Key): int =
+  ## Bucket index is the shared prefix length with `selfId`. Tables with fewer
+  ## buckets than the key length don't rescale it: the last bucket absorbs every
+  ## peer sharing at least that long a prefix.
   let
     selfHash =
       if rtable.config.selfIdPreHashed:
@@ -50,10 +58,7 @@ proc bucketIndex*(rtable: RoutingTable, key: Key): int =
     keyHash = key.hashFor(rtable.config.hasher)
     lz = xorDistance(selfHash, keyHash).leadingZeros
 
-  if rtable.config.maxBuckets <= 1:
-    return 0
-
-  return min(((lz * rtable.config.maxBuckets) div 256), rtable.config.maxBuckets - 1)
+  min(lz, bucketCount(rtable.config.maxBuckets) - 1)
 
 proc peerIndexInBucket(bucket: Bucket, nodeId: Key): Opt[int] =
   for i, p in bucket.peers:
@@ -209,9 +214,7 @@ proc randomKeyInBucket*(
     selfId: Key, bucketIndex: int, rng: Rng, maxBuckets: int = DefaultMaxBuckets
 ): Key =
   let
-    index = clamp(bucketIndex, 0, (maxBuckets - 1))
-    buckets = max(1, maxBuckets)
-    leadingZeros = index * 256 div buckets
+    leadingZeros = clamp(bucketIndex, 0, bucketCount(maxBuckets) - 1)
     (byteIdx, rem) = divmod(leadingZeros, 8)
     boundBitIdx = 7 - rem
 

tests/libp2p/kademlia/test_routing_table.nim

@@ -169,6 +169,33 @@ suite "KadDHT Routing Table":
       idx == TargetBucket
       rid != selfId
 
+  test "bucketIndex is the prefix length, not a rescale of it":
+    let selfId = testKey(0)
+    var rt =
+      RoutingTable.new(selfId, RoutingTableConfig.new(hasher = Opt.some(noOpHasher)))
+
+    for lz in 0 .. 8:
+      let key = randomKeyInBucket(selfId, lz, rng())
+      check rt.bucketIndex(key) == lz
+
+  test "small maxBuckets keeps prefix resolution and saturates the last bucket":
+    let selfId = testKey(0)
+    const MaxBuckets = 16
+    var rt = RoutingTable.new(
+      selfId,
+      RoutingTableConfig.new(hasher = Opt.some(noOpHasher), maxBuckets = MaxBuckets),
+    )
+
+    # Prefix lengths below the cap each get their own bucket.
+    for lz in 0 ..< MaxBuckets - 1:
+      let key = randomKeyInBucket(selfId, lz, rng(), MaxBuckets)
+      check rt.bucketIndex(key) == lz
+
+    # Deeper prefixes all fall into the last bucket.
+    for lz in MaxBuckets - 1 .. MaxBuckets + 8:
+      let key = randomKeyInBucket(selfId, lz, rng(), DefaultMaxBuckets)
+      check rt.bucketIndex(key) == MaxBuckets - 1
+
   test "randomKey returns none for empty bucket":
     var bucket: Bucket
     check randomKey(bucket, rng()).isNone()