feat: allow overriding actual shard values in auto-sharding

Auto-sharding currently produces shard ids in the range
[0..NUM_SHARDS_IN_NETWORK-1], derived from the content topic hash. To
interoperate with Status, which uses a specific set of shard values, we
need to be able to force the actual shard values while keeping the
auto-sharding distribution.

This adds an optional shard override array (length NUM_SHARDS_IN_NETWORK):
auto-sharding still hashes the content topic to an index, but the shard id
placed in the pubsub topic becomes shardOverride[index] instead of the
index itself.

Wired through:
- Sharding.getGenZeroShard remaps the computed index via shardOverride
- ShardingConf gains a shardOverride field; default subscribe shards and
  shard validation use the actual (override) shard values
- WakuConfBuilder.withShardOverride + buildShardingConf threading
- CLI arg --shard-override (repeatable)
- mountAutoSharding accepts the override (validates length)
- api_conf AutoShardingConfig gains shardOverride (+ JSON)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

logos_delivery/waku/api/api_conf.nim

@@ -16,6 +16,10 @@ export json_serialization, json_options
 
 type AutoShardingConfig* = object
   numShardsInCluster*: uint16
+  # Optional list of forced shard values. When set it must hold exactly
+  # `numShardsInCluster` values; auto-sharding selects an index and uses the
+  # value at that index as the shard. Empty means shards are `[0..n-1]`.
+  shardOverride*: seq[uint16]
 
 type RlnConfig* = object
   contractAddress*: string
