feat(analytics): add demand-adaptive token-bucket flush governor

ajaysubra · claude · ajaysubra · commit e4d114b8d0e9 · 2026-06-23T22:02:02.000-04:00
Replace the rigid fixed-interval flush decision with a token-bucket gate so
the SDK adapts to demand: bursts of activity after an idle period flush
immediately (the bucket banks up to `flushTokenBucketCapacity` tokens), while
sustained activity is capped at the long-term refill rate. Refill rate is
derived from the existing network-aware `flushInterval` (wifi/cell), and the
bucket is frozen while offline.

Also add a queue-depth early-flush trigger (mirrors the Android SDK's batch
flush depth) so large bursts flush early instead of waiting for the next
interval tick; the token bucket still gates whether the early flush proceeds.

Defaults preserve today's steady-state cadence — no behavior change unless the
queue bursts or sustained load would otherwise exceed the refill rate.

Token-bucket state is transient (not persisted), so the bucket starts full on
each launch and the first post-launch flush proceeds immediately.

Co-Authored-By: Claude Opus 4.8 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/Sources/KlaviyoSwift/StateManagement/KlaviyoState+FlushTokenBucket.swift b/Sources/KlaviyoSwift/StateManagement/KlaviyoState+FlushTokenBucket.swift
@@ -0,0 +1,56 @@
+//
+//  KlaviyoState+FlushTokenBucket.swift
+//
+//  Demand-adaptive flush governor: a token-bucket rate limiter plus a queue-depth
+//  early-flush trigger. Together these let the SDK absorb bursts of activity after
+//  idle periods while capping the long-term flush rate, instead of flushing on a
+//  rigid fixed interval regardless of demand.
+//
+//  Copyright (c) 2026 Klaviyo
+//  Licensed under the MIT License. See LICENSE file in the project root for full license information.
+//
+
+import Foundation
+
+extension KlaviyoState {
+    /// `true` when the queue has grown large enough to warrant an early flush attempt
+    /// rather than waiting for the next flush-interval tick. Suppressed when the
+    /// `flushInterval` is non-finite (i.e. offline) so we don't kick off doomed requests.
+    var shouldFlushForQueueDepth: Bool {
+        flushInterval.isFinite && queue.count >= StateManagementConstants.flushDepth
+    }
+
+    /// Refills the flush-token bucket based on the time elapsed since the last refill, then
+    /// attempts to consume a single token.
+    ///
+    /// Tokens accrue at a rate of `1 / flushInterval` per second (the same network-aware
+    /// cadence used by the flush timer) and are capped at `flushTokenBucketCapacity`. Because
+    /// the bucket can bank up to `capacity` tokens during quiet periods, a burst of flushes
+    /// immediately after an idle stretch is allowed through, while sustained activity is
+    /// throttled to the long-term refill rate. When offline (`flushInterval` is non-finite)
+    /// the bucket is frozen — no tokens accrue until connectivity returns.
+    ///
+    /// - Parameter currentTime: The current time, injected for testability.
+    /// - Returns: `true` if a token was available and consumed (the caller may flush);
+    ///   `false` if the bucket is empty (the caller should defer until tokens refill).
+    mutating func consumeFlushToken(currentTime: Date) -> Bool {
+        if let lastRefill = lastFlushTokenRefill,
+           flushInterval.isFinite, flushInterval > 0 {
+            let elapsed = currentTime.timeIntervalSince(lastRefill)
+            if elapsed > 0 {
+                let refillRate = 1.0 / flushInterval // tokens per second
+                availableFlushTokens = min(
+                    StateManagementConstants.flushTokenBucketCapacity,
+                    availableFlushTokens + elapsed * refillRate
+                )
+            }
+        }
+        lastFlushTokenRefill = currentTime
+
+        guard availableFlushTokens >= 1.0 else {
+            return false
+        }
+        availableFlushTokens -= 1.0
+        return true
+    }
+}
diff --git a/Sources/KlaviyoSwift/StateManagement/KlaviyoState.swift b/Sources/KlaviyoSwift/StateManagement/KlaviyoState.swift
@@ -62,6 +62,13 @@ struct KlaviyoState: Equatable, Codable {
     var pendingProfile: [Profile.ProfileKey: AnyEncodable]?
     var isProcessingDeepLink = false
 
+    // token bucket related stuff
+    // These govern the flush cadence (see `consumeFlushToken(now:)`). They are intentionally
+    // transient — left out of `CodingKeys` below — so the bucket starts full on every launch,
+    // letting the first post-launch flush proceed immediately.
+    var availableFlushTokens = StateManagementConstants.flushTokenBucketCapacity
+    var lastFlushTokenRefill: Date?
+
     enum CodingKeys: CodingKey {
         case apiKey
         case email
diff --git a/Sources/KlaviyoSwift/StateManagement/StateManagement.swift b/Sources/KlaviyoSwift/StateManagement/StateManagement.swift
@@ -24,6 +24,19 @@ enum StateManagementConstants {
     static let wifiFlushInterval = 10.0
     static let maxQueueSize = 200
     static let initialAttempt = 1
+
+    /// Maximum number of flushes the token bucket can hold. Each flush consumes one
+    /// token; tokens refill over time at a rate derived from `flushInterval` (the
+    /// network-aware cadence). A capacity greater than 1 lets the SDK absorb a burst
+    /// of flushes after an idle period while still capping the long-term flush rate at
+    /// the refill rate — see `KlaviyoState.consumeFlushToken(currentTime:)`.
+    static let flushTokenBucketCapacity = 5.0
+
+    /// Queue depth that triggers an early flush attempt instead of waiting for the next
+    /// flush-interval tick. Mirrors the Android SDK's batch flush depth so the burst
+    /// behavior is consistent across platforms. The token bucket still gates whether the
+    /// early flush actually proceeds.
+    static let flushDepth = 25
 }
 
 /// Describes how the state machine should handle retrying a request after a failure.
@@ -318,6 +331,14 @@ struct KlaviyoReducer: ReducerProtocol {
                 return .none
             }
 
+            // Gate the flush on the token bucket: this enforces the long-term flush rate
+            // while still allowing bursts immediately after idle periods. If no token is
+            // available we defer — the next flush-interval tick (or queue-depth trigger)
+            // will retry once tokens have refilled.
+            guard state.consumeFlushToken(currentTime: environment.date()) else {
+                return .none
+            }
+
             state.requestsInFlight.append(contentsOf: state.queue)
             state.queue.removeAll()
             state.flushing = true
@@ -513,7 +534,11 @@ struct KlaviyoReducer: ReducerProtocol {
                 state.enqueueRequest(request: request)
             }
 
-            let baseEffect = shouldPrioritize ? EffectTask<KlaviyoAction>.task { .flushQueue } : .none
+            // Prioritized events flush immediately; otherwise flush early only once the queue
+            // reaches the depth threshold. The token bucket in `.flushQueue` decides whether the
+            // attempt actually proceeds, so this stays within the long-term rate limit.
+            let shouldFlushNow = shouldPrioritize || state.shouldFlushForQueueDepth
+            let baseEffect = shouldFlushNow ? EffectTask<KlaviyoAction>.task { .flushQueue } : .none
             return .merge([
                 baseEffect,
                 .fireAndForget { KlaviyoInternal.publishEvent(event) }
@@ -531,7 +556,7 @@ struct KlaviyoReducer: ReducerProtocol {
 
             state.enqueueRequest(request: request)
 
-            return .none
+            return state.shouldFlushForQueueDepth ? EffectTask<KlaviyoAction>.task { .flushQueue } : .none
 
         case let .enqueueProfile(profile):
             guard case .initialized = state.initalizationState
diff --git a/Tests/KlaviyoSwiftTests/FlushTokenBucketTests.swift b/Tests/KlaviyoSwiftTests/FlushTokenBucketTests.swift
@@ -0,0 +1,150 @@
+//
+//  FlushTokenBucketTests.swift
+//
+//  Tests for the demand-adaptive (token-bucket) flush governor and the
+//  queue-depth early-flush trigger.
+//
+//  Copyright (c) 2026 Klaviyo
+//  Licensed under the MIT License. See LICENSE file in the project root for full license information.
+//
+
+@testable import KlaviyoCore
+@testable import KlaviyoSwift
+import Foundation
+import XCTest
+
+final class FlushTokenBucketTests: XCTestCase {
+    @MainActor
+    override func setUp() async throws {
+        environment = KlaviyoEnvironment.test()
+        klaviyoSwiftEnvironment = KlaviyoSwiftEnvironment.test()
+    }
+
+    private func sampleRequest(name: String = "test") -> KlaviyoRequest {
+        KlaviyoRequest(endpoint: .createEvent("foo", CreateEventPayload(data: .init(name: name))))
+    }
+
+    // MARK: - consumeFlushToken(currentTime:)
+
+    func test_consumeFlushToken_fullBucketAllowsFlushAndDecrementsByOne() {
+        var state = KlaviyoState(queue: [])
+        state.flushInterval = StateManagementConstants.wifiFlushInterval
+
+        let flushTime = Date(timeIntervalSince1970: 1000)
+        XCTAssertTrue(state.consumeFlushToken(currentTime: flushTime))
+        XCTAssertEqual(
+            state.availableFlushTokens,
+            StateManagementConstants.flushTokenBucketCapacity - 1,
+            accuracy: 0.0001
+        )
+        XCTAssertEqual(state.lastFlushTokenRefill, flushTime)
+    }
+
+    func test_consumeFlushToken_emptyBucketIsDeniedUntilEnoughTimeElapses() {
+        var state = KlaviyoState(queue: [])
+        state.flushInterval = StateManagementConstants.wifiFlushInterval // 10s -> 0.1 token/sec
+        state.availableFlushTokens = 0
+        let start = Date(timeIntervalSince1970: 1000)
+        state.lastFlushTokenRefill = start
+
+        // After 5s only half a token has accrued -> still denied.
+        XCTAssertFalse(state.consumeFlushToken(currentTime: start.addingTimeInterval(5)))
+        XCTAssertEqual(state.availableFlushTokens, 0.5, accuracy: 0.0001)
+
+        // 10 more seconds accrues a full token -> allowed, then decremented.
+        XCTAssertTrue(state.consumeFlushToken(currentTime: start.addingTimeInterval(15)))
+        XCTAssertEqual(state.availableFlushTokens, 0.5, accuracy: 0.0001)
+    }
+
+    func test_consumeFlushToken_refillIsCappedAtCapacity() {
+        var state = KlaviyoState(queue: [])
+        state.flushInterval = StateManagementConstants.wifiFlushInterval
+        state.availableFlushTokens = 4
+        let start = Date(timeIntervalSince1970: 1000)
+        state.lastFlushTokenRefill = start
+
+        // 100s would add 10 tokens, but the bucket is capped at capacity before consuming.
+        XCTAssertTrue(state.consumeFlushToken(currentTime: start.addingTimeInterval(100)))
+        XCTAssertEqual(
+            state.availableFlushTokens,
+            StateManagementConstants.flushTokenBucketCapacity - 1,
+            accuracy: 0.0001
+        )
+    }
+
+    func test_consumeFlushToken_doesNotRefillWhenOffline() {
+        var state = KlaviyoState(queue: [])
+        state.flushInterval = .infinity // offline: timer paused, bucket frozen
+        state.availableFlushTokens = 0
+        let start = Date(timeIntervalSince1970: 1000)
+        state.lastFlushTokenRefill = start
+
+        XCTAssertFalse(state.consumeFlushToken(currentTime: start.addingTimeInterval(10_000)))
+        XCTAssertEqual(state.availableFlushTokens, 0, accuracy: 0.0001)
+        XCTAssertEqual(state.lastFlushTokenRefill, start.addingTimeInterval(10_000))
+    }
+
+    // MARK: - shouldFlushForQueueDepth
+
+    func test_shouldFlushForQueueDepth_belowThresholdIsFalse() {
+        let queue = Array(repeating: sampleRequest(), count: StateManagementConstants.flushDepth - 1)
+        var state = KlaviyoState(queue: queue)
+        state.flushInterval = StateManagementConstants.wifiFlushInterval
+        XCTAssertFalse(state.shouldFlushForQueueDepth)
+    }
+
+    func test_shouldFlushForQueueDepth_atThresholdIsTrue() {
+        let queue = Array(repeating: sampleRequest(), count: StateManagementConstants.flushDepth)
+        var state = KlaviyoState(queue: queue)
+        state.flushInterval = StateManagementConstants.wifiFlushInterval
+        XCTAssertTrue(state.shouldFlushForQueueDepth)
+    }
+
+    func test_shouldFlushForQueueDepth_isFalseWhenOffline() {
+        let queue = Array(repeating: sampleRequest(), count: StateManagementConstants.flushDepth)
+        var state = KlaviyoState(queue: queue)
+        state.flushInterval = .infinity
+        XCTAssertFalse(state.shouldFlushForQueueDepth)
+    }
+
+    // MARK: - flushQueue token gate (reducer)
+
+    @MainActor
+    func test_flushQueue_withDepletedBucketDoesNotFlush() async {
+        var initialState = INITIALIZED_TEST_STATE()
+        initialState.flushing = false
+        initialState.availableFlushTokens = 0
+        // Refilled "now", so no time has elapsed at flush time -> no token accrues.
+        initialState.lastFlushTokenRefill = environment.date()
+        initialState.queue = [initialState.buildProfileRequest(
+            apiKey: initialState.apiKey!,
+            anonymousId: initialState.anonymousId!
+        )]
+        let store = TestStore(initialState: initialState, reducer: KlaviyoReducer())
+
+        // Bucket is empty and nothing refills, so the flush is deferred:
+        // no state change and, crucially, no `.sendRequest` is emitted.
+        _ = await store.send(.flushQueue)
+    }
+
+    // MARK: - queue-depth early-flush trigger (reducer)
+
+    @MainActor
+    func test_enqueueAggregateEventReachingQueueDepthTriggersFlush() async {
+        var initialState = INITIALIZED_TEST_STATE()
+        initialState.flushing = false
+        // Empty the bucket so the triggered flush is gated. This keeps the test focused on the
+        // depth trigger itself rather than the downstream send cascade.
+        initialState.availableFlushTokens = 0
+        initialState.queue = Array(
+            repeating: sampleRequest(name: "filler"),
+            count: StateManagementConstants.flushDepth - 1
+        )
+        let store = TestStore(initialState: initialState, reducer: KlaviyoReducer())
+        store.exhaustivity = .off
+
+        // Enqueuing one more request reaches `flushDepth`, which schedules an early flush.
+        await store.send(.enqueueAggregateEvent(Data("agg".utf8)))
+        await store.receive(.flushQueue)
+    }
+}
diff --git a/Tests/KlaviyoSwiftTests/KlaviyoTestUtils.swift b/Tests/KlaviyoSwiftTests/KlaviyoTestUtils.swift
@@ -92,6 +92,16 @@ extension Store where State == KlaviyoState, Action == KlaviyoAction {
     static let test = Store(initialState: .test, reducer: KlaviyoTestReducer())
 }
 
+extension KlaviyoState {
+    /// Mirrors the token-bucket mutation performed by a single successful `.flushQueue`:
+    /// one token is consumed and the refill timestamp is advanced to "now". Use inside
+    /// `TestStore` expectation closures to keep flush assertions DRY.
+    mutating func expectFlushTokenConsumed() {
+        availableFlushTokens -= 1
+        lastFlushTokenRefill = environment.date()
+    }
+}
+
 extension FileClient {
     static let test = FileClient(
         write: { _, _ in },
diff --git a/Tests/KlaviyoSwiftTests/StateManagementTests.swift b/Tests/KlaviyoSwiftTests/StateManagementTests.swift
@@ -140,6 +140,7 @@ class StateManagementTests: XCTestCase {
             $0.flushing = true
             $0.requestsInFlight = $0.queue
             $0.queue = []
+            $0.expectFlushTokenConsumed()
         }
 
         await store.receive(.sendRequest)
@@ -173,6 +174,7 @@ class StateManagementTests: XCTestCase {
             $0.flushing = true
             $0.requestsInFlight = $0.queue
             $0.queue = []
+            $0.expectFlushTokenConsumed()
         }
 
         await store.receive(.sendRequest)
@@ -206,6 +208,7 @@ class StateManagementTests: XCTestCase {
             $0.flushing = true
             $0.requestsInFlight = $0.queue
             $0.queue = []
+            $0.expectFlushTokenConsumed()
         }
 
         await store.receive(.sendRequest)
@@ -313,6 +316,7 @@ class StateManagementTests: XCTestCase {
             $0.flushing = true
             $0.requestsInFlight = $0.queue
             $0.queue = []
+            $0.expectFlushTokenConsumed()
         }
         await store.receive(.sendRequest)
 
@@ -360,6 +364,7 @@ class StateManagementTests: XCTestCase {
             $0.flushing = true
             $0.requestsInFlight = $0.queue
             $0.queue = []
+            $0.expectFlushTokenConsumed()
         }
         await store.receive(.sendRequest)
 
@@ -480,6 +485,7 @@ class StateManagementTests: XCTestCase {
             $0.requestsInFlight = $0.queue
             $0.queue = []
             $0.flushing = true
+            $0.expectFlushTokenConsumed()
             $0.pendingProfile = nil
             request = $0.requestsInFlight[0]
             switch request?.endpoint {
@@ -761,6 +767,7 @@ class StateManagementTests: XCTestCase {
             $0.flushing = true
             $0.requestsInFlight = $0.queue
             $0.queue = []
+            $0.expectFlushTokenConsumed()
             XCTAssertEqual($0.requestsInFlight.count, 3, "Should have 3 requests in flight")
             actualGeofenceRequest = $0.requestsInFlight[0]
             if case let .createEvent(_, payload) = actualGeofenceRequest!.endpoint {
diff --git a/Tests/KlaviyoSwiftTests/__Snapshots__/KlaviyoStateTests/testLoadNewKlaviyoState.1.txt b/Tests/KlaviyoSwiftTests/__Snapshots__/KlaviyoStateTests/testLoadNewKlaviyoState.1.txt
@@ -3,12 +3,14 @@
     - some: "00000000-0000-0000-0000-000000000001"
   ▿ apiKey: Optional<String>
     - some: "foo"
+  - availableFlushTokens: 5.0
   - email: Optional<String>.none
   - externalId: Optional<String>.none
   - flushInterval: 10.0
   - flushing: false
   - initalizationState: InitializationState.uninitialized
   - isProcessingDeepLink: false
+  - lastFlushTokenRefill: Optional<Date>.none
   - pendingProfile: Optional<Dictionary<ProfileKey, AnyEncodable>>.none
   - pendingRequests: 0 elements
   - phoneNumber: Optional<String>.none
diff --git a/Tests/KlaviyoSwiftTests/__Snapshots__/KlaviyoStateTests/testStateFileExistsInvalidData.1.txt b/Tests/KlaviyoSwiftTests/__Snapshots__/KlaviyoStateTests/testStateFileExistsInvalidData.1.txt
@@ -3,12 +3,14 @@
     - some: "00000000-0000-0000-0000-000000000001"
   ▿ apiKey: Optional<String>
     - some: "foo"
+  - availableFlushTokens: 5.0
   - email: Optional<String>.none
   - externalId: Optional<String>.none
   - flushInterval: 10.0
   - flushing: false
   - initalizationState: InitializationState.uninitialized
   - isProcessingDeepLink: false
+  - lastFlushTokenRefill: Optional<Date>.none
   - pendingProfile: Optional<Dictionary<ProfileKey, AnyEncodable>>.none
   - pendingRequests: 0 elements
   - phoneNumber: Optional<String>.none
diff --git a/Tests/KlaviyoSwiftTests/__Snapshots__/KlaviyoStateTests/testStateFileExistsInvalidJSON.1.txt b/Tests/KlaviyoSwiftTests/__Snapshots__/KlaviyoStateTests/testStateFileExistsInvalidJSON.1.txt
@@ -3,12 +3,14 @@
     - some: "00000000-0000-0000-0000-000000000001"
   ▿ apiKey: Optional<String>
     - some: "foo"
+  - availableFlushTokens: 5.0
   - email: Optional<String>.none
   - externalId: Optional<String>.none
   - flushInterval: 10.0
   - flushing: false
   - initalizationState: InitializationState.uninitialized
   - isProcessingDeepLink: false
+  - lastFlushTokenRefill: Optional<Date>.none
   - pendingProfile: Optional<Dictionary<ProfileKey, AnyEncodable>>.none
   - pendingRequests: 0 elements
   - phoneNumber: Optional<String>.none
diff --git a/Tests/KlaviyoSwiftTests/__Snapshots__/KlaviyoStateTests/testValidStateFileExists.1.txt b/Tests/KlaviyoSwiftTests/__Snapshots__/KlaviyoStateTests/testValidStateFileExists.1.txt
