diff --git a/logos_delivery/api/conf/messaging_conf.nim b/logos_delivery/api/conf/messaging_conf.nim index 62abea75e..a78305c49 100644 --- a/logos_delivery/api/conf/messaging_conf.nim +++ b/logos_delivery/api/conf/messaging_conf.nim @@ -4,9 +4,9 @@ import results, libp2p/crypto/crypto import logos_delivery/api/conf/kernel_conf import logos_delivery/waku/common/logging import logos_delivery/waku/factory/networks_config -import logos_delivery/messaging/rate_limit_manager/rate_limit_manager +import logos_delivery/messaging/rate_limit_manager/rate_limit_config -export kernel_conf, rate_limit_manager +export kernel_conf, rate_limit_config # `LogosDeliveryMode` and `EntryLayer` are defined at the leaf (`cli_args`) so # they can appear on `WakuNodeConf`; re-exported here via `kernel_conf`. @@ -52,9 +52,12 @@ type MessagingClientConf* = object nodeKey* {.name: "nodekey".}: Opt[crypto.PrivateKey] ## P2P node private key (64-char hex): stable identity / peerId across restarts. rateLimit*: Opt[RateLimitConfig] - ## Per-epoch message rate limit enforced by the send service. `Opt` like - ## every other field so `merge` propagates a caller's override; unset falls + ## Per-epoch message rate limit enforced by the send service; unset falls ## back to `DefaultRateLimitConfig` (rate limiting disabled). + ## + ## Settable only programmatically: as a nested object with no `{.name.}` + ## pragma or `parseCmdArg`, it is not reachable from the JSON config or a + ## CLI flag. proc applyMode*(conf: var WakuNodeConf, mode: LogosDeliveryMode): ConfResult[void] = ## Sets the protocol flags implied by the mode. diff --git a/logos_delivery/messaging/delivery_service/send_service/delivery_task.nim b/logos_delivery/messaging/delivery_service/send_service/delivery_task.nim index f0c9d9942..5caa0919d 100644 --- a/logos_delivery/messaging/delivery_service/send_service/delivery_task.nim +++ b/logos_delivery/messaging/delivery_service/send_service/delivery_task.nim @@ -26,6 +26,11 @@ type DeliveryTask* = ref object firstPropagatedTime*: Opt[Moment] ## Set once on the first successful propagation; never reset on re-publish. ## Anchors the store-validation time cap (see propagationAge). + firstAdmittedTime*: Opt[Moment] + ## Set when the task first passes rate-limit admission; `none` while parked + ## waiting for epoch budget. Guards re-admission on retry and anchors the + ## delivery-timeout reaper, so a task parked for budget is not aged out + ## before it can be sent. propagateEventEmitted*: bool errorDesc*: string @@ -59,30 +64,47 @@ proc new*( func `==`*(r, l: DeliveryTask): bool = if r.isNil() == l.isNil(): - r.isNil() or r.msgHash == l.msgHash + return r.isNil() or r.msgHash == l.msgHash else: - false + return false proc messageAge*(self: DeliveryTask): timer.Duration = let actual = getNanosecondTime(getTime().toUnixFloat()) if self.msg.timestamp >= 0 and self.msg.timestamp < actual: - nanoseconds(actual - self.msg.timestamp) + return nanoseconds(actual - self.msg.timestamp) else: - ZeroDuration + return ZeroDuration proc deliveryAge*(self: DeliveryTask): timer.Duration = if self.state == DeliveryState.SuccessfullyPropagated: - timer.Moment.now() - self.deliveryTime + return timer.Moment.now() - self.deliveryTime else: - ZeroDuration + return ZeroDuration proc propagationAge*(self: DeliveryTask): timer.Duration = ## Time elapsed since the message was first successfully propagated. ## Stable across re-publishes; ZeroDuration until first propagation. if self.firstPropagatedTime.isSome(): - timer.Moment.now() - self.firstPropagatedTime.get() + return timer.Moment.now() - self.firstPropagatedTime.get() else: - ZeroDuration + return ZeroDuration + +proc admissionAge*(self: DeliveryTask): timer.Duration = + ## Time since the task first passed admission; ZeroDuration while never + ## admitted (still parked waiting for epoch budget). + if self.firstAdmittedTime.isSome(): + return timer.Moment.now() - self.firstAdmittedTime.get() + else: + return ZeroDuration + +proc isDeliveryTimedOut*(self: DeliveryTask, maxTime: timer.Duration): bool = + ## True when an admitted task has been trying to deliver longer than `maxTime` + ## without ever propagating. A task never admitted (parked for budget) is + ## exempt: the clock runs from admission time, so waiting for budget does not count + ## against it. + return + self.firstAdmittedTime.isSome() and self.firstPropagatedTime.isNone() and + self.admissionAge() > maxTime proc isEphemeral*(self: DeliveryTask): bool = return self.msg.ephemeral diff --git a/logos_delivery/messaging/delivery_service/send_service/send_service.nim b/logos_delivery/messaging/delivery_service/send_service/send_service.nim index feec01953..18e43a2a3 100644 --- a/logos_delivery/messaging/delivery_service/send_service/send_service.nim +++ b/logos_delivery/messaging/delivery_service/send_service/send_service.nim @@ -48,8 +48,8 @@ type SendService* = ref object of RootObj serviceLoopHandle: Future[void] ## handle that allows to stop the async task sendProcessor: BaseSendProcessor rateLimitManager: RateLimitManager - ## Meters first transmissions against the per-epoch budget; re-publishes - ## of an already-propagated message resend the same bytes and are free. + ## Charges first transmissions against the per-epoch budget; re-publishes + ## are free. waku: Waku checkStoreForMessages: bool @@ -85,7 +85,10 @@ proc new*( preferP2PReliability: bool, waku: Waku, rateLimitManager: RateLimitManager, + sendProcessor: BaseSendProcessor = nil, ): Result[T, string] = + ## `sendProcessor` overrides the relay/lightpush chain built from `waku`, + ## letting a caller drive the scheduler against a scripted delivery outcome. if not waku.hasRelay() and not waku.hasLightpush(): return err( "Could not create SendService. wakuRelay or wakuLightpushClient should be set" @@ -93,8 +96,12 @@ proc new*( let checkStoreForMessages = preferP2PReliability and waku.isStoreMounted() - let sendProcessorChain = setupSendProcessorChain(waku, waku.brokerCtx).valueOr: - return err("failed to setup SendProcessorChain: " & $error) + let sendProcessorChain = + if sendProcessor.isNil(): + setupSendProcessorChain(waku, waku.brokerCtx).valueOr: + return err("failed to setup SendProcessorChain: " & $error) + else: + sendProcessor let sendService = SendService( brokerCtx: waku.brokerCtx, @@ -197,15 +204,14 @@ proc reportTaskResult(self: SendService, task: DeliveryTask) = # rest of the states are intermediate and does not translate to event discard - # Only tasks that never propagated are reported as hard send failures here. - # Propagated-but-not-store-validated tasks are handled (warn + drop, no event) - # in evaluateAndCleanUp. - if task.firstPropagatedTime.isNone() and task.messageAge() > MaxTimeInCache: + # Hard-fail a task admitted but never propagated within MaxTimeInCache. + # Propagated-but-unvalidated tasks are dropped in evaluateAndCleanUp instead. + if task.isDeliveryTimedOut(MaxTimeInCache): error "Failed to send message", requestId = task.requestId, msgHash = task.msgHash.to0xHex(), error = "Message too old", - age = task.messageAge() + age = task.admissionAge() task.state = DeliveryState.FailedToDeliver MessageErrorEvent.emit( self.brokerCtx, @@ -248,17 +254,27 @@ proc evaluateAndCleanUp(self: SendService) = ) ) -proc trySendMessages(self: SendService) {.async.} = +proc admitOnce(self: SendService, task: DeliveryTask): Future[bool] {.async.} = + ## Charges the task's first transmission against the epoch budget, at most + ## once per task (`firstAdmittedTime`); retries then resend for free. Returns + ## false when the task must stay parked for a later epoch. + if task.firstAdmittedTime.isSome(): + return true + + (await self.rateLimitManager.admit(task.msg.payload)).isOkOr: + debug "over rate-limit budget, task waits for the epoch to roll", + requestId = task.requestId, msgHash = task.msgHash.to0xHex() + return false + task.firstAdmittedTime = Opt.some(Moment.now()) + return true + +proc trySendMessages*(self: SendService) {.async.} = let tasksToSend = self.taskCache.filterIt(it.state == DeliveryState.NextRoundRetry) for task in tasksToSend: # Todo, check if it has any perf gain to run them concurrent... - if task.firstPropagatedTime.isNone(): - ## Never propagated: this transmission consumes a fresh epoch slot, so - ## it must be admitted. Tasks that stay over budget remain in - ## `NextRoundRetry` and are retried as the epoch rolls over. - (await self.rateLimitManager.admit(task.msg.payload)).isOkOr: - continue + if not (await self.admitOnce(task)): + continue await self.sendProcessor.process(task) proc serviceLoop(self: SendService) {.async.} = @@ -288,8 +304,8 @@ proc send*(self: SendService, task: DeliveryTask) {.async.} = error "SendService.send: failed to subscribe to content topic", contentTopic = task.msg.contentTopic, error = error - (await self.rateLimitManager.admit(task.msg.payload)).isOkOr: - info "SendService.send: over rate-limit budget, parking task", + if not (await self.admitOnce(task)): + info "SendService.send: parking task for a later round", requestId = task.requestId, msgHash = task.msgHash.to0xHex() task.state = DeliveryState.NextRoundRetry self.addTask(task) diff --git a/logos_delivery/messaging/messaging_client.nim b/logos_delivery/messaging/messaging_client.nim index c33b3ff6a..25804cdf7 100644 --- a/logos_delivery/messaging/messaging_client.nim +++ b/logos_delivery/messaging/messaging_client.nim @@ -6,6 +6,7 @@ import logos_delivery/api/conf/messaging_conf, logos_delivery/api/messaging_client_api, logos_delivery/waku/waku, + logos_delivery/waku/api/publish, logos_delivery/waku/factory/conf_builder/waku_conf_builder, logos_delivery/messaging/delivery_service/[recv_service, send_service], logos_delivery/messaging/rate_limit_manager/rate_limit_manager @@ -19,15 +20,26 @@ type MessagingClient* = ref object recvService*: RecvService started*: bool +proc rlnQuotaProvider(waku: Waku): QuotaProvider = + ## Sources the rate limit manager's epoch and limit from RLN. The closure + ## queries `waku` on each admission, so a node whose RLN mounts after + ## construction upgrades from the wall-clock fallback automatically. + return proc(): Opt[EpochQuota] {.gcsafe, raises: [].} = + let q = waku.currentRlnEpochQuota().valueOr: + return Opt.none(EpochQuota) + return + Opt.some(EpochQuota(epochIndex: q.epochIndex, userMessageLimit: q.messageLimit)) + proc new*( T: type MessagingClient, conf: MessagingClientConf, waku: Waku ): Result[T, string] = ## The messaging layer chains onto Waku: it drives the underlying Waku kernel ## for transport while exposing its own send/recv API. let reliability = conf.reliabilityEnabled.get(DefaultP2pReliability) - let sendService = ?SendService.new( - reliability, waku, RateLimitManager.new(conf.rateLimit.get(DefaultRateLimitConfig)) + let rateLimitManager = ?RateLimitManager.new( + conf.rateLimit.get(DefaultRateLimitConfig), rlnQuotaProvider(waku) ) + let sendService = ?SendService.new(reliability, waku, rateLimitManager) let recvService = RecvService.new(waku) return ok( T( diff --git a/logos_delivery/messaging/rate_limit_manager/quota_source.nim b/logos_delivery/messaging/rate_limit_manager/quota_source.nim new file mode 100644 index 000000000..cf656f105 --- /dev/null +++ b/logos_delivery/messaging/rate_limit_manager/quota_source.nim @@ -0,0 +1,22 @@ +## Epoch + user-message-limit source for the rate limit manager. +## +## Both values are read through one provider call, so a read cannot pair a +## fresh epoch with a stale limit. The provider is a callback, keeping the +## manager free of any `Waku` dependency. + +import std/times +import results + +type + EpochQuota* = object + epochIndex*: uint64 ## Current epoch (`timestamp div` epoch length). + userMessageLimit*: uint64 ## Messages the epoch grants. + + QuotaProvider* = proc(): Opt[EpochQuota] {.gcsafe, raises: [].} + ## `none` when RLN is not mounted — the signal to fall back to the + ## wall clock. + +proc wallClockEpochIndex*(epochPeriodSec: uint64): uint64 = + ## Absolute epoch (`unixTime div epochPeriodSec`), the same derivation RLN + ## uses, so independent nodes agree on the boundary. + return uint64(getTime().toUnix()) div epochPeriodSec diff --git a/logos_delivery/messaging/rate_limit_manager/rate_limit_config.nim b/logos_delivery/messaging/rate_limit_manager/rate_limit_config.nim new file mode 100644 index 000000000..793407462 --- /dev/null +++ b/logos_delivery/messaging/rate_limit_manager/rate_limit_config.nim @@ -0,0 +1,25 @@ +## Configuration for the messaging rate limit manager. +## +## Kept separate from `rate_limit_manager` so `messaging_conf` can depend on +## `RateLimitConfig` without pulling in the manager itself. + +type + RateLimitError* {.pure.} = enum + OverBudget + + RateLimitConfig* = object + enabled*: bool + epochPeriodSec*: uint64 + ## Epoch length in seconds. Shapes only the wall-clock fallback window; + ## ignored once the RLN quota source supplies the period. + messagesPerEpoch*: uint64 + ## Local cap on messages admitted per epoch. When RLN is mounted the cap + ## is clamped to RLN's limit + +const + DefaultEpochPeriodSec* = 600'u64 + DefaultMessagesPerEpoch* = 1'u64 + + DefaultRateLimitConfig* = RateLimitConfig( + epochPeriodSec: DefaultEpochPeriodSec, messagesPerEpoch: DefaultMessagesPerEpoch + ) ## Used when no rate-limit config is supplied; `enabled` defaults false. diff --git a/logos_delivery/messaging/rate_limit_manager/rate_limit_manager.nim b/logos_delivery/messaging/rate_limit_manager/rate_limit_manager.nim index ba4944dca..d0c059622 100644 --- a/logos_delivery/messaging/rate_limit_manager/rate_limit_manager.nim +++ b/logos_delivery/messaging/rate_limit_manager/rate_limit_manager.nim @@ -1,58 +1,80 @@ ## Rate Limit Manager for the Messaging API. ## -## Tracks messages sent per RLN epoch and rejects admission when the -## limit is approached, ensuring RLN compliance on enforcing relays. +## Rate-limits message transmissions against the per-epoch user message limit, +## rejecting admission once the epoch's budget is spent. The epoch rolling +## over refills the budget. ## -## For the skeleton this is a pass-through: every call is admitted. -## Real per-epoch budgeting will use `queue`, `currentEpochStart`, -## `sentInCurrentEpoch`, and `resetEpoch` to park messages and admit -## them as the epoch rolls over. -## -## See: https://lip.logos.co/messaging/raw/reliable-channel-api.html +## The epoch and limit come from a `QuotaProvider` when RLN is mounted; +## otherwise a wall-clock window and the configured limit stand in. Parking and +## retrying over-budget messages is the send service's job — this module only +## answers whether one more transmission fits the current epoch. -import std/times import results, chronos -type - RateLimitError* {.pure.} = enum - OverBudget +import ./rate_limit_config, ./quota_source - RateLimitConfig* = object - enabled*: bool ## spec: rate limiting opt-in; SHOULD be true when RLN active - epochPeriodSec*: int - messagesPerEpoch*: int +export rate_limit_config, quota_source - RateLimitManager* = ref object - config*: RateLimitConfig - queue*: seq[seq[byte]] - currentEpochStart*: Time - sentInCurrentEpoch*: int +type RateLimitManager* = ref object + config*: RateLimitConfig + quotaProvider: QuotaProvider + ## Nil or a `none` result selects the wall-clock fallback. Queried per + ## admission so a late RLN mount upgrades automatically. + currentEpochIndex*: uint64 + sentInCurrentEpoch*: uint64 -const - DefaultEpochPeriodSec* = 600 - DefaultMessagesPerEpoch* = 1 +proc new*( + T: type RateLimitManager, + config: RateLimitConfig, + quotaProvider: QuotaProvider = nil, +): Result[T, string] = + ## Rejects an enabled config with a zero epoch period: the wall-clock + ## fallback derives the epoch as `unixTime div epochPeriodSec`. + if config.enabled and config.epochPeriodSec == 0: + return err("rate limit config: epochPeriodSec must be positive when enabled") - DefaultRateLimitConfig* = RateLimitConfig( - epochPeriodSec: DefaultEpochPeriodSec, messagesPerEpoch: DefaultMessagesPerEpoch - ) ## Used when no rate-limit config is supplied; `enabled` defaults false. + return ok( + T( + config: config, + quotaProvider: quotaProvider, + currentEpochIndex: 0, + sentInCurrentEpoch: 0, + ) + ) -proc new*(T: type RateLimitManager, config: RateLimitConfig): T = - return - T(config: config, queue: @[], currentEpochStart: getTime(), sentInCurrentEpoch: 0) +proc currentQuota(self: RateLimitManager): Opt[EpochQuota] = + if self.quotaProvider.isNil(): + return Opt.none(EpochQuota) + return self.quotaProvider() proc admit*( self: RateLimitManager, msg: seq[byte] ): Future[Result[void, RateLimitError]] {.async: (raises: []).} = - ## Skeleton behaviour: admits immediately. Real per-epoch budgeting - ## will consult `config`, `sentInCurrentEpoch`, and the elapsed - ## `epochPeriodSec` window before admitting or parking `msg`. + ## Charges one message against the current epoch's limit, rolling the window + ## first when the epoch has advanced. A disabled config admits everything. + if not self.config.enabled: + return ok() + + let quota = self.currentQuota() + + let epochIndex = + if quota.isSome(): + quota.get().epochIndex + else: + wallClockEpochIndex(self.config.epochPeriodSec) + + # RLN can only tighten the configured cap, never widen it: exceeding RLN's + # limit would fail later at proof generation. + var limit = self.config.messagesPerEpoch + if quota.isSome() and quota.get().userMessageLimit < limit: + limit = quota.get().userMessageLimit + + if epochIndex != self.currentEpochIndex: + self.currentEpochIndex = epochIndex + self.sentInCurrentEpoch = 0 + + if self.sentInCurrentEpoch >= limit: + return err(RateLimitError.OverBudget) + + self.sentInCurrentEpoch.inc() return ok() - -proc dequeueReady*(self: RateLimitManager): seq[seq[byte]] = - ## Returns the set of queued messages that may be dispatched now - ## without exceeding the configured rate limit. - discard - -proc resetEpoch*(self: RateLimitManager) = - self.currentEpochStart = getTime() - self.sentInCurrentEpoch = 0 diff --git a/logos_delivery/waku/api/publish.nim b/logos_delivery/waku/api/publish.nim index d1dffd5a0..8347c6057 100644 --- a/logos_delivery/waku/api/publish.nim +++ b/logos_delivery/waku/api/publish.nim @@ -42,6 +42,17 @@ proc relayPushHandler*(self: Waku): PushMessageHandler = ## in (legacy)lightpushPublish; this handler only validates and republishes. return getRelayPushHandler(self.node.wakuRelay) +proc currentRlnEpochQuota*(self: Waku): Opt[tuple[epochIndex, messageLimit: uint64]] = + ## RLN's current epoch index and user message limit, read together so the + ## pair cannot straddle an epoch boundary. + if self.node.rln.isNil(): + return Opt.none(tuple[epochIndex, messageLimit: uint64]) + + let limit = self.node.rln.groupManager.userMessageLimit.valueOr: + return Opt.none(tuple[epochIndex, messageLimit: uint64]) + + return Opt.some((fromEpoch(self.node.rln.getCurrentEpoch()), uint64(limit))) + proc lightpushPeerAvailable*(self: Waku, shard: PubsubTopic): bool = ## True if a lightpush service peer is available for `shard`. return self.node.peerManager.selectPeer(WakuLightPushCodec, Opt.some(shard)).isSome() diff --git a/tests/messaging/test_all.nim b/tests/messaging/test_all.nim index 3cf508033..c38b65c4e 100644 --- a/tests/messaging/test_all.nim +++ b/tests/messaging/test_all.nim @@ -1,3 +1,4 @@ {.used.} -import ./test_rate_limit_manager +import + ./test_rate_limit_manager, ./test_delivery_task_reaping, ./test_send_service_scheduler diff --git a/tests/messaging/test_delivery_task_reaping.nim b/tests/messaging/test_delivery_task_reaping.nim new file mode 100644 index 000000000..1bb385074 --- /dev/null +++ b/tests/messaging/test_delivery_task_reaping.nim @@ -0,0 +1,36 @@ +{.used.} + +import results, chronos, testutils/unittests + +import logos_delivery/messaging/delivery_service/send_service/delivery_task + +const MaxTime = chronos.minutes(1) + +proc taskWith(admitted, propagated: Opt[Moment]): DeliveryTask = + ## Builds a DeliveryTask directly (bypassing `new`, which needs a broker). + return DeliveryTask(firstAdmittedTime: admitted, firstPropagatedTime: propagated) + +suite "DeliveryTask - delivery-timeout reaping": + test "a task parked for budget (never admitted) is exempt, however old": + ## Budget-parked tasks must survive to the epoch roll, not be aged out. + let task = taskWith(Opt.none(Moment), Opt.none(Moment)) + check not task.isDeliveryTimedOut(MaxTime) + + test "an admitted, never-propagated task past the window times out": + let task = taskWith(Opt.some(Moment.now() - chronos.minutes(2)), Opt.none(Moment)) + check task.isDeliveryTimedOut(MaxTime) + + test "an admitted task still within the window does not time out": + let task = taskWith(Opt.some(Moment.now()), Opt.none(Moment)) + check not task.isDeliveryTimedOut(MaxTime) + + test "a propagated task is never timed out here (store validation owns it)": + let task = + taskWith(Opt.some(Moment.now() - chronos.minutes(2)), Opt.some(Moment.now())) + check not task.isDeliveryTimedOut(MaxTime) + + test "the timeout clock runs from admission, not message creation": + ## A just-admitted task has a fresh clock, even after a long budget wait. + let task = taskWith(Opt.some(Moment.now()), Opt.none(Moment)) + check task.admissionAge() < MaxTime + check not task.isDeliveryTimedOut(MaxTime) diff --git a/tests/messaging/test_rate_limit_manager.nim b/tests/messaging/test_rate_limit_manager.nim index 895fcb49d..4d0dc8592 100644 --- a/tests/messaging/test_rate_limit_manager.nim +++ b/tests/messaging/test_rate_limit_manager.nim @@ -1,26 +1,98 @@ {.used.} -import chronos, testutils/unittests, stew/byteutils +import results, chronos, testutils/unittests, stew/byteutils import logos_delivery/messaging/rate_limit_manager/rate_limit_manager +proc fixedQuota(epochIndex, userMessageLimit: uint64): QuotaProvider = + ## A quota source pinned to one epoch, so limit-boundary tests don't touch + ## the wall clock. + return proc(): Opt[EpochQuota] {.gcsafe, raises: [].} = + return + Opt.some(EpochQuota(epochIndex: epochIndex, userMessageLimit: userMessageLimit)) + suite "RateLimitManager - admission": asyncTest "admit is a pass-through when disabled": - let rl = RateLimitManager.new( - RateLimitConfig(enabled: false, epochPeriodSec: 600, messagesPerEpoch: 1) - ) + let rl = RateLimitManager + .new(RateLimitConfig(enabled: false, epochPeriodSec: 600, messagesPerEpoch: 1)) + .expect("RateLimitManager.new") for _ in 0 ..< 10: - let res = await rl.admit("payload".toBytes()) - check res.isOk() + check (await rl.admit("payload".toBytes())).isOk() - asyncTest "admit is a pass-through in the skeleton even when enabled": - ## Documents the current skeleton behaviour: per-epoch enforcement is - ## not wired yet, so every call is admitted regardless of the - ## configured budget. This test flips to red as soon as real - ## enforcement lands, at which point it should be replaced with - ## budget-boundary assertions. - let rl = RateLimitManager.new( - RateLimitConfig(enabled: true, epochPeriodSec: 600, messagesPerEpoch: 1) - ) + asyncTest "a zero cap admits nothing": + let rl = RateLimitManager + .new(RateLimitConfig(enabled: true, epochPeriodSec: 600, messagesPerEpoch: 0)) + .expect("RateLimitManager.new") + let res = await rl.admit("a".toBytes()) + check: + res.isErr() + res.error == RateLimitError.OverBudget + + test "construction validates the epoch period only when enabled": + check: + RateLimitManager + .new(RateLimitConfig(enabled: true, epochPeriodSec: 0, messagesPerEpoch: 1)) + .isErr() + RateLimitManager.new(RateLimitConfig()).isOk() + + asyncTest "admits up to the message limit then rejects with OverBudget": + ## Fixed-epoch quota so the window cannot roll mid-test. + let rl = RateLimitManager + .new( + RateLimitConfig(enabled: true, epochPeriodSec: 600, messagesPerEpoch: 3), + fixedQuota(epochIndex = 42, userMessageLimit = 100), + ) + .expect("RateLimitManager.new") + for i in 0 ..< 3: + check (await rl.admit(("msg" & $i).toBytes())).isOk() + let res = await rl.admit("over".toBytes()) + check: + res.isErr() + res.error == RateLimitError.OverBudget + + asyncTest "budget refills when the epoch rolls over": + ## Drive the roll through the provider — no sleeps, no flake. + var epoch = 1'u64 + let rl = RateLimitManager + .new( + RateLimitConfig(enabled: true, epochPeriodSec: 600, messagesPerEpoch: 1), + proc(): Opt[EpochQuota] {.gcsafe, raises: [].} = + Opt.some(EpochQuota(epochIndex: epoch, userMessageLimit: 100)), + ) + .expect("RateLimitManager.new") check (await rl.admit("first".toBytes())).isOk() - check (await rl.admit("second".toBytes())).isOk() + check (await rl.admit("second".toBytes())).isErr() + epoch = 2 + check (await rl.admit("third".toBytes())).isOk() + check (await rl.admit("fourth".toBytes())).isErr() + + asyncTest "RLN user message limit clamps a looser configured cap": + ## config cap 5, RLN grants 2 — the lower RLN limit wins. + let rl = RateLimitManager + .new( + RateLimitConfig(enabled: true, epochPeriodSec: 600, messagesPerEpoch: 5), + fixedQuota(epochIndex = 7, userMessageLimit = 2), + ) + .expect("RateLimitManager.new") + check (await rl.admit("a".toBytes())).isOk() + check (await rl.admit("b".toBytes())).isOk() + check (await rl.admit("c".toBytes())).isErr() + + asyncTest "config cap can tighten below the RLN limit": + ## config cap 1, RLN grants 100 — the config cap wins. + let rl = RateLimitManager + .new( + RateLimitConfig(enabled: true, epochPeriodSec: 600, messagesPerEpoch: 1), + fixedQuota(epochIndex = 7, userMessageLimit = 100), + ) + .expect("RateLimitManager.new") + check (await rl.admit("a".toBytes())).isOk() + check (await rl.admit("b".toBytes())).isErr() + + asyncTest "falls back to the wall-clock window when no quota source is set": + ## No provider: rate limiting still enforces within a single wall-clock epoch. + let rl = RateLimitManager + .new(RateLimitConfig(enabled: true, epochPeriodSec: 600, messagesPerEpoch: 1)) + .expect("RateLimitManager.new") + check (await rl.admit("first".toBytes())).isOk() + check (await rl.admit("second".toBytes())).isErr() diff --git a/tests/messaging/test_send_service_scheduler.nim b/tests/messaging/test_send_service_scheduler.nim new file mode 100644 index 000000000..c7e2ae948 --- /dev/null +++ b/tests/messaging/test_send_service_scheduler.nim @@ -0,0 +1,150 @@ +{.used.} + +import std/net +import chronos, chronicles, testutils/unittests, results, stew/byteutils + +import + logos_delivery/waku/waku, + logos_delivery/waku/waku_core, + logos_delivery/api/types, + logos_delivery/api/conf/messaging_conf, + logos_delivery/waku/factory/waku_conf, + logos_delivery/messaging/rate_limit_manager/rate_limit_manager, + logos_delivery/messaging/delivery_service/send_service/ + [send_service, send_processor, delivery_task] +import ../testlib/testasync + +## Scheduler-level coverage for the send service's rate-limit seam: a task is +## charged exactly once however many rounds it takes, and an over-budget task +## is parked then released when the epoch rolls. A fake processor scripts the +## delivery outcome so the loop runs without network or sleeps. + +type FakeSendProcessor = ref object of BaseSendProcessor + calls: int + script: seq[DeliveryState] + ## State to stamp on the task per invocation; the last entry repeats. + +method process(self: FakeSendProcessor, task: DeliveryTask): Future[void] {.async.} = + let outcome = self.script[min(self.calls, self.script.high)] + inc self.calls + task.state = outcome + if outcome == DeliveryState.SuccessfullyPropagated and + task.firstPropagatedTime.isNone(): + task.firstPropagatedTime = Opt.some(Moment.now()) + +proc testConf(): WakuConf = + var conf = MessagingClientConf() + .toWakuNodeConf(messaging_conf.LogosDeliveryMode.Core).valueOr: + raiseAssert error + conf.listenAddress = parseIpAddress("0.0.0.0") + conf.tcpPort = Port(0) + conf.discv5UdpPort = Port(0) + conf.clusterId = Opt.some(3'u16) + conf.numShardsInNetwork = 1 + conf.rest = false + return conf.toWakuConf().valueOr: + raiseAssert error + +proc fixedEpochQuota(epoch: ptr uint64, userMessageLimit: uint64): QuotaProvider = + ## Quota pinned to whatever `epoch` holds, so a test rolls the epoch by + ## writing through the pointer. + return proc(): Opt[EpochQuota] {.gcsafe, raises: [].} = + return Opt.some(EpochQuota(epochIndex: epoch[], userMessageLimit: userMessageLimit)) + +suite "SendService - rate-limit scheduling": + var waku {.threadvar.}: Waku + + asyncSetup: + waku = (await Waku.new(testConf())).expect("Waku.new") + + asyncTeardown: + ## The node is never started, so stop is best-effort cleanup. + discard await waku.stop() + + proc buildTask(id, payload: string): DeliveryTask = + ## Built directly rather than via `DeliveryTask.new`, which needs a broker + ## provider only registered once the node starts. + let msg = WakuMessage( + contentTopic: "/test/1/scheduler/proto", + payload: payload.toBytes(), + timestamp: 1_700_000_000_000_000_000, + ) + let pubsubTopic = PubsubTopic("/waku/2/rs/3/0") + return DeliveryTask( + requestId: RequestId(id), + pubsubTopic: pubsubTopic, + msg: msg, + msgHash: computeMessageHash(pubsubTopic, msg), + state: DeliveryState.Entry, + ) + + asyncTest "a task is charged once even when delivery takes several rounds": + ## First round fails to propagate, second succeeds. The retry must not draw a + ## second slot: `firstAdmittedTime` guards re-admission. + var epoch = 5'u64 + let manager = RateLimitManager + .new( + RateLimitConfig(enabled: true, epochPeriodSec: 600, messagesPerEpoch: 3), + fixedEpochQuota(addr epoch, userMessageLimit = 100), + ) + .expect("RateLimitManager.new") + let processor = FakeSendProcessor( + script: @[DeliveryState.NextRoundRetry, DeliveryState.SuccessfullyPropagated] + ) + let service = + SendService.new(false, waku, manager, processor).expect("SendService.new") + + let task = buildTask("charge-once", "hi") + await service.send(task) + check: + manager.sentInCurrentEpoch == 1'u64 + task.firstAdmittedTime.isSome() + task.state == DeliveryState.NextRoundRetry + + await service.trySendMessages() + check: + manager.sentInCurrentEpoch == 1'u64 # not re-charged on retry + processor.calls == 2 + task.state == DeliveryState.SuccessfullyPropagated + + asyncTest "an over-budget task is parked, then released when the epoch rolls": + ## Budget of one per epoch. The second send is parked until the epoch rolls, + ## then admitted and delivered. + var epoch = 1'u64 + let manager = RateLimitManager + .new( + RateLimitConfig(enabled: true, epochPeriodSec: 600, messagesPerEpoch: 1), + fixedEpochQuota(addr epoch, userMessageLimit = 100), + ) + .expect("RateLimitManager.new") + let processor = FakeSendProcessor(script: @[DeliveryState.SuccessfullyPropagated]) + let service = + SendService.new(false, waku, manager, processor).expect("SendService.new") + + let first = buildTask("in-budget", "one") + await service.send(first) + check: + first.state == DeliveryState.SuccessfullyPropagated + manager.sentInCurrentEpoch == 1'u64 + + let second = buildTask("over-budget", "two") + await service.send(second) + check: + second.state == DeliveryState.NextRoundRetry # parked + second.firstAdmittedTime.isNone() # never admitted + let callsWhenParked = processor.calls + + # Same epoch: still over budget, so the parked task is not handed to the + # processor. + await service.trySendMessages() + check: + second.state == DeliveryState.NextRoundRetry + second.firstAdmittedTime.isNone() + processor.calls == callsWhenParked + + # Epoch rolls: budget refills, the parked task is admitted and delivered. + epoch = 2'u64 + await service.trySendMessages() + check: + second.firstAdmittedTime.isSome() + second.state == DeliveryState.SuccessfullyPropagated