From 38773dd754b14b828ba12bc105bb13d6ab427cd6 Mon Sep 17 00:00:00 2001 From: Stevenson Michel <130018170+thoven87@users.noreply.github.com> Date: Sat, 10 Oct 2026 19:28:43 -0400 Subject: [PATCH 1/2] Add rate-limit units scaling support Extend `RateLimit` with a `units` parameter (default `1`) so one dispatch can consume multiple tokens from a shared bucket. `slotParams` now scales `intervalMs` by `units`, and values are clamped to at least 1 to avoid invalid configuration. Expanded tests cover pure `slotParams` behavior (base interval, scaling, key stability, clamping, per-key buckets) and an integration scenario proving a high-unit dispatch advances the shared bucket cursor by multiple slots. --- Sources/Strand/Activity.swift | 17 +++- Tests/StrandTests/RateLimitTests.swift | 124 ++++++++++++++++++++++++- 2 files changed, 137 insertions(+), 4 deletions(-) diff --git a/Sources/Strand/Activity.swift b/Sources/Strand/Activity.swift index a73588b..2a4724e 100644 --- a/Sources/Strand/Activity.swift +++ b/Sources/Strand/Activity.swift @@ -120,14 +120,25 @@ public struct RateLimit: Sendable { /// per-interval drip begins. Default `0` = no burst. public var burst: Int + /// Number of rate-limit tokens this dispatch consumes from the shared bucket. + /// Default `1` = one dispatch charges one slot. Set to the number of + /// external API calls made inside the activity so the bucket stays + /// calibrated to actual outbound call volume rather than dispatch count. + /// + /// Example: a batch activity that calls a third-party API 50 times per run + /// should set `units: 50` so the cursor advances 50× further on each + /// dispatch, correctly throttling peer activities on the same bucket. + public var units: Int + /// Returns `nil` when `limit` is not positive; no crash. - public init?(limit: Double, period: Duration = .seconds(1), key: String? = nil, sharedKey: String? = nil, burst: Int = 0) { + public init?(limit: Double, period: Duration = .seconds(1), key: String? = nil, sharedKey: String? = nil, burst: Int = 0, units: Int = 1) { guard limit > 0 else { return nil } self.limit = limit self.period = period self.key = key self.sharedKey = sharedKey self.burst = burst + self.units = max(1, units) } // MARK: - Internal @@ -139,11 +150,11 @@ public struct RateLimit: Sendable { /// or as the prefix in `"ActivityName:entityKey"` (per-entity bucket). /// - Returns: `slotKey` — the row key in `strand.rate_limit_slots`; /// `intervalMs` — milliseconds between consecutive slots, - /// derived from `period / limit` and clamped to at least 1 ms; + /// derived from `(period / limit) × units` and clamped to at least 1 ms; /// `burstSlots` — number of tasks eligible for immediate dispatch. internal func slotParams(for activityName: String) -> (slotKey: String, intervalMs: Int, burstSlots: Int) { let slotKey = sharedKey ?? key.map { "\(activityName):\($0)" } ?? activityName - let intervalMs = max(1, Int(Double(period.milliseconds) / limit)) + let intervalMs = max(1, Int(Double(period.milliseconds) / limit)) * units return (slotKey: slotKey, intervalMs: intervalMs, burstSlots: max(0, burst)) } } diff --git a/Tests/StrandTests/RateLimitTests.swift b/Tests/StrandTests/RateLimitTests.swift index ea95a9c..9af3eb1 100644 --- a/Tests/StrandTests/RateLimitTests.swift +++ b/Tests/StrandTests/RateLimitTests.swift @@ -76,7 +76,69 @@ private func fetchAvailableAt( return result } -// MARK: - Tests +// MARK: - Unit tests + +/// Pure in-process tests for `RateLimit.slotParams` — no database required. +/// +/// Uses GitHub's Search API limit (30 req/min authenticated) as a concrete, +/// publicly documented example: +/// base interval = 60,000 ms / 30 = 2,000 ms per slot +/// +/// The batch scenario — fetching detail records for each of 50 search results +/// in one dispatch — is realistic: search returns 50 items, the next step +/// calls the detail endpoint for every item (50 outbound calls). +@Suite("RateLimit.slotParams — units scaling") +struct RateLimitSlotParamsTests { + + // GitHub Search API (authenticated): 30 requests / minute. + // One search request per dispatch → interval = 60,000 ms / 30 = 2,000 ms. + @Test("units=1 produces the base intervalMs") + func unitsOneIsBaseInterval() { + let rl = RateLimit(limit: 30, period: .minutes(1), sharedKey: "github")! + let params = rl.slotParams(for: "SearchReposActivity") + #expect(params.intervalMs == 2_000) + #expect(params.slotKey == "github") + } + + // Fetching detail records for each of the 50 search results = 50 API calls + // per dispatch. units=50 advances the cursor by 50 × 2,000 ms = 100,000 ms, + // so the bucket correctly throttles the next activity by the full cost. + @Test("units=50 scales intervalMs by 50, same shared bucket") + func unitsFiftyScalesInterval() { + let rl = RateLimit(limit: 30, period: .minutes(1), sharedKey: "github", units: 50)! + let params = rl.slotParams(for: "FetchRepoDetailsBatchActivity") + #expect(params.intervalMs == 100_000) + #expect(params.slotKey == "github") // sharedKey is unchanged by units + } + + // A single-call dispatch (units=1) and a 50-call batch dispatch (units=50) + // must drain from the same "github" bucket so their combined load is correctly + // accounted for. + @Test("units=1 and units=50 resolve to the same slot key") + func unitsDontChangeSlotKey() { + let single = RateLimit(limit: 30, period: .minutes(1), sharedKey: "github", units: 1)! + let batch = RateLimit(limit: 30, period: .minutes(1), sharedKey: "github", units: 50)! + #expect(single.slotParams(for: "A").slotKey == batch.slotParams(for: "B").slotKey) + } + + @Test("units clamped to 1: zero or negative units treated as 1") + func unitsClampedToOne() { + let rl = RateLimit(limit: 30, period: .minutes(1), sharedKey: "github", units: 0)! + #expect(rl.units == 1) + #expect(rl.slotParams(for: "A").intervalMs == 2_000) + } + + @Test("per-key bucket: units scales intervalMs without changing the key prefix") + func unitsWithPerKeyBucket() { + let rl = RateLimit(limit: 10, period: .seconds(1), key: "customer:42", units: 5)! + let params = rl.slotParams(for: "PaymentActivity") + // base = 1,000 ms / 10 = 100 ms; × 5 units = 500 ms + #expect(params.intervalMs == 500) + #expect(params.slotKey == "PaymentActivity:customer:42") + } +} + +// MARK: - Integration tests @Suite("Integration — Rate limiting", .tags(.integration), .serialized) struct RateLimitTests { @@ -392,4 +454,64 @@ struct RateLimitTests { ) } } + + // ── 7 ──────────────────────────────────────────────────────────────────────────────── + // units > 1: a high-unit dispatch must advance the shared-bucket cursor + // by N × intervalMs rather than 1 × intervalMs. + // + // Setup: all three tasks share the same sharedKey (unique per test run). + // e1 — units=1, limit=1/10s: immediate; cursor → NOW+10 s + // e2 — units=5, limit=1/10s: available ≈ NOW+10 s; cursor → NOW+60 s + // e3 — units=1, limit=1/10s: available ≈ NOW+60 s + // + // Gap e2→e3 must be ~50 s (5 × 10 s), proving the units=5 dispatch consumed + // five slots’ worth of budget from the shared bucket. + @Test("units > 1: dispatch consumes N slots from the shared bucket") + func unitsConsumeMultipleSlots() async throws { + try await withTestEnvironment { client in + // Unique key so this test’s bucket doesn’t bleed into other runs. + let key = "units-test-\(UUID().uuidString)" + let single = RateLimit(limit: 1, period: .seconds(10), sharedKey: key)! + let five = RateLimit(limit: 1, period: .seconds(10), sharedKey: key, units: 5)! + + let e1 = try await client.enqueueActivity( + EchoActivity.self, input: .init(value: "1"), options: .init(rateLimit: single) + ) + let e2 = try await client.enqueueActivity( + EchoActivity.self, input: .init(value: "2"), options: .init(rateLimit: five) + ) + let e3 = try await client.enqueueActivity( + EchoActivity.self, input: .init(value: "3"), options: .init(rateLimit: single) + ) + + let dates = try await fetchAvailableAt( + postgres: client.postgres, + taskIDs: [e1.taskID, e2.taskID, e3.taskID], + logger: client.logger + ) + + let at1 = try #require(dates[e1.taskID]) + let at2 = try #require(dates[e2.taskID]) + let at3 = try #require(dates[e3.taskID]) + let now = Date.now + + // e1: first in a cold bucket — immediately available. + #expect(at1 <= now.addingTimeInterval(0.5), "e1 should be immediate") + + // e2: delayed by one single-unit interval from e1 (~10 s). + let gap12 = at2.timeIntervalSince(at1) + #expect( + gap12 >= 9.0 && gap12 <= 11.0, + "gap e1→e2 should be ~10 s (1 unit × 10 s), got \(gap12) s" + ) + + // e3: delayed by five intervals from e2 (~50 s) because units=5 + // advanced the cursor by 5 × 10 s = 50 s. + let gap23 = at3.timeIntervalSince(at2) + #expect( + gap23 >= 49.0 && gap23 <= 51.0, + "gap e2→e3 should be ~50 s (5 units × 10 s), got \(gap23) s" + ) + } + } } From 37118ce922b178454c3118678413c9aadc79cbc6 Mon Sep 17 00:00:00 2001 From: Stevenson Michel <130018170+thoven87@users.noreply.github.com> Date: Sat, 10 Oct 2026 20:50:58 -0400 Subject: [PATCH 2/2] Add typed dynamic keys for rate limits Refactors `RateLimit` into its own file and extends it with `.keyed(...)` factories that derive per-entity keys from typed input via closures or KeyPaths. Enqueue paths now resolve dynamic keys while typed input is still available (`StrandClient` and `WorkflowContext`) before encoding/scheduling. Slot-key behavior was also clarified for `sharedKey` + `key` composition, and tests were expanded to cover dynamic key extraction, shared bucket partitioning, no-op resolution, and end-to-end per-entity bucket timing. --- Sources/Strand/Activity.swift | 85 ---------- Sources/Strand/Client.swift | 3 +- Sources/Strand/RateLimit.swift | 215 +++++++++++++++++++++++++ Sources/Strand/WorkflowContext.swift | 24 ++- Tests/StrandTests/RateLimitTests.swift | 149 ++++++++++++++++- 5 files changed, 383 insertions(+), 93 deletions(-) create mode 100644 Sources/Strand/RateLimit.swift diff --git a/Sources/Strand/Activity.swift b/Sources/Strand/Activity.swift index 2a4724e..dc66e6d 100644 --- a/Sources/Strand/Activity.swift +++ b/Sources/Strand/Activity.swift @@ -74,91 +74,6 @@ extension ActivityCancellationType: PostgresCodable { } } -// MARK: - RateLimit - -/// Leaky-bucket rate limit for an activity. When set, each enqueue atomically -/// claims a time slot; the run waits in PENDING until its slot arrives. -/// -/// Only enforced when explicitly set — `nil` (the default) means no rate limit -/// and the activity is enqueued at full speed, exactly as without this option. -/// -/// - `limit` / `period`: together define the slot interval. -/// `.init(limit: 10, period: .seconds(1))` → one slot every 100 ms. -/// `.init(limit: 2, period: .minutes(1))` → one slot every 30 s. -/// -/// - `key`: optional partition key that gives each entity its own independent -/// bucket. `nil` = one global bucket shared by all executions of this -/// activity type on this queue. Use the workflow's own computed value — -/// there is no server-side expression language; Swift code is the expression. -/// -/// ```swift -/// // Global: 10 payment activities per second across all customers -/// options: .init(rateLimit: .init(limit: 10, period: .seconds(1))) -/// -/// // Per-customer: 2 enrichment calls per minute per customer ID -/// options: .init( -/// rateLimit: .init(limit: 2, period: .minutes(1), key: "customer:\(customerID)") -/// ) -/// ``` -public struct RateLimit: Sendable { - /// Maximum executions allowed within `period`. - public var limit: Double - /// Time window over which `limit` applies. Default: `.seconds(1)`. - public var period: Duration - /// Optional partition key. `nil` = one global bucket for this activity - /// type on this queue. A non-nil value gives each entity its own bucket. - public var key: String? - - /// When set, this string is used verbatim as the slot key — no activity-name - /// prefix is applied. Use this to share a single rate-limit bucket across - /// multiple activity types (e.g. all HHAX activities sharing "hhax"). - /// When both `sharedKey` and `key` are set, `sharedKey` wins. - public var sharedKey: String? - - /// Number of tasks that may be dispatched immediately from a cold bucket (or - /// from any bucket that still has burst tokens remaining) before the regular - /// per-interval drip begins. Default `0` = no burst. - public var burst: Int - - /// Number of rate-limit tokens this dispatch consumes from the shared bucket. - /// Default `1` = one dispatch charges one slot. Set to the number of - /// external API calls made inside the activity so the bucket stays - /// calibrated to actual outbound call volume rather than dispatch count. - /// - /// Example: a batch activity that calls a third-party API 50 times per run - /// should set `units: 50` so the cursor advances 50× further on each - /// dispatch, correctly throttling peer activities on the same bucket. - public var units: Int - - /// Returns `nil` when `limit` is not positive; no crash. - public init?(limit: Double, period: Duration = .seconds(1), key: String? = nil, sharedKey: String? = nil, burst: Int = 0, units: Int = 1) { - guard limit > 0 else { return nil } - self.limit = limit - self.period = period - self.key = key - self.sharedKey = sharedKey - self.burst = burst - self.units = max(1, units) - } - - // MARK: - Internal - - /// Slot-allocation parameters for the `strand.rate_limit_slots` SQL. - /// - /// - Parameter activityName: the activity's registered name. - /// Used as the bucket prefix when `key` is `nil` (global bucket), - /// or as the prefix in `"ActivityName:entityKey"` (per-entity bucket). - /// - Returns: `slotKey` — the row key in `strand.rate_limit_slots`; - /// `intervalMs` — milliseconds between consecutive slots, - /// derived from `(period / limit) × units` and clamped to at least 1 ms; - /// `burstSlots` — number of tasks eligible for immediate dispatch. - internal func slotParams(for activityName: String) -> (slotKey: String, intervalMs: Int, burstSlots: Int) { - let slotKey = sharedKey ?? key.map { "\(activityName):\($0)" } ?? activityName - let intervalMs = max(1, Int(Double(period.milliseconds) / limit)) * units - return (slotKey: slotKey, intervalMs: intervalMs, burstSlots: max(0, burst)) - } -} - // MARK: - ActivityOptions /// Per-activity execution configuration: timeout, retry, routing, priority, and fairness. diff --git a/Sources/Strand/Client.swift b/Sources/Strand/Client.swift index 7de762b..e382883 100644 --- a/Sources/Strand/Client.swift +++ b/Sources/Strand/Client.swift @@ -104,7 +104,8 @@ public struct StrandClient: Sendable { input: A.Input, options: ActivityOptions = .init() ) async throws -> EnqueueResult { - let rlParams = options.rateLimit.map { $0.slotParams(for: A.name) } + // Resolve any dynamic key extractor while the typed input is still available. + let rlParams = options.rateLimit.map { $0.resolving(input: input).slotParams(for: A.name) } let resolvedID = options.id ?? A.generateActivityID() return try await _enqueue( queue: options.queue ?? queueName, diff --git a/Sources/Strand/RateLimit.swift b/Sources/Strand/RateLimit.swift new file mode 100644 index 0000000..bfc8ec1 --- /dev/null +++ b/Sources/Strand/RateLimit.swift @@ -0,0 +1,215 @@ +#if canImport(FoundationEssentials) +import FoundationEssentials +#else +import Foundation +#endif + +// MARK: - RateLimit + +/// Leaky-bucket rate limit applied when enqueuing an activity, workflow, or +/// generic task. When set, each enqueue atomically claims a time slot; the +/// run waits in PENDING until its slot arrives. +/// +/// Only enforced when explicitly set — `nil` (the default) means no rate limit +/// and the task is enqueued at full speed, exactly as without this option. +/// +/// - `limit` / `period`: together define the slot interval. +/// `.init(limit: 10, period: .seconds(1))` → one slot every 100 ms. +/// `.init(limit: 2, period: .minutes(1))` → one slot every 30 s. +/// +/// - `key`: static partition key — one bucket per unique value. +/// +/// ```swift +/// // Global: 10 payment activities per second across all customers +/// options: .init(rateLimit: .init(limit: 10, period: .seconds(1))) +/// +/// // Per-customer static key — computed at the call site +/// options: .init( +/// rateLimit: .init(limit: 2, period: .minutes(1), key: "customer:\(customerID)") +/// ) +/// ``` +public struct RateLimit: Sendable { + /// Maximum executions allowed within `period`. + public var limit: Double + /// Time window over which `limit` applies. Default: `.seconds(1)`. + public var period: Duration + /// Optional partition key. `nil` = one global bucket for this task + /// type on this queue. A non-nil value gives each entity its own bucket. + public var key: String? + + /// When set, this string is used as the slot-key prefix — no task-name + /// prefix is applied. Use this to share a rate-limit bucket across multiple + /// task types (e.g. all payment tasks sharing `"payments"`). + /// + /// - When only `sharedKey` is set the slot key is `sharedKey` verbatim + /// (one global bucket shared by all matching tasks). + /// - When BOTH `sharedKey` and `key` are set the slot key is + /// `"sharedKey:key"` — per-entity buckets within the shared namespace. + /// Example: `sharedKey: "payments", key: "tenant-acme"` → `"payments:tenant-acme"`. + public var sharedKey: String? + + /// Number of tasks that may be dispatched immediately from a cold bucket (or + /// from any bucket that still has burst tokens remaining) before the regular + /// per-interval drip begins. Default `0` = no burst. + public var burst: Int + + /// Number of rate-limit tokens this dispatch consumes from the shared bucket. + /// Default `1` = one dispatch charges one slot. Set to the number of + /// external API calls made inside the activity so the bucket stays + /// calibrated to actual outbound call volume rather than dispatch count. + /// + /// Example: a batch activity that calls a third-party API 50 times per run + /// should set `units: 50` so the cursor advances 50× further on each + /// dispatch, correctly throttling peer tasks on the same bucket. + public var units: Int + + /// Derives the per-entity partition key from the task's typed input at + /// enqueue time. Set via the ``keyed(_:period:sharedKey:burst:units:by:)`` + /// factory, which wraps the `any Sendable` cast in a typed closure and + /// adds a debug assertion for type mismatches. + /// + /// `internal` — callers should use `keyed(...)`, not set this directly. + /// + /// `nonisolated(unsafe)` because the closure is called exactly once, + /// synchronously, inside `resolving(input:)` before anything crosses an + /// isolation boundary. The closure must be a pure function; `keyed` + /// enforces this by construction (KeyPath reads and field accesses). + nonisolated(unsafe) var keyExtractor: ((any Sendable) -> String?)? + + /// Returns `nil` when `limit` is not positive; no crash. + public init?(limit: Double, period: Duration = .seconds(1), key: String? = nil, sharedKey: String? = nil, burst: Int = 0, units: Int = 1) { + guard limit > 0 else { return nil } + self.limit = limit + self.period = period + self.key = key + self.sharedKey = sharedKey + self.burst = burst + self.units = max(1, units) + self.keyExtractor = nil + } + + // MARK: - Dynamic key factories + + /// Creates a rate limit whose per-entity key is derived from the task's + /// input at enqueue time via a typed closure. + /// + /// The closure is evaluated while `Input` is still a concrete Swift + /// value — before it is JSON-encoded. This lets you declare the key schema + /// once rather than repeating `key: input.tenantID` at every call site. + /// + /// ```swift + /// // Declare once — Input type is inferred from the closure: + /// static let rateLimitOpts = ActivityOptions( + /// rateLimit: .keyed(limit: 100, period: .minutes(1)) { + /// (input: ProcessPaymentInput) in input.customerID + /// } + /// ) + /// + /// // Every call site just passes the shared options: + /// try await context.runActivity(ProcessPayment.self, input: input, options: .rateLimitOpts) + /// ``` + static func keyed( + limit: Double, + period: Duration = .seconds(1), + sharedKey: String? = nil, + burst: Int = 0, + units: Int = 1, + by keyExtractor: @escaping (Input) -> String? + ) -> RateLimit? { + // Optional.map on the failable init handles the limit <= 0 guard; + // the transform sets keyExtractor on the valid copy and returns it. + RateLimit(limit: limit, period: period, sharedKey: sharedKey, burst: burst, units: units).map { base in + var rl = base + rl.keyExtractor = { anyInput in + guard let typed = anyInput as? Input else { + assertionFailure( + "RateLimit keyExtractor: expected \(Input.self), got \(type(of: anyInput)). " + + "Ensure the RateLimit is used with the matching task's Input type." + ) + return nil + } + return keyExtractor(typed) + } + return rl + } + } + + /// KeyPath overload — derives the per-entity key from a non-optional property. + /// + /// ```swift + /// rateLimit: .keyed(limit: 100, period: .minutes(1), by: \ProcessPaymentInput.customerID) + /// ``` + static func keyed( + limit: Double, + period: Duration = .seconds(1), + sharedKey: String? = nil, + burst: Int = 0, + units: Int = 1, + by keyPath: KeyPath + ) -> RateLimit? { + keyed(limit: limit, period: period, sharedKey: sharedKey, burst: burst, units: units) { + (input: Input) in input[keyPath: keyPath] + } + } + + /// KeyPath overload — derives the per-entity key from an optional property. + /// Returns `nil` (global bucket) when the property itself is `nil`. + /// + /// ```swift + /// rateLimit: .keyed(limit: 100, period: .minutes(1), by: \SyncVisitsPageInput.officeID) + /// ``` + static func keyed( + limit: Double, + period: Duration = .seconds(1), + sharedKey: String? = nil, + burst: Int = 0, + units: Int = 1, + by keyPath: KeyPath + ) -> RateLimit? { + keyed(limit: limit, period: period, sharedKey: sharedKey, burst: burst, units: units) { + (input: Input) in input[keyPath: keyPath] + } + } + + // MARK: - Internal + + /// Slot-allocation parameters for the `strand.rate_limit_slots` SQL. + /// + /// - Parameter taskName: the task's registered name (activity or workflow name). + /// Used as the bucket prefix when `key` is `nil` (global bucket), + /// or as the prefix in `"TaskName:entityKey"` (per-entity bucket). + /// - Returns: `slotKey` — the row key in `strand.rate_limit_slots`; + /// `intervalMs` — milliseconds between consecutive slots, + /// derived from `(period / limit) × units` and clamped to at least 1 ms; + /// `burstSlots` — number of tasks eligible for immediate dispatch. + /// + /// - Note: Uses `self.key` directly. When a `keyExtractor` is present, + /// call sites must resolve it first via `resolving(input:)` so the + /// per-entity key is written into `self.key` before this is invoked. + internal func slotParams(for taskName: String) -> (slotKey: String, intervalMs: Int, burstSlots: Int) { + let slotKey: String + if let sk = sharedKey { + // sharedKey as namespace: append per-entity key when present. + slotKey = key.map { "\(sk):\($0)" } ?? sk + } else { + slotKey = key.map { "\(taskName):\($0)" } ?? taskName + } + let intervalMs = max(1, Int(Double(period.milliseconds) / limit)) * units + return (slotKey: slotKey, intervalMs: intervalMs, burstSlots: max(0, burst)) + } + + /// Resolves `key` in-place from `input` via `keyExtractor`. + /// No-op when no extractor is set. + internal mutating func resolve(input: any Sendable) { + guard let extractor = keyExtractor else { return } + key = extractor(input) + } + + /// Returns a copy with `key` resolved from `input`. + /// Used where a mutating call isn’t possible (e.g. inside `Optional.map`). + internal func resolving(input: any Sendable) -> RateLimit { + var copy = self + copy.resolve(input: input) + return copy + } +} diff --git a/Sources/Strand/WorkflowContext.swift b/Sources/Strand/WorkflowContext.swift index 5822787..79a0564 100644 --- a/Sources/Strand/WorkflowContext.swift +++ b/Sources/Strand/WorkflowContext.swift @@ -902,6 +902,10 @@ public struct WorkflowContext: Sendable { // throw CancellationError here rather than scheduling a new activity that would // never be needed. Fast paths above are unaffected — already-completed work replays. try Task.checkCancellation() + // Resolve any dynamic rate-limit key while the typed input is still available. + // WorkflowScheduling only has a ByteBuffer at that point, so this is the + // last moment a KeyPath / closure extractor can be evaluated. + let resolvedOptions = options.resolvingRateLimit(input) let inputBuffer = try _impl.codec.encode(input) // Idempotency key is always derived from the run UUID and sequence number. // Run UUID (not task UUID) scopes the key to this run so that continueAsNew @@ -921,7 +925,7 @@ public struct WorkflowContext: Sendable { .scheduleActivity( name: A.name, input: inputBuffer, - options: options, + options: resolvedOptions, seqNum: seqNum, idempotencyKey: idempotencyKey ) @@ -1869,3 +1873,21 @@ public struct ChildWorkflowHandle: Sendable { try await _result() } } + +// MARK: - fileprivate helpers + +extension ActivityOptions { + /// Returns a copy of these options with the rate-limit key resolved from + /// `input` via `RateLimit.keyExtractor`, or `self` unchanged when no + /// extractor is set. + /// + /// Only needed in `WorkflowContext.runActivity`, which must pass a fully + /// resolved `ActivityOptions` into the `WorkflowCommand` — the typed input + /// is unavailable once it has been encoded to a `ByteBuffer`. + fileprivate func resolvingRateLimit(_ input: any Sendable) -> ActivityOptions { + guard rateLimit?.keyExtractor != nil else { return self } + var copy = self + copy.rateLimit?.resolve(input: input) + return copy + } +} diff --git a/Tests/StrandTests/RateLimitTests.swift b/Tests/StrandTests/RateLimitTests.swift index 9af3eb1..812e811 100644 --- a/Tests/StrandTests/RateLimitTests.swift +++ b/Tests/StrandTests/RateLimitTests.swift @@ -108,7 +108,7 @@ struct RateLimitSlotParamsTests { let rl = RateLimit(limit: 30, period: .minutes(1), sharedKey: "github", units: 50)! let params = rl.slotParams(for: "FetchRepoDetailsBatchActivity") #expect(params.intervalMs == 100_000) - #expect(params.slotKey == "github") // sharedKey is unchanged by units + #expect(params.slotKey == "github") // sharedKey is unchanged by units } // A single-call dispatch (units=1) and a 50-call batch dispatch (units=50) @@ -117,10 +117,84 @@ struct RateLimitSlotParamsTests { @Test("units=1 and units=50 resolve to the same slot key") func unitsDontChangeSlotKey() { let single = RateLimit(limit: 30, period: .minutes(1), sharedKey: "github", units: 1)! - let batch = RateLimit(limit: 30, period: .minutes(1), sharedKey: "github", units: 50)! + let batch = RateLimit(limit: 30, period: .minutes(1), sharedKey: "github", units: 50)! #expect(single.slotParams(for: "A").slotKey == batch.slotParams(for: "B").slotKey) } + // keyed factory: closure + // keyed(limit:period:by:) with a typed closure derives the slot key from the + // activity's input at enqueue time instead of requiring it at the call site. + @Test("keyed factory: closure derives slot key from typed input") + func keyedClosureDrivesSlotKey() { + struct UserInput: Sendable { let userID: String } + + let rl = RateLimit.keyed(limit: 30, period: .minutes(1)) { + (input: UserInput) in input.userID + }! + + let resolved = rl.resolving(input: UserInput(userID: "user-42")) + let params = resolved.slotParams(for: "ProcessPayment") + // slot key = "ProcessPayment:user-42" (no sharedKey, so activity-name prefix applies) + #expect(params.slotKey == "ProcessPayment:user-42") + #expect(params.intervalMs == 2_000) + } + + // keyed factory: KeyPath (non-optional) + @Test("keyed factory: non-optional KeyPath derives slot key from typed input") + func keyedKeyPathNonOptional() { + struct UserInput: Sendable { let userID: String } + + let rl = RateLimit.keyed(limit: 30, period: .minutes(1), by: \UserInput.userID)! + + let resolved = rl.resolving(input: UserInput(userID: "user-99")) + #expect(resolved.slotParams(for: "A").slotKey == "A:user-99") + } + + // keyed factory: KeyPath (optional) — nil value falls back to global bucket + @Test("keyed factory: optional KeyPath falls back to global bucket when nil") + func keyedKeyPathOptional() { + struct PageInput: Sendable { let officeID: String? } + + let rl = RateLimit.keyed(limit: 30, period: .minutes(1), by: \PageInput.officeID)! + + // Non-nil value → per-entity bucket + let resolvedSome = rl.resolving(input: PageInput(officeID: "office-7")) + #expect(resolvedSome.slotParams(for: "SyncPage").slotKey == "SyncPage:office-7") + + // Nil value → global bucket (activity-name key, no partition) + let resolvedNil = rl.resolving(input: PageInput(officeID: nil)) + #expect(resolvedNil.slotParams(for: "SyncPage").slotKey == "SyncPage") + } + + // keyed + sharedKey — extractor partitions WITHIN the shared bucket + @Test("keyed factory: extractor partitions within a sharedKey bucket") + func keyedWithSharedKey() { + struct TenantInput: Sendable { let tenantID: String } + + // Two different activity types, same sharedKey, same extractor. + let rl = RateLimit.keyed(limit: 100, period: .minutes(1), sharedKey: "api") { + (input: TenantInput) in input.tenantID + }! + + // sharedKey wins over the derived key for the bucket PREFIX, + // but the partition is appended: "api:tenant-A". + let paramsA = rl.resolving(input: TenantInput(tenantID: "tenant-A")).slotParams(for: "ActivityX") + let paramsB = rl.resolving(input: TenantInput(tenantID: "tenant-B")).slotParams(for: "ActivityY") + + // Different tenants → different slot keys despite different activity names. + #expect(paramsA.slotKey == "api:tenant-A") + #expect(paramsB.slotKey == "api:tenant-B") + #expect(paramsA.slotKey != paramsB.slotKey) + } + + // resolving(input:) is a no-op when no extractor is present + @Test("resolving(input:) is a no-op without a keyExtractor") + func resolvingNoOpWithoutExtractor() { + let rl = RateLimit(limit: 30, period: .minutes(1), key: "static-key")! + let resolved = rl.resolving(input: "anything") + #expect(resolved.slotParams(for: "A").slotKey == "A:static-key") + } + @Test("units clamped to 1: zero or negative units treated as 1") func unitsClampedToOne() { let rl = RateLimit(limit: 30, period: .minutes(1), sharedKey: "github", units: 0)! @@ -472,16 +546,22 @@ struct RateLimitTests { // Unique key so this test’s bucket doesn’t bleed into other runs. let key = "units-test-\(UUID().uuidString)" let single = RateLimit(limit: 1, period: .seconds(10), sharedKey: key)! - let five = RateLimit(limit: 1, period: .seconds(10), sharedKey: key, units: 5)! + let five = RateLimit(limit: 1, period: .seconds(10), sharedKey: key, units: 5)! let e1 = try await client.enqueueActivity( - EchoActivity.self, input: .init(value: "1"), options: .init(rateLimit: single) + EchoActivity.self, + input: .init(value: "1"), + options: .init(rateLimit: single) ) let e2 = try await client.enqueueActivity( - EchoActivity.self, input: .init(value: "2"), options: .init(rateLimit: five) + EchoActivity.self, + input: .init(value: "2"), + options: .init(rateLimit: five) ) let e3 = try await client.enqueueActivity( - EchoActivity.self, input: .init(value: "3"), options: .init(rateLimit: single) + EchoActivity.self, + input: .init(value: "3"), + options: .init(rateLimit: single) ) let dates = try await fetchAvailableAt( @@ -514,4 +594,61 @@ struct RateLimitTests { ) } } + + // ── 8 ──────────────────────────────────────────────────────────────────────────────── + // keyExtractor: two activities with DIFFERENT inputs derive DIFFERENT slot keys + // from the same RateLimit declaration, landing in independent per-entity buckets. + // + // Setup: one shared RateLimit created with .keyed(..., by: \EchoInput.value). + // Each of the three enqueues passes a different EchoInput.value. + // + // Expected available_at schedule (limit: 1/10 s, each input is its own bucket): + // e1 (value: "a") → immediate bucket "EchoActivity:a" + // e2 (value: "b") → immediate bucket "EchoActivity:b" (independent) + // e3 (value: "a") → ~10 s delayed bucket "EchoActivity:a" (same as e1) + @Test("keyExtractor: different inputs land in independent per-entity buckets") + func keyExtractorCreatesPerEntityBuckets() async throws { + try await withTestEnvironment { client in + let rl = RateLimit.keyed(limit: 1, period: .seconds(10), by: \EchoInput.value)! + + let e1 = try await client.enqueueActivity( + EchoActivity.self, + input: .init(value: "a"), + options: .init(rateLimit: rl) + ) + let e2 = try await client.enqueueActivity( + EchoActivity.self, + input: .init(value: "b"), + options: .init(rateLimit: rl) + ) + // Same input value as e1 — should share e1’s bucket and be delayed. + let e3 = try await client.enqueueActivity( + EchoActivity.self, + input: .init(value: "a"), + options: .init(rateLimit: rl) + ) + + let dates = try await fetchAvailableAt( + postgres: client.postgres, + taskIDs: [e1.taskID, e2.taskID, e3.taskID], + logger: client.logger + ) + + let at1 = try #require(dates[e1.taskID]) + let at2 = try #require(dates[e2.taskID]) + let at3 = try #require(dates[e3.taskID]) + let now = Date.now + + // e1 and e2 have different keys — both immediately available. + #expect(at1 <= now.addingTimeInterval(0.5), "e1 (bucket \"a\") should be immediate") + #expect(at2 <= now.addingTimeInterval(0.5), "e2 (bucket \"b\") should be immediate") + + // e3 shares e1’s bucket — delayed ~10 s. + let gap13 = at3.timeIntervalSince(at1) + #expect( + gap13 >= 9.0 && gap13 <= 11.0, + "e3 should be ~10 s after e1 (same bucket \"a\"), got \(gap13) s" + ) + } + } }