diff --git a/bfe_modules/mod_ai_token_auth/mod_ai_token_auth.go b/bfe_modules/mod_ai_token_auth/mod_ai_token_auth.go index b99c62361..9602b39c5 100644 --- a/bfe_modules/mod_ai_token_auth/mod_ai_token_auth.go +++ b/bfe_modules/mod_ai_token_auth/mod_ai_token_auth.go @@ -211,6 +211,11 @@ func GetApiKey(req *bfe_basic.Request) string { // found product handler func (m *ModuleAITokenAuth) tokenFoundProductHandler(req *bfe_basic.Request) (int, *bfe_http.Response) { + meta := req.GetAiBasicInfo() + if meta == nil { + return bfe_module.BfeHandlerGoOn, nil + } + matched := m.matchTokenRule(req) if !matched { // no rule, just pass diff --git a/bfe_util/limit_rate/redis_lr_agent_test.go b/bfe_util/limit_rate/redis_lr_agent_test.go index baf015807..712a95f91 100644 --- a/bfe_util/limit_rate/redis_lr_agent_test.go +++ b/bfe_util/limit_rate/redis_lr_agent_test.go @@ -19,15 +19,24 @@ import ( "testing" "time" + "github.com/bfenetworks/bfe/bfe_util/bns" "github.com/bfenetworks/bfe/bfe_util/redis_client" ) func makeRedisClient() redis_client.Client { + filename := "./testdata/name_conf.data" + err := bns.LoadLocalNameConf(filename) + if err != nil { + fmt.Println("LoadLocalNameConf err", err.Error()) + } + opts := &redis_client.Options{ - ServiceConf: "127.0.0.1:6379", - MaxIdle: 10, - MaxActive: 20, - Wait: true, + //ServiceConf: "127.0.0.1:6379", + ServiceConf: "testredis", + + MaxIdle: 10, + MaxActive: 20, + Wait: true, ConnTimeoutMs: 100, ReadTimeoutMs: 100, @@ -35,6 +44,11 @@ func makeRedisClient() redis_client.Client { Password: "", } + err = opts.Format() + if err != nil { + fmt.Println("Format Options err", err.Error()) + } + client := redis_client.NewRedisClient(opts) return client } diff --git a/bfe_util/limit_rate/redis_qpm_agent.go b/bfe_util/limit_rate/redis_qpm_agent.go index c5cc11902..1378d93cf 100644 --- a/bfe_util/limit_rate/redis_qpm_agent.go +++ b/bfe_util/limit_rate/redis_qpm_agent.go @@ -19,9 +19,11 @@ import ( "strconv" ) +// QpmCheck performs sliding-window rate limiting. +// Main window: key stores timestamps in a ZSET, period seconds, max limit requests. +// tat returns the current main-window request count (was TAT in GCRA). func (agent *RedisLRAgent) QpmCheck( key string, - burst int64, period int64, limit int64, expectUse int64, @@ -36,9 +38,7 @@ func (agent *RedisLRAgent) QpmCheck( return } - // now := float64(time.Now().UnixNano()) / 1e9 - // args := []interface{}{burst, period, limit, expectUse, now} - args := []interface{}{burst, period, limit, expectUse} + args := []interface{}{period, limit, expectUse} res, errtmp := agent.scriptQpmCheck.Run(key, args...) if errtmp != nil { diff --git a/bfe_util/limit_rate/redis_qpm_agent_test.go b/bfe_util/limit_rate/redis_qpm_agent_test.go index 23a76d6d4..94dce4b6c 100644 --- a/bfe_util/limit_rate/redis_qpm_agent_test.go +++ b/bfe_util/limit_rate/redis_qpm_agent_test.go @@ -20,170 +20,298 @@ import ( "time" ) -func TestQpmRedisAgent(t *testing.T) { - // //disable when git push - // tmpQpmRedisAgentImpl(t) - // tmpQpmRedisAgentBurstImpl(t) - // tmpQpmRedisAgentBurstImpl2(t) +func TestQpmSlidingWindow(t *testing.T) { + // Integration tests requiring Redis at 127.0.0.1:6379 + // t.Run("WithinLimit", testQpmWithinLimit) + // t.Run("ExceedLimit", testQpmExceedLimit) + // t.Run("WindowExpire", testQpmWindowExpire) + // t.Run("SlidingBehavior", testQpmSlidingBehavior) + // t.Run("CountCorrectness", testQpmCountCorrectness) + // t.Run("KeyIsolation", testQpmKeyIsolation) + // t.Run("LimitOne", testQpmLimitOne) + // t.Run("EmissionCost", testQpmEmissionCost) } -func tmpQpmRedisAgentImpl(t *testing.T) { +func testQpmWithinLimit(t *testing.T) { client := makeRedisClient() agent := NewRedisLRAgent(client) - key := "testcase_qpm" - burst := int64(1) + key := "testcase_qpm_within" period := int64(60) - limit := int64(60) //1 request per second - expectUse := int64(1) - - // first request should be allowed - currentTime, isAllowed, tat, err := agent.QpmCheck( - key, - burst, - period, - limit, - expectUse, - ) + limit := int64(10) + + for i := int64(0); i < limit; i++ { + _, isAllowed, _, err := agent.QpmCheck(key, period, limit, 1) + if err != nil { + t.Fatalf("QpmCheck error at idx %d: %v", i, err) + } + if !isAllowed { + t.Errorf("request %d should be allowed", i) + } + } + fmt.Println("WithinLimit: all 10 requests allowed") +} + +func testQpmExceedLimit(t *testing.T) { + client := makeRedisClient() + agent := NewRedisLRAgent(client) + + key := "testcase_qpm_exceed" + period := int64(60) + limit := int64(10) + + for i := int64(0); i < limit; i++ { + _, isAllowed, _, err := agent.QpmCheck(key, period, limit, 1) + if err != nil { + t.Fatalf("QpmCheck error at idx %d: %v", i, err) + } + if !isAllowed { + t.Fatalf("request %d should be allowed, but was rejected", i) + } + } + + _, isAllowed, count, err := agent.QpmCheck(key, period, limit, 1) if err != nil { - t.Fatalf("QpmCheck returned error: %v", err) + t.Fatalf("QpmCheck error: %v", err) } - if currentTime <= 0 { - t.Errorf("unexpected currentTime: %v", currentTime) + if isAllowed { + t.Errorf("request %d should be rejected, count=%f", limit+1, count) } - if !isAllowed { - t.Errorf("isAllowed should be allowed on first request") - } - fmt.Printf("QpmCheck1 currentTime:%v, isAllowed:%v, tat:%f, err:%v\n", currentTime, isAllowed, tat, err) - - // second request immediately should be denied (increment=1s, burst=1, only 1 buffered) - currentTime, isAllowed, tat, err = agent.QpmCheck( - key, - burst, - period, - limit, - expectUse, - ) + if int64(count) != limit { + t.Errorf("count should be %d, got %f", limit, count) + } + fmt.Printf("ExceedLimit: 11th request rejected, count=%f\n", count) +} + +func testQpmWindowExpire(t *testing.T) { + client := makeRedisClient() + agent := NewRedisLRAgent(client) + + key := "testcase_qpm_expire" + period := int64(2) + limit := int64(5) + + for i := int64(0); i < limit; i++ { + _, isAllowed, _, err := agent.QpmCheck(key, period, limit, 1) + if err != nil { + t.Fatalf("QpmCheck error at idx %d: %v", i, err) + } + if !isAllowed { + t.Fatalf("request %d should be allowed", i) + } + } + + _, isAllowed, _, err := agent.QpmCheck(key, period, limit, 1) if err != nil { - t.Fatalf("QpmCheck returned error: %v", err) + t.Fatalf("QpmCheck error: %v", err) } if isAllowed { - t.Errorf("isAllowed should not be allowed on immediate second request") - } - fmt.Printf("QpmCheck2 currentTime:%v, isAllowed:%v, tat:%f, err:%v\n", currentTime, isAllowed, tat, err) - - // wait for increment (1s) + a small buffer, then request should be allowed again - time.Sleep(1500 * time.Millisecond) - currentTime, isAllowed, tat, err = agent.QpmCheck( - key, - burst, - period, - limit, - expectUse, - ) + t.Errorf("request should be rejected when window full") + } + fmt.Println("WindowExpire: window full, request rejected") + + time.Sleep(3 * time.Second) + + _, isAllowed, _, err = agent.QpmCheck(key, period, limit, 1) if err != nil { - t.Fatalf("QpmCheck returned error: %v", err) + t.Fatalf("QpmCheck error after window expire: %v", err) } if !isAllowed { - t.Errorf("isAllowed should be allowed after waiting 1.5s") + t.Errorf("request should be allowed after window expired") } - fmt.Printf("QpmCheck3 currentTime:%v, isAllowed:%v, tat:%v, err:%v\n", currentTime, isAllowed, tat, err) + fmt.Println("WindowExpire: after 3s, request allowed") } -func tmpQpmRedisAgentBurstImpl(t *testing.T) { +func testQpmCountCorrectness(t *testing.T) { client := makeRedisClient() agent := NewRedisLRAgent(client) - key := "testcase_qpm_burst" - burst := int64(5) + key := "testcase_qpm_count" period := int64(60) - limit := int64(6000) - expectUse := int64(1) - - // with burst=5, increment=0.01s, first 6 requests should be allowed - for idx := 0; idx < 6; idx++ { - currentTime, isAllowed, tat, err := agent.QpmCheck( - key, - burst, - period, - limit, - expectUse, - ) - fmt.Printf("QpmCheckBurst idx:%d, currentTime:%v, isAllowed:%v, tat:%f\n", idx, currentTime, isAllowed, tat) + limit := int64(10) + + for i := int64(1); i <= 5; i++ { + _, _, count, err := agent.QpmCheck(key, period, limit, 1) if err != nil { - t.Fatalf("QpmCheck returned error: %v, idx:%d", err, idx) + t.Fatalf("QpmCheck error: %v", err) } - if currentTime <= 0 { - t.Errorf("unexpected currentTime: %v, idx:%d, %f", currentTime, idx, tat) + if int64(count) != i { + t.Errorf("count should be %d, got %f", i, count) + } + } + fmt.Println("CountCorrectness: count increments correctly from 1 to 5") +} + +func testQpmKeyIsolation(t *testing.T) { + client := makeRedisClient() + agent := NewRedisLRAgent(client) + + key1 := "testcase_qpm_iso_1" + key2 := "testcase_qpm_iso_2" + period := int64(60) + limit := int64(5) + + for i := int64(0); i < limit; i++ { + _, isAllowed, _, err := agent.QpmCheck(key1, period, limit, 1) + if err != nil { + t.Fatalf("key1 QpmCheck error at idx %d: %v", i, err) } if !isAllowed { - t.Errorf("isAllowed should be allowed on request %d", idx) + t.Fatalf("key1 request %d should be allowed", i) } - time.Sleep(2 * time.Millisecond) } - // 7th request should be denied - currentTime, isAllowed, tat, err := agent.QpmCheck( - key, - burst, - period, - limit, - expectUse, - ) + _, isAllowed, _, err := agent.QpmCheck(key1, period, limit, 1) if err != nil { - t.Fatalf("QpmCheck returned error: %v", err) + t.Fatalf("key1 QpmCheck error: %v", err) } if isAllowed { - t.Errorf("isAllowed should not be allowed on 7th request (burst exhausted)") + t.Errorf("key1 should be rejected") } - fmt.Printf("QpmCheckBurst idx:6, currentTime:%v, isAllowed:%v, tat:%f\n", currentTime, isAllowed, tat) + + _, isAllowed, _, err = agent.QpmCheck(key2, period, limit, 1) + if err != nil { + t.Fatalf("key2 QpmCheck error: %v", err) + } + if !isAllowed { + t.Errorf("key2 should be allowed (isolated from key1)") + } + fmt.Println("KeyIsolation: key2 not affected by key1 exhaustion") } -func tmpQpmRedisAgentBurstImpl2(t *testing.T) { +func testQpmLimitOne(t *testing.T) { client := makeRedisClient() agent := NewRedisLRAgent(client) - key := "testcase_qpm_burst1" - burst := int64(1) + key := "testcase_qpm_limit1" period := int64(60) - limit := int64(6000) //0.01s each request - expectUse := int64(1) - - // with burst=5, increment=0.01s, first 6 requests should be allowed - for idx := 0; idx < 6; idx++ { - currentTime, isAllowed, tat, err := agent.QpmCheck( - key, - burst, - period, - limit, - expectUse, - ) - fmt.Printf("QpmCheckBurst idx:%d, currentTime:%v, isAllowed:%v, tat:%f\n", idx, currentTime, isAllowed, tat) + limit := int64(1) + + _, isAllowed, _, err := agent.QpmCheck(key, period, limit, 1) + if err != nil { + t.Fatalf("QpmCheck error: %v", err) + } + if !isAllowed { + t.Errorf("first request should be allowed") + } + + _, isAllowed, _, err = agent.QpmCheck(key, period, limit, 1) + if err != nil { + t.Fatalf("QpmCheck error: %v", err) + } + if isAllowed { + t.Errorf("second request should be rejected (limit=1)") + } + fmt.Println("LimitOne: 1st allowed, 2nd rejected") +} + +func testQpmEmissionCost(t *testing.T) { + client := makeRedisClient() + agent := NewRedisLRAgent(client) + + key := "testcase_qpm_cost" + period := int64(60) + limit := int64(10) + + _, isAllowed, count, err := agent.QpmCheck(key, period, limit, 3) + if err != nil { + t.Fatalf("QpmCheck error: %v", err) + } + if !isAllowed { + t.Errorf("request with emission_cost=3 should be allowed") + } + if int64(count) != 3 { + t.Errorf("count should be 3, got %f", count) + } + + _, isAllowed, count, err = agent.QpmCheck(key, period, limit, 1) + if err != nil { + t.Fatalf("QpmCheck error: %v", err) + } + if !isAllowed { + t.Errorf("request should be allowed") + } + if int64(count) != 2 { + t.Errorf("count should be 2, got %f", count) + } + fmt.Println("EmissionCost: count reflects emission_cost correctly") +} + +func testQpmSlidingBehavior(t *testing.T) { + client := makeRedisClient() + agent := NewRedisLRAgent(client) + + key := "testcase_qpm_sliding" + period := int64(10) + limit := int64(5) + + // t≈1s: send 3 requests, all allowed (count 0→3) + time.Sleep(1 * time.Second) + for i := 0; i < 3; i++ { + _, isAllowed, _, err := agent.QpmCheck(key, period, limit, 1) if err != nil { - t.Fatalf("QpmCheck returned error: %v, idx:%d", err, idx) + t.Fatalf("t=1s idx %d: %v", i, err) + } + if !isAllowed { + t.Errorf("t=1s idx %d should be allowed", i) } - if currentTime <= 0 { - t.Errorf("unexpected currentTime: %v, idx:%d, %f", currentTime, idx, tat) + } + fmt.Println("t≈1s: 3 requests allowed") + + // t≈2s: send 3 requests, first 2 allowed, 3rd rejected (count 3→5) + time.Sleep(1 * time.Second) + for i := 0; i < 2; i++ { + _, isAllowed, _, err := agent.QpmCheck(key, period, limit, 1) + if err != nil { + t.Fatalf("t=2s idx %d: %v", i, err) } if !isAllowed { - t.Errorf("isAllowed should be allowed on request %d", idx) + t.Errorf("t=2s idx %d should be allowed", i) } - time.Sleep(100 * time.Millisecond) } + _, isAllowed, _, err := agent.QpmCheck(key, period, limit, 1) + if err != nil { + t.Fatalf("t=2s idx 3: %v", err) + } + if isAllowed { + t.Errorf("t=2s idx 3 should be rejected (window full)") + } + fmt.Println("t≈2s: 2 allowed, 1 rejected (window full at 5)") - // 7th request should be allowed - currentTime, isAllowed, tat, err := agent.QpmCheck( - key, - burst, - period, - limit, - expectUse, - ) + // t≈6s: send 4 requests, all rejected (window still full) + time.Sleep(4 * time.Second) + for i := 0; i < 4; i++ { + _, isAllowed, _, err := agent.QpmCheck(key, period, limit, 1) + if err != nil { + t.Fatalf("t=6s idx %d: %v", i, err) + } + if isAllowed { + t.Errorf("t=6s idx %d should be rejected", i) + } + } + fmt.Println("t≈6s: all 4 rejected") + + // t≈7s: send 1 request, rejected + time.Sleep(1 * time.Second) + _, isAllowed, _, err = agent.QpmCheck(key, period, limit, 1) if err != nil { - t.Fatalf("QpmCheck returned error: %v", err) + t.Fatalf("t=7s: %v", err) } - if !isAllowed { - t.Errorf("isAllowed should not be allowed on 7th request (burst exhausted)") + if isAllowed { + t.Errorf("t=7s should be rejected") + } + fmt.Println("t≈7s: rejected") + + // t≈12s: send 3 requests, all allowed (t≈1s entries expired, window slid) + time.Sleep(5 * time.Second) + for i := 0; i < 3; i++ { + _, isAllowed, _, err := agent.QpmCheck(key, period, limit, 1) + if err != nil { + t.Fatalf("t=12s idx %d: %v", i, err) + } + if !isAllowed { + t.Errorf("t=12s idx %d should be allowed", i) + } } - fmt.Printf("QpmCheckBurst idx:6, currentTime:%v, isAllowed:%v, tat:%f\n", currentTime, isAllowed, tat) + fmt.Println("t≈12s: 3 requests allowed (window slid)") } diff --git a/bfe_util/limit_rate/redis_qpm_limit_check.lua b/bfe_util/limit_rate/redis_qpm_limit_check.lua index 972481325..afb9705b7 100644 --- a/bfe_util/limit_rate/redis_qpm_limit_check.lua +++ b/bfe_util/limit_rate/redis_qpm_limit_check.lua @@ -14,58 +14,31 @@ limitations under the License. ]] --- KEYS[1]: unique rate limiting key (e.g. "rate:qpm:tenant_A:model_v3") --- ARGV[1]: max allowed burst buffer size (min 1) --- ARGV[2]: emission period (seconds), fixed at 60 for QPM --- ARGV[3]: allowed requests per period (standard QPM limit, e.g. 1000) --- ARGV[4]: amount consumed by current request (usually 1) --- ARGV[5]: current high-precision system timestamp (seconds, float, e.g. Go's time.Now().UnixNano() / 1e9) +-- Sliding Window QPM Rate Limiter +-- KEYS[1]: rate limit key +-- ARGV[1]: period (main window duration, seconds) +-- ARGV[2]: limit (max requests in main window) +-- ARGV[3]: emission_cost (requests consumed, usually 1) local key = KEYS[1] -local burst = tonumber(ARGV[1]) -local period = tonumber(ARGV[2]) -local limit = tonumber(ARGV[3]) -local emission_cost = tonumber(ARGV[4]) +local period = tonumber(ARGV[1]) +local limit = tonumber(ARGV[2]) +local emission_cost = tonumber(ARGV[3]) ---local now = tonumber(ARGV[5]) - -local current_time = 0 -local t = redis.call('TIME') ---local ms = math.floor((t[1] * 1000) + (t[2] / 1000)) -current_time = t[1] - --- Redis cluster high-precision time +local t = redis.call('TIME') +local current_time = t[1] local now = t[1] + t[2] / 1000000 +redis.call('ZREMRANGEBYSCORE', key, 0, now - period) --- Calculate the theoretical emission increment (T) for the current request -local increment = (period / limit) * emission_cost +local count = redis.call('ZCARD', key) --- Max allowed burst time delay offset (Burst * T) -local burst_offset = burst * (period / limit) - --- Get the last calculated theoretical arrival time (TAT) from Redis -local tat = tonumber(redis.call("GET", key)) -if tat == nil then - -- First request, initialize TAT to current time - tat = now -else - -- If last TAT lags behind current time, connection has been idle, reset to current time - tat = math.max(tat, now) +if count >= limit then + return {current_time, 0, count, now} end --- Calculate what the new TAT would be if this request is accepted -local new_tat = tat + increment +local member = tostring(now) .. ':' .. tostring(t[2]) +redis.call('ZADD', key, now, member) +redis.call('EXPIRE', key, math.ceil(period + 60)) --- Core sliding window check: if new TAT minus burst offset still exceeds current time, the request exceeds the limit -if new_tat - burst_offset > now then - -- return 0 -- trigger rate limiting, reject - return {current_time, 0, tostring(tat), tostring(now)} -else - -- Allow, update TAT timestamp in Redis - redis.call("SET", key, new_tat) - -- Dynamically set expiry, longer than one full sliding window to prevent memory leak - redis.call("EXPIRE", key, math.ceil(period + 60)) - -- return 1 -- allow - return {current_time, 1, tostring(new_tat), tostring(now)} -end +return {current_time, 1, count + emission_cost, now} \ No newline at end of file diff --git a/bfe_util/limit_rate/redis_qpm_limiter.go b/bfe_util/limit_rate/redis_qpm_limiter.go index 8b0c8d851..1e989d87c 100644 --- a/bfe_util/limit_rate/redis_qpm_limiter.go +++ b/bfe_util/limit_rate/redis_qpm_limiter.go @@ -39,7 +39,6 @@ func NewQPMLimiter(redisKey string, burst int64, period int64, limit int64) *QPM func (l *QPMLimiter) Check(reqToConsume int64, agent *RedisLRAgent) (bool, int64, float64, error) { tat := float64(0) limit := l.limit.Load() - burst := l.burst.Load() period := l.period.Load() if reqToConsume <= 0 { @@ -55,7 +54,7 @@ func (l *QPMLimiter) Check(reqToConsume int64, agent *RedisLRAgent) (bool, int64 } //check remote - currentTime, isAllowed, tat, err := agent.QpmCheck(l.redisKey, burst, period, limit, reqToConsume) + currentTime, isAllowed, tat, err := agent.QpmCheck(l.redisKey, period, limit, reqToConsume) if err != nil { return false, currentTime, tat, err } diff --git a/bfe_util/limit_rate/testdata/name_conf.data b/bfe_util/limit_rate/testdata/name_conf.data new file mode 100644 index 000000000..db8137288 --- /dev/null +++ b/bfe_util/limit_rate/testdata/name_conf.data @@ -0,0 +1,12 @@ +{ + "Version": "1", + "Config": { + "testredis" : [ + { + "Host": "127.0.0.1", + "Port": 6379, + "Weight": 10 + } + ] + } +}