memory-store.js
2.08 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
/**
* 内存 store：实现 `RateLimitStore`，单机 / 单测用。
*
* 用 `Map<key, BucketState>` 持桶状态，所有桶数学复用 `token-bucket.js`——保证与
* 未来 redis 实现行为一致。JS 单线程下「读-补-扣」在一次同步函数调用内完成、其间
* 不交还事件循环，因此天然原子：并发的 take 顺序执行、不会放行超过 capacity。
* redis 实现需用 Lua/单命令复刻这一原子性。
*
* @typedef {import('./rules.js').RateRule} RateRule
* @typedef {import('./token-bucket.js').BucketState} BucketState
*
* RateLimitStore 接口（§3，本包对其唯一要求）：
* {
* take(key, cost, rule, nowMs): Promise<{ allowed, remaining, retryAfterMs }>,
* peek(key, rule, nowMs): Promise<{ remaining, retryAfterMs }>
* }
*/
import * as bucket from './token-bucket.js'
/**
* 创建一个内存 store。
* @returns {{
* take(key: string, cost: number, rule: RateRule, nowMs: number): Promise<{ allowed: boolean, remaining: number, retryAfterMs: number }>,
* peek(key: string, rule: RateRule, nowMs: number): Promise<{ remaining: number, retryAfterMs: number }>
* }}
*/
export function createMemoryStore() {
/** @type {Map<string, BucketState>} */
const buckets = new Map()
/**
* @param {string} key
* @param {RateRule} rule
* @param {number} nowMs
* @returns {BucketState}
*/
function stateOf(key, rule, nowMs) {
return buckets.get(key) ?? bucket.initialState(rule, nowMs)
}
return {
async take(key, cost, rule, nowMs) {
const state = stateOf(key, rule, nowMs)
const r = bucket.take(state, cost, rule, nowMs)
buckets.set(key, r.next)
return { allowed: r.allowed, remaining: r.remaining, retryAfterMs: r.retryAfterMs }
},
async peek(key, rule, nowMs) {
const state = stateOf(key, rule, nowMs)
// 只读：不写回，不消耗令牌。cost 固定为 1（判定「下一次单位调用是否会放行」）。
const r = bucket.peek(state, 1, rule, nowMs)
return { remaining: r.remaining, retryAfterMs: r.retryAfterMs }
},
}
}