diff --git a/logos_delivery/waku/api/api_conf.nim b/logos_delivery/waku/api/api_conf.nim index 5810b1eb2..846563064 100644 --- a/logos_delivery/waku/api/api_conf.nim +++ b/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 ---------- diff --git a/logos_delivery/waku/factory/conf_builder/waku_conf_builder.nim b/logos_delivery/waku/factory/conf_builder/waku_conf_builder.nim index dc17fb4ff..c4b6d3739 100644 --- a/logos_delivery/waku/factory/conf_builder/waku_conf_builder.nim +++ b/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(@[]) diff --git a/logos_delivery/waku/factory/networks_config.nim b/logos_delivery/waku/factory/networks_config.nim index 3d1e296fe..f0cef9d8d 100644 --- a/logos_delivery/waku/factory/networks_config.nim +++ b/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) diff --git a/logos_delivery/waku/factory/node_factory.nim b/logos_delivery/waku/factory/node_factory.nim index 30e37850a..db040c322 100644 --- a/logos_delivery/waku/factory/node_factory.nim +++ b/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") diff --git a/logos_delivery/waku/node/waku_node.nim b/logos_delivery/waku/node/waku_node.nim index 2ad7dc601..0cc7d6a62 100644 --- a/logos_delivery/waku/node/waku_node.nim +++ b/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() diff --git a/logos_delivery/waku/waku_core/topics/sharding.nim b/logos_delivery/waku/waku_core/topics/sharding.nim index 704180554..9f7c1b769 100644 --- a/logos_delivery/waku/waku_core/topics/sharding.nim +++ b/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. diff --git a/tests/waku_core/topics/test_sharding.nim b/tests/waku_core/topics/test_sharding.nim index 5979269e8..e73302626 100644 --- a/tests/waku_core/topics/test_sharding.nim +++ b/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 = diff --git a/tools/confutils/cli_args.nim b/tools/confutils/cli_args.nim index 5228dfff7..757d7f9a8 100644 --- a/tools/confutils/cli_args.nim +++ b/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: