200 lines
8.1 KiB
Go
200 lines
8.1 KiB
Go
package middleware
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"strconv"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/i18n"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"github.com/QuantumNous/new-api/setting"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestModelRedisRateLimitUsesUTCRegardlessOfLocalTimezone(t *testing.T) {
|
|
redisServer, redisClient := useRateLimitMiniRedis(t)
|
|
previousLocation := time.Local
|
|
time.Local = time.FixedZone("test-utc-plus-eight", 8*60*60)
|
|
t.Cleanup(func() { time.Local = previousLocation })
|
|
|
|
ctx := context.Background()
|
|
recordKey := "rateLimit:model-utc-record"
|
|
recordRedisRequest(ctx, redisClient, recordKey, 2)
|
|
recorded, err := redisClient.LIndex(ctx, recordKey, 0).Result()
|
|
require.NoError(t, err)
|
|
recordedAt, err := time.Parse(modelRateLimitTimeFormat, recorded)
|
|
require.NoError(t, err)
|
|
assert.WithinDuration(t, time.Now().UTC(), recordedAt, 2*time.Second)
|
|
|
|
checkKey := "rateLimit:model-utc-check"
|
|
withinWindow := time.Now().UTC().Add(-30 * time.Second).Format(modelRateLimitTimeFormat)
|
|
_, err = redisServer.Push(checkKey, withinWindow, withinWindow)
|
|
require.NoError(t, err)
|
|
allowed, err := checkRedisRateLimit(ctx, redisClient, checkKey, 2, 60)
|
|
require.NoError(t, err)
|
|
assert.False(t, allowed, "an existing UTC timestamp inside the window must remain limited on a non-UTC host")
|
|
}
|
|
|
|
var modelRateLimitTestUsers atomic.Int64
|
|
|
|
func TestModelRateLimitStreamFailuresDoNotConsumeSuccessLimit(t *testing.T) {
|
|
require.NoError(t, i18n.Init())
|
|
for _, backend := range []string{"memory", "redis"} {
|
|
for _, totalLimit := range []int{0, 2} {
|
|
t.Run(fmt.Sprintf("%s/total=%d", backend, totalLimit), func(t *testing.T) {
|
|
subject := strconv.Itoa(7200000 + int(modelRateLimitTestUsers.Add(1)))
|
|
handler := memoryRateLimitHandler(subject, 60, totalLimit, 1)
|
|
if backend == "redis" {
|
|
useRateLimitMiniRedis(t)
|
|
handler = redisRateLimitHandler(subject, 60, totalLimit, 1)
|
|
}
|
|
router := gin.New()
|
|
router.GET("/:outcome", handler, func(c *gin.Context) {
|
|
status := relaycommon.NewStreamStatus()
|
|
if c.Param("outcome") == "failed" {
|
|
status.MarkFailed("server_error", "", 0)
|
|
} else {
|
|
status.MarkCompleted()
|
|
}
|
|
common.SetContextKey(c, constant.ContextKeyResponseStreamStatus, status)
|
|
c.Status(http.StatusOK)
|
|
})
|
|
assert.Equal(t, http.StatusOK, performRateLimitRequest(router, "/failed", "127.0.0.1:1000").Code)
|
|
if totalLimit > 0 {
|
|
assert.Equal(t, http.StatusOK, performRateLimitRequest(router, "/failed", "127.0.0.1:1000").Code)
|
|
} else {
|
|
assert.Equal(t, http.StatusOK, performRateLimitRequest(router, "/completed", "127.0.0.1:1000").Code)
|
|
}
|
|
assert.Equal(t, http.StatusTooManyRequests, performRateLimitRequest(router, "/completed", "127.0.0.1:1000").Code)
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestModelMemoryRateLimitReservesConcurrentSuccessAdmission(t *testing.T) {
|
|
require.NoError(t, i18n.Init())
|
|
subject := strconv.Itoa(7200000 + int(modelRateLimitTestUsers.Add(1)))
|
|
entered, release, finished := make(chan struct{}), make(chan struct{}), make(chan struct{})
|
|
router := gin.New()
|
|
router.GET("/:outcome", memoryRateLimitHandler(subject, 60, 0, 1), func(c *gin.Context) {
|
|
if c.Param("outcome") == "failed" {
|
|
close(entered)
|
|
<-release
|
|
status := relaycommon.NewStreamStatus()
|
|
status.MarkFailed("server_error", "", 0)
|
|
common.SetContextKey(c, constant.ContextKeyResponseStreamStatus, status)
|
|
}
|
|
c.Status(http.StatusOK)
|
|
})
|
|
go func() {
|
|
defer close(finished)
|
|
assert.Equal(t, http.StatusOK, performRateLimitRequest(router, "/failed", "127.0.0.1:1000").Code)
|
|
}()
|
|
<-entered
|
|
assert.Equal(t, http.StatusTooManyRequests, performRateLimitRequest(router, "/completed", "127.0.0.1:1000").Code)
|
|
close(release)
|
|
<-finished
|
|
assert.Equal(t, http.StatusOK, performRateLimitRequest(router, "/completed", "127.0.0.1:1000").Code)
|
|
assert.Equal(t, http.StatusTooManyRequests, performRateLimitRequest(router, "/completed", "127.0.0.1:1000").Code)
|
|
}
|
|
|
|
// The memory limiter answers like the Redis limiter: an OpenAI error body
|
|
// naming the limit that was reached, not an empty 429.
|
|
func TestModelMemoryRateLimitRejectionsReturnOpenAIError(t *testing.T) {
|
|
require.NoError(t, i18n.Init())
|
|
for _, tc := range []struct {
|
|
name string
|
|
totalLimit, successLimit int
|
|
messageKey string
|
|
}{
|
|
{name: "total", totalLimit: 1, messageKey: i18n.MsgRateLimitTotalReached},
|
|
{name: "success", successLimit: 1, messageKey: i18n.MsgRateLimitReached},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
subject := strconv.Itoa(7200000 + int(modelRateLimitTestUsers.Add(1)))
|
|
router := gin.New()
|
|
router.GET("/", func(c *gin.Context) {
|
|
c.Set(common.RequestIdKey, "req-limited")
|
|
}, memoryRateLimitHandler(subject, 60, tc.totalLimit, tc.successLimit), func(c *gin.Context) {
|
|
c.Status(http.StatusOK)
|
|
})
|
|
require.Equal(t, http.StatusOK, performRateLimitRequest(router, "/", "127.0.0.1:1000").Code)
|
|
|
|
resp := performRateLimitRequest(router, "/", "127.0.0.1:1000")
|
|
assert.Equal(t, http.StatusTooManyRequests, resp.Code)
|
|
message := i18n.Translate("", tc.messageKey, map[string]any{"Minutes": setting.ModelRequestRateLimitDurationMinutes, "Max": 1})
|
|
want, err := common.Marshal(gin.H{"error": gin.H{
|
|
"message": common.MessageWithRequestId(message, "req-limited"),
|
|
"type": "new_api_error",
|
|
"code": "",
|
|
}})
|
|
require.NoError(t, err)
|
|
assert.JSONEq(t, string(want), resp.Body.String())
|
|
})
|
|
}
|
|
}
|
|
|
|
// A group with its own limit counts each user's requests on its own; groups
|
|
// without one keep sharing the user's global limit.
|
|
func TestModelRateLimitCountsConfiguredGroupsSeparately(t *testing.T) {
|
|
require.NoError(t, i18n.Init())
|
|
previousEnabled, previousMinutes := setting.ModelRequestRateLimitEnabled, setting.ModelRequestRateLimitDurationMinutes
|
|
previousTotal, previousSuccess := setting.ModelRequestRateLimitCount, setting.ModelRequestRateLimitSuccessCount
|
|
previousGroups, previousRedisEnabled := setting.ModelRequestRateLimitGroup, common.RedisEnabled
|
|
t.Cleanup(func() {
|
|
setting.ModelRequestRateLimitEnabled, setting.ModelRequestRateLimitDurationMinutes = previousEnabled, previousMinutes
|
|
setting.ModelRequestRateLimitCount, setting.ModelRequestRateLimitSuccessCount = previousTotal, previousSuccess
|
|
setting.ModelRequestRateLimitGroup, common.RedisEnabled = previousGroups, previousRedisEnabled
|
|
})
|
|
setting.ModelRequestRateLimitEnabled, setting.ModelRequestRateLimitDurationMinutes = true, 1
|
|
common.RedisEnabled = false
|
|
|
|
// Groups a and b have their own limits; c and d have none and share the global one.
|
|
steps := []struct {
|
|
group string
|
|
want int
|
|
}{
|
|
{"a", http.StatusOK}, {"a", http.StatusTooManyRequests},
|
|
{"b", http.StatusOK}, {"b", http.StatusTooManyRequests},
|
|
{"c", http.StatusOK}, {"d", http.StatusTooManyRequests},
|
|
}
|
|
// The Redis token bucket binds to the first Redis client it sees and another
|
|
// test already drives it, so the total limit runs on the memory backend only.
|
|
for _, tc := range []struct {
|
|
backend, limit string
|
|
counts [2]int // [total, success] for groups a and b and for the global limit
|
|
}{
|
|
{"memory", "success", [2]int{0, 1}},
|
|
{"redis", "success", [2]int{0, 1}},
|
|
{"memory", "total", [2]int{1, 1000}},
|
|
} {
|
|
t.Run(tc.backend+"/"+tc.limit, func(t *testing.T) {
|
|
if tc.backend == "redis" {
|
|
useRateLimitMiniRedis(t)
|
|
}
|
|
setting.ModelRequestRateLimitCount, setting.ModelRequestRateLimitSuccessCount = tc.counts[0], tc.counts[1]
|
|
setting.ModelRequestRateLimitGroup = map[string][2]int{"a": tc.counts, "b": tc.counts}
|
|
userID := 7200000 + int(modelRateLimitTestUsers.Add(1))
|
|
router := gin.New()
|
|
router.GET("/:group", func(c *gin.Context) {
|
|
c.Set("id", userID)
|
|
common.SetContextKey(c, constant.ContextKeyTokenGroup, c.Param("group"))
|
|
}, ModelRequestRateLimit(), func(c *gin.Context) {
|
|
c.Status(http.StatusOK)
|
|
})
|
|
|
|
for _, step := range steps {
|
|
assert.Equal(t, step.want, performRateLimitRequest(router, "/"+step.group, "127.0.0.1:1000").Code, "group %s", step.group)
|
|
}
|
|
})
|
|
}
|
|
}
|