@@ -213,6 +217,8 @@ proc toWakuConf*(
   b.withShardingConf(ShardingConfKind.AutoSharding)
   let autoShardingConfig = protocolsConfig.autoShardingConfig
   b.withNumShardsInCluster(autoShardingConfig.numShardsInCluster)
+  if autoShardingConfig.shardOverride.len > 0:
+    b.withShardOverride(autoShardingConfig.shardOverride)
 
   # Process entry nodes - supports enrtree:, enr:, and multiaddress formats
   if protocolsConfig.entryNodes.len > 0:
@@ -279,24 +285,32 @@ proc toWakuConf*(
 proc writeValue*(w: var JsonWriter, val: AutoShardingConfig) {.raises: [IOError].} =
   w.beginRecord()
   w.writeField("numShardsInCluster", val.numShardsInCluster)
+  w.writeField("shardOverride", val.shardOverride)
   w.endRecord()
 
 proc readValue*(
     r: var JsonReader, val: var AutoShardingConfig
 ) {.raises: [SerializationError, IOError].} =
-  var numShardsInCluster: Option[uint16]
+  var
+    numShardsInCluster: Option[uint16]
+    shardOverride: Option[seq[uint16]]
 
   for fieldName in readObjectFields(r):
     case fieldName
     of "numShardsInCluster":
       numShardsInCluster = some(r.readValue(uint16))
+    of "shardOverride":
+      shardOverride = some(r.readValue(seq[uint16]))
     else:
       r.raiseUnexpectedField(fieldName, "AutoShardingConfig")
 
   if numShardsInCluster.isNone():
     r.raiseUnexpectedValue("Missing required field 'numShardsInCluster'")
 
-  val = AutoShardingConfig(numShardsInCluster: numShardsInCluster.get())
+  val = AutoShardingConfig(
+    numShardsInCluster: numShardsInCluster.get(),
+    shardOverride: shardOverride.get(@[]),
+  )
 
 # ---------- RlnConfig ----------
 

logos_delivery/waku/factory/conf_builder/waku_conf_builder.nim

@@ -106,6 +106,7 @@ type WakuConfBuilder* = object
   clusterId: Option[uint16]
   shardingConf: Option[ShardingConfKind]
   numShardsInCluster: Option[uint16]
+  shardOverride: Option[seq[uint16]]
   subscribeShards: Option[seq[uint16]]
   protectedShards: Option[seq[ProtectedShard]]
   contentTopics: Option[seq[string]]
@@ -206,6 +207,9 @@ proc withShardingConf*(b: var WakuConfBuilder, shardingConf: ShardingConfKind) =
 proc withNumShardsInCluster*(b: var WakuConfBuilder, numShardsInCluster: uint16) =
   b.numShardsInCluster = some(numShardsInCluster)
 
+proc withShardOverride*(b: var WakuConfBuilder, shardOverride: seq[uint16]) =
+  b.shardOverride = some(shardOverride)
+
 proc withSubscribeShards*(b: var WakuConfBuilder, shards: seq[uint16]) =
   b.subscribeShards = some(shards)
 
@@ -345,6 +349,7 @@ proc nodeKey(
 proc buildShardingConf(
     bShardingConfKind: Option[ShardingConfKind],
     bNumShardsInCluster: Option[uint16],
+    bShardOverride: Option[seq[uint16]],
     bSubscribeShards: Option[seq[uint16]],
 ): (ShardingConf, seq[uint16]) =
   case bShardingConfKind.get(DefaultShardingConfKind)
@@ -352,10 +357,15 @@ proc buildShardingConf(
     (ShardingConf(kind: StaticSharding), bSubscribeShards.get(@[]))
   of AutoSharding:
     let numShardsInCluster = bNumShardsInCluster.get(DefaultNumShardsInCluster)
-    let shardingConf =
-      ShardingConf(kind: AutoSharding, numShardsInCluster: numShardsInCluster)
-    let upperShard = uint16(numShardsInCluster - 1)
-    (shardingConf, bSubscribeShards.get(toSeq(0.uint16 .. upperShard)))
+    let shardingConf = ShardingConf(
+      kind: AutoSharding,
+      numShardsInCluster: numShardsInCluster,
+      shardOverride: bShardOverride.get(@[]),
+    )
+    # Subscribe to every actual shard by default. With a shard override the
+    # actual shards are the override values, otherwise the indices
+    # `[0..numShardsInCluster-1]`.
+    (shardingConf, bSubscribeShards.get(shardingConf.shards()))
 
 template checkSetPresetValueToField[T](
     field: var Option[T], presetVal: T, msg: static string
@@ -603,7 +613,8 @@ proc build*(
       builder.clusterId.get().uint16
 
   let (shardingConf, subscribeShards) = buildShardingConf(
-    builder.shardingConf, builder.numShardsInCluster, builder.subscribeShards
+    builder.shardingConf, builder.numShardsInCluster, builder.shardOverride,
+    builder.subscribeShards,
   )
   let protectedShards = builder.protectedShards.get(@[])
 

logos_delivery/waku/factory/networks_config.nim

@@ -15,6 +15,12 @@ type
     case kind*: ShardingConfKind
     of AutoSharding:
       numShardsInCluster*: uint16
+      # Optional list of forced shard values used by autosharding. When set it
+      # must hold exactly `numShardsInCluster` values: autosharding selects an
+      # index in `[0..numShardsInCluster-1]` and the actual shard used is the
+      # value at that index. Empty means shards are the indices themselves
+      # (`[0..numShardsInCluster-1]`).
+      shardOverride*: seq[uint16]
     of StaticSharding:
       discard
 
@@ -123,6 +129,21 @@ proc LogosTestConf*(T: type NetworkPresetConf): NetworkPresetConf =
     ],
   )
 
+proc shards*(shardingConf: ShardingConf): seq[uint16] =
+  ## The actual set of shard ids autosharding can produce, and which a node
+  ## subscribes to by default. When a shard override is configured these are the
+  ## override values, otherwise the indices `[0..numShardsInCluster-1]`.
+  case shardingConf.kind
+  of StaticSharding:
+    return @[]
+  of AutoSharding:
+    if shardingConf.shardOverride.len > 0:
+      return shardingConf.shardOverride
+    var allShards = newSeq[uint16](shardingConf.numShardsInCluster.int)
+    for i in 0 ..< shardingConf.numShardsInCluster.int:
+      allShards[i] = uint16(i)
+    return allShards
+
 proc validateShards*(
     shardingConf: ShardingConf, shards: seq[uint16]
 ): Result[void, string] =
@@ -131,11 +152,23 @@ proc validateShards*(
     return ok()
   of AutoSharding:
     let numShardsInCluster = shardingConf.numShardsInCluster
+    # A shard override, when provided, must hold exactly one value per shard.
+    if shardingConf.shardOverride.len > 0 and
+        shardingConf.shardOverride.len != numShardsInCluster.int:
+      let msg =
+        "shardOverride must hold exactly numShardsInCluster (" & $numShardsInCluster &
+        ") values, got: " & $shardingConf.shardOverride.len
+      error "validateShards failed", error = msg
+      return err(msg)
+
+    # Valid shards are the actual shard values (override values when configured,
+    # otherwise the indices `[0..numShardsInCluster-1]`).
+    let validShards = shardingConf.shards()
     for shard in shards:
-      if shard >= numShardsInCluster:
+      if shard notin validShards:
         let msg =
-          "validateShards invalid shard: " & $shard & " when numShardsInCluster: " &
-          $numShardsInCluster
+          "validateShards invalid shard: " & $shard & " when valid shards are: " &
+          $validShards
         error "validateShards failed", error = msg
         return err(msg)
 

logos_delivery/waku/factory/node_factory.nim

@@ -246,7 +246,10 @@ proc setupProtocols(
     node.setupStoreResume()
 
   if conf.shardingConf.kind == AutoSharding:
-    node.mountAutoSharding(conf.clusterId, conf.shardingConf.numShardsInCluster).isOkOr:
+    node.mountAutoSharding(
+      conf.clusterId, conf.shardingConf.numShardsInCluster,
+      conf.shardingConf.shardOverride,
+    ).isOkOr:
       return err("failed to mount waku auto sharding: " & error)
   else:
     warn("Auto sharding is disabled")

logos_delivery/waku/node/waku_node.nim

@@ -297,11 +297,24 @@ proc mountMetadata*(
 
 ## Waku AutoSharding
 proc mountAutoSharding*(
-    node: WakuNode, clusterId: uint16, shardCount: uint32
+    node: WakuNode,
+    clusterId: uint16,
+    shardCount: uint32,
+    shardOverride: seq[uint16] = @[],
 ): Result[void, string] =
-  info "mounting auto sharding", clusterId = clusterId, shardCount = shardCount
-  node.wakuAutoSharding =
-    some(Sharding(clusterId: clusterId, shardCountGenZero: shardCount))
+  if shardOverride.len > 0 and shardOverride.len != int(shardCount):
+    return err(
+      "shard override must hold exactly shardCount (" & $shardCount & ") values, got: " &
+        $shardOverride.len
+    )
+
+  info "mounting auto sharding",
+    clusterId = clusterId, shardCount = shardCount, shardOverride = shardOverride
+  node.wakuAutoSharding = some(
+    Sharding(
+      clusterId: clusterId, shardCountGenZero: shardCount, shardOverride: shardOverride
+    )
+  )
 
   return ok()
 

logos_delivery/waku/waku_core/topics/sharding.nim

@@ -14,9 +14,23 @@ type Sharding* = object
   clusterId*: uint16
   # TODO: generations could be stored in a table here
   shardCountGenZero*: uint32
+  # Optional override for the actual shard ids used by autosharding. When non
+  # empty, it must hold exactly `shardCountGenZero` values: autosharding still
+  # hashes the content topic to an index in `[0..shardCountGenZero-1]`, but the
+  # shard id placed in the pubsub topic is `shardOverride[index]` instead of the
+  # index itself. This allows forcing a specific set of shard values (e.g. to
+  # interoperate with Status) while keeping the autosharding distribution.
+  shardOverride*: seq[uint16]
 
-proc new*(T: type Sharding, clusterId: uint16, shardCount: uint32): T =
-  return Sharding(clusterId: clusterId, shardCountGenZero: shardCount)
+proc new*(
+    T: type Sharding,
+    clusterId: uint16,
+    shardCount: uint32,
+    shardOverride: seq[uint16] = @[],
+): T =
+  return Sharding(
+    clusterId: clusterId, shardCountGenZero: shardCount, shardOverride: shardOverride
+  )
 
 proc getGenZeroShard*(s: Sharding, topic: NsContentTopic, count: int): RelayShard =
   let bytes = toBytes(topic.application) & toBytes(topic.version)
@@ -26,9 +40,17 @@ proc getGenZeroShard*(s: Sharding, topic: NsContentTopic, count: int): RelayShar
   # We only use the last 64 bits of the hash as having more shards is unlikely.
   let hashValue = uint64.fromBytesBE(hash.data[24 .. 31])
 
-  let shard = hashValue mod uint64(count)
+  let index = int(hashValue mod uint64(count))
 
-  RelayShard(clusterId: s.clusterId, shardId: uint16(shard))
+  # The hashed value is an index into the shard space. By default the shard id
+  # equals that index, but a configured override remaps it to a forced value.
+  let shardId =
+    if s.shardOverride.len == count:
+      s.shardOverride[index]
+    else:
+      uint16(index)
+
+  RelayShard(clusterId: s.clusterId, shardId: shardId)
 
 proc getShard*(s: Sharding, topic: NsContentTopic): Result[RelayShard, string] =
   ## Compute the (pubsub topic) shard to use for this content topic.

tests/waku_core/topics/test_sharding.nim

@@ -65,6 +65,50 @@ suite "Autosharding":
         shard9 == RelayShard(clusterId: ClusterId, shardId: 7)
         shard10 == RelayShard(clusterId: ClusterId, shardId: 3)
 
+    test "Generate Gen0 Shard with shardOverride":
+      # Given a sharding with a forced shard mapping (index -> value)
+      let shardOverride = @[10'u16, 11, 12, 13, 14, 15, 16, 17]
+      let sharding = Sharding(
+        clusterId: ClusterId,
+        shardCountGenZero: GenerationZeroShardsCount,
+        shardOverride: shardOverride,
+      )
+
+      let
+        nsContentTopic1 = NsContentTopic.parse(contentTopicShort).value()
+          # hashes to index 3
+        nsContentTopic3 = NsContentTopic.parse(contentTopicShort2).value()
+          # hashes to index 6
+        nsContentTopic9 = NsContentTopic.parse(contentTopicFull4).value()
+          # hashes to index 7
+
+      # When we generate gen0 shards from them
+      let
+        shard1 = sharding.getGenZeroShard(nsContentTopic1, GenerationZeroShardsCount)
+        shard3 = sharding.getGenZeroShard(nsContentTopic3, GenerationZeroShardsCount)
+        shard9 = sharding.getGenZeroShard(nsContentTopic9, GenerationZeroShardsCount)
+
+      # Then the computed index is remapped through the override array
+      check:
+        shard1 == RelayShard(clusterId: ClusterId, shardId: 13) # index 3 -> 13
+        shard3 == RelayShard(clusterId: ClusterId, shardId: 16) # index 6 -> 16
+        shard9 == RelayShard(clusterId: ClusterId, shardId: 17) # index 7 -> 17
+
+    test "shardOverride is ignored when length differs from shard count":
+      # Given a shardOverride whose length does not match the shard count
+      let sharding = Sharding(
+        clusterId: ClusterId,
+        shardCountGenZero: GenerationZeroShardsCount,
+        shardOverride: @[10'u16, 11, 12], # wrong length, must be ignored
+      )
+
+      let nsContentTopic1 = NsContentTopic.parse(contentTopicShort).value()
+      let shard1 = sharding.getGenZeroShard(nsContentTopic1, GenerationZeroShardsCount)
+
+      # Then it falls back to using the index as the shard id
+      check:
+        shard1 == RelayShard(clusterId: ClusterId, shardId: 3)
+
   suite "getShard from NsContentTopic":
     test "Generate Gen0 Shard with topic.generation==none":
       let sharding =

tools/confutils/cli_args.nim

@@ -360,6 +360,12 @@ hence would have reachability issues.""",
       name: "shard"
     .}: seq[uint16]
 
+    shardOverride* {.
+      desc:
+        "Override the actual shard values used by auto-sharding. Provide exactly NUM_SHARDS_IN_NETWORK values; auto-sharding selects an index into this array and uses the value at that index as the shard. Argument may be repeated. Defaults to [0..NUM_SHARDS_IN_NETWORK-1] when unset",
+      name: "shard-override"
+    .}: seq[uint16]
+
     contentTopics* {.
       desc: "Default content topic to subscribe to. Argument may be repeated.",
       name: "content-topic"
@@ -1101,6 +1107,11 @@ proc toWakuConf*(n: WakuNodeConf): ConfResult[WakuConf] =
   elif networkPresetConf.isNone():
     b.withShardingConf(StaticSharding)
 
+  # It is not possible to pass an empty sequence on the CLI
+  # If this is empty, it means the user did not specify a shard override
+  if n.shardOverride.len != 0:
+    b.withShardOverride(n.shardOverride)
+
   # It is not possible to pass an empty sequence on the CLI
   # If this is empty, it means the user did not specify any shards
   if n.shards.len != 0: