Skip to content

Commit e0aa3c6

Browse files
committed
Update BackpressureMiddlerware to use RedisClient
1 parent 04ebe11 commit e0aa3c6

File tree

2 files changed

+60
-48
lines changed

2 files changed

+60
-48
lines changed
Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
// Copyright Dave Verwer, Sven A. Schmidt, and other contributors.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
import Dependencies
16+
import Redis
17+
import Vapor
18+
19+
20+
final class BackpressureMiddleware: AsyncMiddleware {
21+
22+
let slidingWindow: Duration
23+
let countLimit: Int
24+
25+
public func respond(to request: Request, chainingTo next: AsyncResponder) async throws -> Response {
26+
guard let cfray = request.headers.first(name: "cf-ray") else {
27+
return try await next.respond(to: request)
28+
}
29+
let epochSeconds = Int(Date().timeIntervalSince1970)
30+
let combinedKey = "\(cfray):\(epochSeconds)"
31+
32+
let slidingWindow = Int(slidingWindow.components.seconds)
33+
34+
@Dependency(\.redis) var redis
35+
36+
guard var countOverWindow = try? await redis.increment(key: combinedKey) else {
37+
request.logger.log(level: .warning, "BackpressureMiddleware failed to increment key '\(combinedKey)'.")
38+
return try await next.respond(to: request)
39+
}
40+
41+
for i in epochSeconds - slidingWindow ..< epochSeconds {
42+
let key = "\(cfray):\(i)"
43+
if let value = (try? await redis.get(key)).flatMap(Int.init) {
44+
countOverWindow += value
45+
}
46+
}
47+
48+
if countOverWindow >= countLimit {
49+
request.logger.log(level: .warning, "BackpressureMiddleware acting on request with cf-ray '\(cfray)' due to \(countOverWindow) requests in the last \(slidingWindow) seconds.")
50+
throw Abort(.tooManyRequests)
51+
} else {
52+
return try await next.respond(to: request)
53+
}
54+
}
55+
56+
init(slidingWindow: Duration, countLimit: Int) {
57+
self.slidingWindow = slidingWindow
58+
self.countLimit = countLimit
59+
}
60+
}

Sources/App/Core/RateCatcherMiddleware.swift

Lines changed: 0 additions & 48 deletions
This file was deleted.

0 commit comments

Comments
 (0)