之前使用了JDK自帶的 Map 實現了自定義限速的簡單需求。在當時的實現當中，有一個被隱藏的小設計，就是如果是用使用異步線程，用來根據配置給請求次數數據重置。如此這樣，校驗方法會非常簡單方便。
對於普通 Java 項目來說，如果使用異步線程處理，除了 deamon 進程以外，其他實現的確有點麻煩。即使 deamon 線程也很難做到完全的實用性，所以才使用了上篇文章的實現方案。
之前提到過一個非常有趣的高性能本地緩存 Caffeine 剛好能解決這個問題，可以通過緩存過期或者定時刷新功能來實現定時刷新的需求。這裏我選擇了定時刷新功能，這種選擇會限制限流配置的種類，無法進行 2/3s , 10/2s 配置，我最終選擇 TPS 進行配置，全部使用 1s 爲限制週期。
代碼
主要思路如下：
- 數據結構選擇：使用了兩種數據結構來實現限流功能：使用了一個 Map 來存儲每個請求的限流配置，以請求的標識符作爲鍵，以該請求的每秒事務數（TPS）作爲值；使用了 Caffeine 緩存來存儲每個請求的計數器，其中鍵爲請求的標識符，值爲一個原子整數，用於記錄請求的處理數量。
- 限流判斷邏輯：具體邏輯如下：首先從緩存中獲取對應請求的計數器；判斷當前計數器的值是否大於等於該請求的配置的每秒事務數（TPS）；如果超過了配置的值，則表示需要限流，返回 true。否則，遞增計數器並返回 false，表示不需要限流。
- 動態配置：動態添加請求的限流配置，將請求的標識符和對應的每秒事務數（TPS）添加到配置中，實現了動態配置的功能。
- 使用 Caffeine 緩存：使用了 Caffeine 緩存來存儲請求的計數器，可以配置緩存的過期時間（1秒），當緩存過期時會自動刷新。這樣可以確保計數器在一定時間內有效，避免長時間未使用的請求佔用內存。
代碼如下：
import com.github.benmanes.caffeine.cache.Caffeine
import com.github.benmanes.caffeine.cache.LoadingCache
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicInteger
/**
* 限流工具,基於Caffeine實現,支持動態配置,根據TPS限流
*/
class TpsLimit {
Map<String, Integer> qpsConfig = [:]
LoadingCache<Object, AtomicInteger> build = Caffeine.newBuilder().refreshAfterWrite(1, TimeUnit.SECONDS).build((key) -> {
return new AtomicInteger()
})
/**
* 是否限流
* @param key
* @return
*/
boolean isLimit(String key) {
AtomicInteger atomicInteger = build.get(key)
if (atomicInteger.get() >= qpsConfig.get(key, 1)) {
return true
}
atomicInteger.incrementAndGet()
return false
}
/**
* 添加限流配置
* @param key
* @param qps
* @return
*/
def addConfig(String key, Integer qps) {
qpsConfig.put(key, qps)
}
}
測試
測試腳本如下，與前一篇文章大同小異：
import com.funtester.httpclient.FunHttp
import com.funtester.utils.TpsLimit
class Routine extends FunHttp {
static void main(String[] args) {
def limit = new TpsLimit()
limit.addConfig("test", 1)
1000.times {
sleep(0.1)
fun {
def limit1 = limit.isLimit("t4est")
if (!limit1) {
output("未限流")
}
}
} }
}
控制檯輸出：
22:19:20:545 F-1 未限流
22:19:20:644 F-2 未限流
22:19:22:094 F-8 未限流
22:19:22:195 F-1 未限流
22:19:24:048 F-3 未限流
22:19:24:150 F-4 未限流
22:19:25 uptime:6 s
22:19:25 finished: 49 task
可以看出，按照默認配置1 TPS的配置實現。