Make backpressure configurable via BACKPRESSURE_CONFIG

Author: finestructure

Sources/App/Core/BackpressureMiddleware.swift

@@ -18,9 +18,7 @@ import Vapor
 
 
 final class BackpressureMiddleware: AsyncMiddleware {
-
-    let slidingWindow: Duration
-    let countLimit: Int
+    let configuration: Configuration
 
     public func respond(to request: Request, chainingTo next: AsyncResponder) async throws -> Response {
         guard let cfray = request.headers.first(name: "cf-ray") else {
@@ -29,7 +27,7 @@ final class BackpressureMiddleware: AsyncMiddleware {
         let epochSeconds = Int(Date().timeIntervalSince1970)
         let combinedKey = "\(cfray):\(epochSeconds)"
 
-        let slidingWindow = Int(slidingWindow.components.seconds)
+        let slidingWindow = configuration.slidingWindow
 
         @Dependency(\.redis) var redis
 
@@ -48,16 +46,41 @@ final class BackpressureMiddleware: AsyncMiddleware {
             }
         }
 
-        if countOverWindow >= countLimit {
+
+        if countOverWindow >= configuration.requestLimit {
             request.logger.log(level: .warning, "BackpressureMiddleware acting on request with cf-ray '\(cfray)' due to \(countOverWindow) requests in the last \(slidingWindow) seconds.")
-            throw Abort(.tooManyRequests)
+
+            switch configuration.mode {
+                case .block:
+                    throw Abort(.tooManyRequests)
+                case .delay(let duration):
+                    try await Task.sleep(for: .init(seconds: duration))
+                    return try await next.respond(to: request)
+            }
         } else {
             return try await next.respond(to: request)
         }
     }
 
-    init(slidingWindow: Duration, countLimit: Int) {
-        self.slidingWindow = slidingWindow
-        self.countLimit = countLimit
+    init(configuration: Configuration) {
+        self.configuration = configuration
+    }
+
+    struct Configuration: Codable {
+        var mode: Mode
+        var requestLimit: Int
+        var slidingWindow: Int  // in seconds
+
+        enum Mode: Codable {
+            case block
+            case delay(seconds: Double)
+        }
+    }
+}
+
+
+extension Duration {
+    init(seconds: Double) {
+        self = .microseconds(seconds/1e-6)
     }
 }

Sources/App/Core/Dependencies/EnvironmentClient.swift

@@ -31,6 +31,7 @@ struct EnvironmentClient {
     var awsDocsBucket: @Sendable () -> String?
     var awsReadmeBucket: @Sendable () -> String?
     var awsSecretAccessKey: @Sendable () -> String?
+    var backpressureConfiguration: @Sendable () -> BackpressureMiddleware.Configuration?
     var builderToken: @Sendable () -> String?
     var buildTimeout: @Sendable () -> Int = { XCTFail("buildTimeout"); return 10 }
     var buildTriggerAllowList: @Sendable () -> [Package.Id] = { XCTFail("buildTriggerAllowList"); return [] }
@@ -72,6 +73,10 @@ extension EnvironmentClient: DependencyKey {
             awsDocsBucket: { Environment.get("AWS_DOCS_BUCKET") },
             awsReadmeBucket: { Environment.get("AWS_README_BUCKET") },
             awsSecretAccessKey: { Environment.get("AWS_SECRET_ACCESS_KEY") },
+            backpressureConfiguration: {
+                Environment.decode("BACKPRESSURE_CONFIG",
+                                   as: BackpressureMiddleware.Configuration.self)
+            },
             builderToken: { Environment.get("BUILDER_TOKEN") },
             buildTimeout: { Environment.get("BUILD_TIMEOUT").flatMap(Int.init) ?? 10 },
             buildTriggerAllowList: {
@@ -137,6 +142,7 @@ extension EnvironmentClient: TestDependencyKey {
         // mechanism for a few heavily used dependencies at the moment.
         var mock = Self()
         mock.appVersion = { "test" }
+        mock.backpressureConfiguration = { nil }
         mock.current = { .development }
         return mock
     }

Sources/App/configure.swift

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import Dependencies
 import Fluent
 import FluentPostgresDriver
 import Vapor
@@ -37,6 +38,11 @@ public func configure(_ app: Application) async throws -> String {
     // app.http.server.configuration.responseCompression = .enabled
     // app.http.server.configuration.requestDecompression = .enabled
 
+    @Dependency(\.environment) var environment
+    if let config = environment.backpressureConfiguration() {
+        print(config)
+        app.middleware.use(BackpressureMiddleware(configuration: config))
+    }
     app.middleware.use(FileMiddleware(publicDirectory: app.directory.publicDirectory))
     app.middleware.use(ErrorMiddleware())
 

app.yml

@@ -28,6 +28,7 @@ x-shared: &shared
     AWS_DOCS_BUCKET: ${AWS_DOCS_BUCKET}
     AWS_README_BUCKET: ${AWS_README_BUCKET}
     AWS_SECRET_ACCESS_KEY: ${AWS_SECRET_ACCESS_KEY}
+    BACKPRESSURE_CONFIG: ${BACKPRESSURE_CONFIG}
     BUILD_TIMEOUT: ${BUILD_TIMEOUT}
     BUILD_TRIGGER_ALLOW_LIST: ${BUILD_TRIGGER_ALLOW_LIST}
     BUILD_TRIGGER_DOWNSCALING: ${BUILD_TRIGGER_DOWNSCALING}