feat: Client-side rate limiting (#564)

Author: jbelkins

Sources/ClientRuntime/Retries/DefaultRetryStrategy/ClientSideRateLimiter.swift

@@ -0,0 +1,155 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import struct Foundation.TimeInterval
+import struct Foundation.Date
+import func Foundation.pow
+
+actor ClientSideRateLimiter {
+
+    // these are constants defined in Retry Behavior 2.0
+    let minFillRate: Double = 0.5
+    let minCapacity: Double = 1.0
+    let smooth: Double = 0.8
+    let beta = 0.7
+    let scaleConstant = 0.4
+
+    // these are state variables explicitly declared in Retry Behavior 2.0
+    var fillRate: Double = 0.0
+    var maxCapacity: Double = 0.0
+    var currentCapacity: Double = 0.0
+    var lastTimestamp: TimeInterval? = 0.0
+    var enabled = false
+    var measuredTXRate: Double = 0.0
+    var lastTXRateBucket: Double
+    var requestCount: Int = 0
+    var lastMaxRate: Double = 0.0
+    var lastThrottleTime: TimeInterval
+
+    // not explicitly included as state in Retry Behavior 2.0, but it said to cache the
+    // value when lastMaxRate changes
+    var timeWindow: Double = 0.0
+
+    // Returns the current time when called.
+    // Exposed so time may be mocked for testing.
+    var clock: () -> TimeInterval
+
+    /// Creates a new client-side rate limiter.
+    ///
+    /// Parameters are for use during testing.  To create this type for actual use, call `.init()`.
+    /// - Parameters:
+    ///   - lastMaxRate: The last max rate to set.  For testing use only.
+    ///   - lastThrottleTime: The last throttle time to set.  For testing use only.
+    ///   - clock: An anonymous closure that provides the current time in the form of a timestamp.  Defaults to actual time.  For testing use only.
+    init(
+        lastMaxRate: Double = 0.0,
+        lastThrottleTime: TimeInterval? = nil,
+        clock: @escaping () -> TimeInterval = { Date().timeIntervalSinceReferenceDate }
+    ) {
+        self.lastMaxRate = lastMaxRate
+        self.lastThrottleTime = lastThrottleTime ?? clock()
+        self.lastTXRateBucket = Self.floor(clock())
+        self.clock = clock
+    }
+
+    // The following functions are built exactly as described in Retry Behavior 2.0.
+
+    func tokenBucketAcquire(amount: Double) -> TimeInterval? {
+        if !enabled { return nil }
+        tokenBucketRefill()
+        if amount <= currentCapacity {
+            currentCapacity -= amount
+            return nil
+        } else {
+            let delay = (amount - currentCapacity) / fillRate
+            currentCapacity -= amount
+            return delay
+        }
+    }
+
+    private func tokenBucketRefill() {
+        let timestamp = clock()
+        guard let lastTimestamp = lastTimestamp else {
+            self.lastTimestamp = timestamp
+            return
+        }
+        let fillAmount = (timestamp - lastTimestamp) * fillRate
+        currentCapacity = min(maxCapacity, currentCapacity + fillAmount)
+        self.lastTimestamp = timestamp
+    }
+
+    private func tokenBucketUpdateRate(newRPS: Double) {
+        tokenBucketRefill()
+        fillRate = max(newRPS, minFillRate)
+        maxCapacity = max(newRPS, minCapacity)
+        currentCapacity = min(currentCapacity, maxCapacity)
+    }
+
+    private func tokenBucketEnable() {
+        enabled = true
+    }
+
+    private func updateMeasuredRate() {
+        let t = clock()
+        let timeBucket = Self.floor(t * 2.0) / 2.0
+        requestCount += 1
+        if timeBucket > lastTXRateBucket {
+            let currentRate = Double(requestCount) / (timeBucket - lastTXRateBucket)
+            measuredTXRate = (currentRate * smooth) + (measuredTXRate * (1.0 - smooth))
+            requestCount = 0
+            lastTXRateBucket = timeBucket
+        }
+    }
+
+    // Exposed internally for use while testing.
+    func updateClientSendingRate(isThrottling: Bool) {
+        updateMeasuredRate()
+        let calculatedRate: Double
+        if isThrottling {
+            let rateToUse = enabled ? min(measuredTXRate, fillRate) : measuredTXRate
+            lastMaxRate = rateToUse
+            calculateTimeWindow()
+            lastThrottleTime = clock()
+            calculatedRate = cubicThrottle(rateToUse: rateToUse)
+            tokenBucketEnable()
+        } else {
+            calculateTimeWindow()
+            calculatedRate = cubicSuccess(timestamp: clock())
+        }
+        let newRate = min(calculatedRate, 2.0 * measuredTXRate)
+        tokenBucketUpdateRate(newRPS: newRate)
+    }
+
+    // Exposed internally for use while testing.
+    func calculateTimeWindow() {
+        timeWindow = pow(lastMaxRate * (1.0 - beta) / scaleConstant, 1.0 / 3.0)
+    }
+
+    // Exposed internally for use while testing.
+    func cubicSuccess(timestamp: TimeInterval) -> Double {
+        let dt = timestamp - lastThrottleTime
+        return scaleConstant * pow(dt - timeWindow, 3.0) + lastMaxRate
+    }
+
+    // Exposed internally for use while testing.
+    func cubicThrottle(rateToUse: Double) -> Double {
+        return rateToUse * beta
+    }
+
+    private static func floor(_ time: TimeInterval) -> TimeInterval {
+        time.rounded(.down)
+    }
+
+    // The following functions are not described in Retry Behavior 2.0 but are
+    // used to set test conditions.
+
+    func setLastMaxRate(_ newValue: Double) { lastMaxRate = newValue }
+
+    func setLastThrottleTime(_ newValue: Double) { lastThrottleTime = newValue }
+
+    func setClock(_ newClock: @escaping () -> TimeInterval) { clock = newClock }
+}

Sources/ClientRuntime/Retries/DefaultRetryStrategy/DefaultRetryStrategy.swift

@@ -5,7 +5,7 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-import Foundation
+import struct Foundation.TimeInterval
 
 public struct DefaultRetryStrategy: RetryStrategy {
     public typealias Token = DefaultRetryToken
@@ -17,27 +17,26 @@ public struct DefaultRetryStrategy: RetryStrategy {
     /// Used to inject a mock during unit tests that simulates sleeping.
     /// The default `sleeper` function actually sleeps asynchronously.
     var sleeper: (TimeInterval) async throws -> Void = { delay in
+        guard delay > 0.0 else { return }
         try await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000.0))
     }
 
     public init(options: RetryStrategyOptions) {
         self.options = options
-        self.quotaRepository = RetryQuotaRepository(
-            availableCapacity: options.availableCapacity,
-            maxCapacity: options.maxCapacity
-        )
+        self.quotaRepository = RetryQuotaRepository(options: options)
     }
 
     public func acquireInitialRetryToken(tokenScope: String) async throws -> DefaultRetryToken {
         let quota = await quotaRepository.quota(partitionID: tokenScope)
+        let rateLimitDelay = await quota.getRateLimitDelay()
+        try await sleeper(rateLimitDelay)
         return DefaultRetryToken(quota: quota)
     }
 
     public func refreshRetryTokenForRetry(tokenToRenew: DefaultRetryToken, errorInfo: RetryErrorInfo) async throws {
-        let delay = errorInfo.retryAfterHint ??
+        let backoffDelay = errorInfo.retryAfterHint ??
             options.backoffStrategy.computeNextBackoffDelay(attempt: tokenToRenew.retryCount)
         tokenToRenew.retryCount += 1
-        tokenToRenew.delay = delay
         if tokenToRenew.retryCount > options.maxRetriesBase {
             throw RetryError.maxAttemptsReached
         }
@@ -46,7 +45,10 @@ public struct DefaultRetryStrategy: RetryStrategy {
         } else {
             throw RetryError.insufficientQuota
         }
-        try await sleeper(tokenToRenew.delay ?? 0.0)
+        let isThrottling = errorInfo.errorType == .throttling
+        await tokenToRenew.quota.updateClientSendingRate(isThrottling: isThrottling)
+        let rateLimitDelay = await tokenToRenew.quota.getRateLimitDelay()
+        try await sleeper(backoffDelay + rateLimitDelay)
     }
 
     public func recordSuccess(token: DefaultRetryToken) async {

Sources/ClientRuntime/Retries/DefaultRetryStrategy/DefaultRetryToken.swift

@@ -11,15 +11,12 @@ import struct Foundation.TimeInterval
 ///
 /// The retry token contains all the state relevant to one request that is needed to manage retry
 /// until the request succeeds or fails after zero or more retries.
-public class DefaultRetryToken: RetryToken {
+public final class DefaultRetryToken: RetryToken {
 
     /// The number of retry attempts that have been made using this token.
     /// Defaults to zero at the initial attempt, goes up by one for each subsequent attempt.
     public internal(set) var retryCount: Int = 0
 
-    /// The delay, in seconds, to the next retry.
-    public internal(set) var delay: TimeInterval?
-
     /// The amount of quota capacity amount held by this token, if any.
     ///
     /// Tokens have nil capacity amount when created.  Quota value is set to a prescribed value when attempting a retry.
@@ -28,6 +25,11 @@ public class DefaultRetryToken: RetryToken {
     /// The quota for this token.  More than one token (i.e. for requests against the same endpoint) may share a quota.
     let quota: RetryQuota
 
+    /// Creates a new retry token.
+    ///
+    /// The quota for this token may be shared with other tokens if other requests are made against an endpoint with
+    /// the same partition ID.
+    /// - Parameter quota: The retry quota associated with this token's request.
     init(quota: RetryQuota) {
         self.quota = quota
     }

Sources/ClientRuntime/Retries/DefaultRetryStrategy/RetryQuota.swift

@@ -5,6 +5,8 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
+import struct Foundation.TimeInterval
+
 /// Keeps the retry quota count for one partition ID.
 ///
 /// Is shared across all requests with the same partition ID; typically this also correlates to one network connection.
@@ -29,14 +31,36 @@ final actor RetryQuota {
     /// The number of tokens this quota currently holds.
     var availableCapacity: Int
 
+    /// The rate limiter to be used, if any.
+    private var rateLimiter: ClientSideRateLimiter?
+
     /// Sets the current capacity in this quota.  To be used for testing only.
     func setAvailableCapacity(_ availableCapacity: Int) { self.availableCapacity = availableCapacity }
 
     /// Creates a new quota, optionally with reduced available capacity (used for testing.)
     /// `maxCapacity` cannot be set less than available.
-    init(availableCapacity: Int, maxCapacity: Int) {
+    /// - Parameters:
+    ///   - availableCapacity: The number of tokens in this quota at creation.
+    ///   - maxCapacity: <#maxCapacity description#>
+    ///   - rateLimitingMode: <#rateLimitingMode description#>
+    init(
+        availableCapacity: Int,
+        maxCapacity: Int,
+        rateLimitingMode: RetryStrategyOptions.RateLimitingMode = .standard
+    ) {
         self.availableCapacity = availableCapacity
         self.maxCapacity = max(maxCapacity, availableCapacity)
+        self.rateLimiter = rateLimitingMode == .adaptive ? ClientSideRateLimiter() : nil
+    }
+
+    /// Creates a new quota with settings from the passed options.
+    /// - Parameter options: The retry strategy options from which to configure this retry quota
+    convenience init(options: RetryStrategyOptions) {
+        self.init(
+            availableCapacity: options.availableCapacity,
+            maxCapacity: options.maxCapacity,
+            rateLimitingMode: options.rateLimitingMode
+        )
     }
 
     /// Deducts the proper number of tokens from available & returns them.
@@ -59,4 +83,12 @@ final actor RetryQuota {
         availableCapacity += capacityAmount ?? Self.noRetryIncrement
         availableCapacity = min(availableCapacity, maxCapacity)
     }
+
+    func getRateLimitDelay() async -> TimeInterval {
+        await rateLimiter?.tokenBucketAcquire(amount: 1.0) ?? 0.0
+    }
+
+    func updateClientSendingRate(isThrottling: Bool) async {
+        await rateLimiter?.updateClientSendingRate(isThrottling: isThrottling)
+    }
 }

Sources/ClientRuntime/Retries/DefaultRetryStrategy/RetryQuotaRepository.swift

@@ -7,13 +7,11 @@
 
 /// Holds multiple quotas, keyed by partition IDs.
 actor RetryQuotaRepository {
-    let maxCapacity: Int
-    let availableCapacity: Int
+    let options: RetryStrategyOptions
     private var quotas = [String: RetryQuota]()
 
-    init(availableCapacity: Int, maxCapacity: Int) {
-        self.availableCapacity = availableCapacity
-        self.maxCapacity = maxCapacity
+    init(options: RetryStrategyOptions) {
+        self.options = options
     }
 
     /// Returns the quota for the given partition ID.
@@ -26,7 +24,7 @@ actor RetryQuotaRepository {
         if let quota = quotas[partitionID] {
             return quota
         } else {
-            let newQuota = RetryQuota(availableCapacity: availableCapacity, maxCapacity: maxCapacity)
+            let newQuota = RetryQuota(options: options)
             quotas[partitionID] = newQuota
             return newQuota
         }

Sources/ClientRuntime/Retries/RetryStrategyOptions.swift

@@ -13,6 +13,26 @@ public struct RetryStrategyOptions {
     /// This is more of a hint since a custom retry strategy could be aware of certain operational contexts ("partition fail over")
     public let maxRetriesBase: Int
 
+    /// Sets the mode used for rate limiting requests in response to throttling.
+    public enum RateLimitingMode {
+
+        /// Requests may be sent immediately, and are not delayed for rate limiting when throttling is detected.
+        ///
+        /// This is default retry behavior.
+        case standard
+
+        /// Initial and retry requests may be delayed by an additional amount when throttling is detected.
+        ///
+        /// This is sometimes called "adaptive" or "client-side rate limiting" mode, and is available opt-in.
+        case adaptive
+    }
+
+    /// The mode to be used for rate-limiting requests.
+    ///
+    /// In `standard` mode, requests are only delayed according to the backoff strategy in use.  In `adaptive` mode, requests are
+    /// delayed when the server indicates that requests are being throttled.
+    public let rateLimitingMode: RateLimitingMode
+
     /// Sets the initial available capacity for this retry strategy's quotas.
     ///
     /// Used only during testing, production uses the default values.
@@ -33,11 +53,13 @@ public struct RetryStrategyOptions {
         backoffStrategy: RetryBackoffStrategy = ExponentialBackoffStrategy(),
         maxRetriesBase: Int = 2,
         availableCapacity: Int = 500,
-        maxCapacity: Int = 500
+        maxCapacity: Int = 500,
+        rateLimitingMode: RateLimitingMode = .standard
     ) {
         self.backoffStrategy = backoffStrategy
         self.maxRetriesBase = maxRetriesBase
         self.availableCapacity = availableCapacity
         self.maxCapacity = maxCapacity
+        self.rateLimitingMode = rateLimitingMode
     }
 }

Sources/ClientRuntime/Retries/RetryToken.swift

@@ -14,7 +14,4 @@ public protocol RetryToken: AnyObject {
 
     /// The number of retries (i.e. NOT including the initial attempt) that this token has made.
     var retryCount: Int { get }
-
-    /// The delay for this request (TODO: not used, maybe get rid of this?)
-    var delay: TimeInterval? { get }
 }

Tests/ClientRuntimeTests/Retry/DefaultRetryErrorInfoProviderTests.swift

@@ -9,7 +9,7 @@ import Foundation
 import XCTest
 @testable import ClientRuntime
 
-class DefaultRetryErrorInfoProviderTests: XCTestCase {
+final class DefaultRetryErrorInfoProviderTests: XCTestCase {
 
     // MARK: - Modeled errors