diff --git a/docs/config-zh.md b/docs/config-zh.md
index 30277e2..07c5b67 100644
--- a/docs/config-zh.md
+++ b/docs/config-zh.md
@@ -43,6 +43,7 @@
| `llms[].embedding_model` | `string` | LLM 的 Embedding 模型。例如 `text-embedding-3-small`。如果用于 Embedding，则不能为空。如果此 LLM 被使用，则不能与 `model` 同时为空。**注意：** 初次使用后请勿直接修改，应添加新的 LLM 配置。 | | 条件性必需 |
| `llms[].tts_model` | `string` | LLM 的文本转语音 (TTS) 模型。 | | 否 |
| `llms[].temperature` | `float32` | LLM 的温度 (0-2)。 | `0.0` | 否 |
+| `llms[].rpm` | `int` | 每分钟请求数限制 (Requests Per Minute)。用于控制 API 调用频率，避免超过提供商的速率限制。`0` 表示不限流。 | `0` | 否 |
### Jina AI 配置 (`jina`)
diff --git a/docs/config.md b/docs/config.md
index 5230c54..a6d7f3e 100644
--- a/docs/config.md
+++ b/docs/config.md
@@ -43,6 +43,7 @@ This section defines the list of available Large Language Models. At least one L
| `llms[].embedding_model` | `string` | Embedding model of the LLM. E.g., `text-embedding-3-small`. Cannot be empty if used for embedding. If this LLM is used, cannot be empty along with `model`. **Note:** Do not modify directly after initial use; add a new LLM configuration instead. | | Conditionally Required |
| `llms[].tts_model` | `string` | The Text-to-Speech (TTS) model of the LLM. | | No |
| `llms[].temperature` | `float32` | Temperature of the LLM (0-2). | `0.0` | No |
+| `llms[].rpm` | `int` | Requests per minute limit. Controls API call frequency to avoid exceeding provider rate limits. `0` means no rate limiting. | `0` | No |
### Jina AI Configuration (`jina`)
diff --git a/pkg/api/api.go b/pkg/api/api.go
index 69f231e..fd458a5 100644
--- a/pkg/api/api.go
+++ b/pkg/api/api.go
@@ -61,13 +61,16 @@ type API interface {
) (resp *QueryRSSHubWebsitesResponse, err error)
QueryRSSHubRoutes(ctx context.Context, req *QueryRSSHubRoutesRequest) (resp *QueryRSSHubRoutesResponse, err error)
+ AddFeedSource(ctx context.Context, req *AddFeedSourceRequest) (resp *AddFeedSourceResponse, err error)
+
Write(ctx context.Context, req *WriteRequest) (resp *WriteResponse, err error) // WARN: beta!!!
Query(ctx context.Context, req *QueryRequest) (resp *QueryResponse, err error)
}
type Config struct {
- RSSHubEndpoint string
- LLM string
+ RSSHubEndpoint string
+ RSSHubAccessKey string
+ LLM string
}
func (c *Config) Validate() error {
@@ -78,6 +81,7 @@ func (c *Config) Validate() error {
func (c *Config) From(app *config.App) *Config {
c.RSSHubEndpoint = app.Scrape.RSSHubEndpoint
+ c.RSSHubAccessKey = app.Scrape.RSSHubAccessKey
c.LLM = app.API.LLM
return c
@@ -143,6 +147,12 @@ type RSSHubRoute struct {
Features map[string]any `json:"features,omitempty"`
}
+type AddFeedSourceRequest struct {
+ Source config.ScrapeSource `json:"source"`
+}
+
+type AddFeedSourceResponse struct{}
+
type WriteRequest struct { // Beta.
Feeds []*model.Feed `json:"feeds"`
}
@@ -277,6 +287,19 @@ func (a *api) Reload(app *config.App) error {
return nil
}
+// appendAccessKeyToURL adds the RSSHub access key to the URL if configured
+func (a *api) appendAccessKeyToURL(url string) string {
+ if a.Config().RSSHubAccessKey != "" {
+ if strings.Contains(url, "?") {
+ return url + "&key=" + a.Config().RSSHubAccessKey
+ }
+
+ return url + "?key=" + a.Config().RSSHubAccessKey
+ }
+
+ return url
+}
+
func (a *api) QueryAppConfigSchema(
ctx context.Context,
req *QueryAppConfigSchemaRequest,
@@ -313,7 +336,7 @@ func (a *api) QueryRSSHubCategories(
ctx context.Context,
req *QueryRSSHubCategoriesRequest,
) (resp *QueryRSSHubCategoriesResponse, err error) {
- url := a.Config().RSSHubEndpoint + "/api/namespace"
+ url := a.appendAccessKeyToURL(a.Config().RSSHubEndpoint + "/api/namespace")
// New request.
forwardReq, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
@@ -357,7 +380,7 @@ func (a *api) QueryRSSHubWebsites(
return nil, ErrBadRequest(errors.New("category is required"))
}
- url := a.Config().RSSHubEndpoint + "/api/category/" + req.Category
+ url := a.appendAccessKeyToURL(a.Config().RSSHubEndpoint + "/api/category/" + req.Category)
// New request.
forwardReq, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
@@ -407,7 +430,7 @@ func (a *api) QueryRSSHubRoutes(
return nil, ErrBadRequest(errors.New("website id is required"))
}
- url := a.Config().RSSHubEndpoint + "/api/namespace/" + req.WebsiteID
+ url := a.appendAccessKeyToURL(a.Config().RSSHubEndpoint + "/api/namespace/" + req.WebsiteID)
// New request.
forwardReq, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
@@ -447,6 +470,48 @@ func (a *api) QueryRSSHubRoutes(
return resp, nil
}
+func (a *api) AddFeedSource(
+ ctx context.Context,
+ req *AddFeedSourceRequest,
+) (resp *AddFeedSourceResponse, err error) {
+ ctx = telemetry.StartWith(ctx, append(a.TelemetryLabels(), telemetrymodel.KeyOperation, "AddFeedSource")...)
+ defer func() { telemetry.End(ctx, err) }()
+
+ // Validate request.
+ if req.Source.Name == "" {
+ return nil, ErrBadRequest(errors.New("source name is required"))
+ }
+ if req.Source.RSS == nil {
+ return nil, ErrBadRequest(errors.New("rss config is required"))
+ }
+ if req.Source.RSS.URL == "" && req.Source.RSS.RSSHubRoutePath == "" {
+ return nil, ErrBadRequest(errors.New("either url or rsshub_route_path is required"))
+ }
+ if req.Source.RSS.URL != "" && req.Source.RSS.RSSHubRoutePath != "" {
+ return nil, ErrBadRequest(errors.New("url and rsshub_route_path cannot be set at the same time"))
+ }
+
+ // Get current config.
+ appConfig := a.Dependencies().ConfigManager.AppConfig()
+
+ // Check if source name already exists.
+ for _, source := range appConfig.Scrape.Sources {
+ if source.Name == req.Source.Name {
+ return nil, ErrBadRequest(errors.New("source name already exists"))
+ }
+ }
+
+ // Add new source.
+ appConfig.Scrape.Sources = append(appConfig.Scrape.Sources, req.Source)
+
+ // Save config.
+ if err := a.Dependencies().ConfigManager.SaveAppConfig(appConfig); err != nil {
+ return nil, ErrInternal(errors.Wrap(err, "save app config"))
+ }
+
+ return &AddFeedSourceResponse{}, nil
+}
+
func (a *api) Write(ctx context.Context, req *WriteRequest) (resp *WriteResponse, err error) {
ctx = telemetry.StartWith(ctx, append(a.TelemetryLabels(), telemetrymodel.KeyOperation, "Write")...)
defer func() { telemetry.End(ctx, err) }()
@@ -584,6 +649,15 @@ func (m *mockAPI) QueryRSSHubRoutes(
return args.Get(0).(*QueryRSSHubRoutesResponse), args.Error(1)
}
+func (m *mockAPI) AddFeedSource(
+ ctx context.Context,
+ req *AddFeedSourceRequest,
+) (resp *AddFeedSourceResponse, err error) {
+ args := m.Called(ctx, req)
+
+ return args.Get(0).(*AddFeedSourceResponse), args.Error(1)
+}
+
func (m *mockAPI) Query(ctx context.Context, req *QueryRequest) (resp *QueryResponse, err error) {
args := m.Called(ctx, req)
diff --git a/pkg/api/http/http.go b/pkg/api/http/http.go
index 419ec18..1ba8b83 100644
--- a/pkg/api/http/http.go
+++ b/pkg/api/http/http.go
@@ -95,6 +95,7 @@ func new(instance string, app *config.App, dependencies Dependencies) (Server, e
router.Handle("/query_rsshub_categories", jsonrpc.API(api.QueryRSSHubCategories))
router.Handle("/query_rsshub_websites", jsonrpc.API(api.QueryRSSHubWebsites))
router.Handle("/query_rsshub_routes", jsonrpc.API(api.QueryRSSHubRoutes))
+ router.Handle("/add_feed_source", jsonrpc.API(api.AddFeedSource))
router.Handle("/query", jsonrpc.API(api.Query))
httpServer := &http.Server{Addr: config.Address, Handler: router}
diff --git a/pkg/api/mcp/mcp.go b/pkg/api/mcp/mcp.go
index 2349768..ca12afd 100644
--- a/pkg/api/mcp/mcp.go
+++ b/pkg/api/mcp/mcp.go
@@ -132,6 +132,7 @@ func new(instance string, app *config.App, dependencies Dependencies) (Server, e
func registerTools(h *mcpserver.MCPServer, s *server) {
registerConfigTools(h, s)
registerRSSHubTools(h, s)
+ registerFeedSourceTools(h, s)
h.AddTool(mcp.NewTool("query",
mcp.WithDescription("Query feeds with semantic search. You can query any latest messages. "+
@@ -203,6 +204,26 @@ func registerRSSHubTools(h *mcpserver.MCPServer, s *server) {
), mcpserver.ToolHandlerFunc(s.queryRSSHubRoutes))
}
+func registerFeedSourceTools(h *mcpserver.MCPServer, s *server) {
+ h.AddTool(mcp.NewTool("add_feed_source",
+ mcp.WithDescription("Add a new feed source to the app config. "+
+ "The source will be added to the scrape.sources list. "+
+ "You should confirm with the user before adding the source."),
+ mcp.WithString("name",
+ mcp.Required(),
+ mcp.Description("The name of the feed source. It must be unique."),
+ ),
+ mcp.WithString("url",
+ mcp.Description("The direct RSS feed URL. e.g. https://tech.meituan.com/feed. "+
+ "Either url or rsshub_route_path must be provided, but not both."),
+ ),
+ mcp.WithString("rsshub_route_path",
+ mcp.Description("The RSSHub route path. e.g. telegram/channel/zrj96. "+
+ "Either url or rsshub_route_path must be provided, but not both."),
+ ),
+ ), mcpserver.ToolHandlerFunc(s.addFeedSource))
+}
+
// --- Implementation code block ---
type server struct {
*component.Base[Config, Dependencies]
@@ -334,6 +355,44 @@ func (s *server) queryRSSHubRoutes(ctx context.Context, req mcp.CallToolRequest)
return s.response(string(b)), nil
}
+func (s *server) addFeedSource(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+ // Parse arguments.
+ name, ok := req.Params.Arguments["name"].(string)
+ if !ok || name == "" {
+ return s.error(errors.New("name is required")), nil
+ }
+
+ url, _ := req.Params.Arguments["url"].(string)
+ rsshubRoutePath, _ := req.Params.Arguments["rsshub_route_path"].(string)
+
+ // Validate that either url or rsshub_route_path is provided.
+ if url == "" && rsshubRoutePath == "" {
+ return s.error(errors.New("either url or rsshub_route_path must be provided")), nil
+ }
+ if url != "" && rsshubRoutePath != "" {
+ return s.error(errors.New("url and rsshub_route_path cannot be set at the same time")), nil
+ }
+
+ // Build the source config.
+ source := config.ScrapeSource{
+ Name: name,
+ RSS: &config.ScrapeSourceRSS{},
+ }
+ if url != "" {
+ source.RSS.URL = url
+ } else {
+ source.RSS.RSSHubRoutePath = rsshubRoutePath
+ }
+
+ // Forward request to API.
+ _, err := s.Dependencies().API.AddFeedSource(ctx, &api.AddFeedSourceRequest{Source: source})
+ if err != nil {
+ return s.error(errors.Wrap(err, "add feed source")), nil
+ }
+
+ return s.response("Feed source added successfully. The source name is: " + name), nil
+}
+
func (s *server) query(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
// Parse arguments.
query, ok := req.Params.Arguments["query"].(string)
diff --git a/pkg/config/config.go b/pkg/config/config.go
index cbfc2a4..f2bdb17 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -92,6 +92,7 @@ type LLM struct {
EmbeddingModel string `yaml:"embedding_model,omitempty" json:"embedding_model,omitempty" desc:"The embedding model of the LLM. e.g. text-embedding-3-small. Can not be empty with model at same time when api.llm is set. NOTE: Once used, do not modify it directly, instead, add a new LLM configuration."`
TTSModel string `yaml:"tts_model,omitempty" json:"tts_model,omitempty" desc:"The TTS model of the LLM."`
Temperature float32 `yaml:"temperature,omitempty" json:"temperature,omitempty" desc:"The temperature (0-2) of the LLM. Default: 0.0"`
+ RPM int `yaml:"rpm,omitempty" json:"rpm,omitempty" desc:"Requests per minute limit for this LLM. 0 means no limit. Default: 0"`
}
type Scrape struct {
diff --git a/pkg/llm/gemini.go b/pkg/llm/gemini.go
index a219860..16b81c2 100644
--- a/pkg/llm/gemini.go
+++ b/pkg/llm/gemini.go
@@ -36,7 +36,8 @@ import (
type gemini struct {
*component.Base[Config, struct{}]
text
- hc *http.Client
+ hc *http.Client
+ rateLimiter RateLimiter
embeddingSpliter embeddingSpliter
}
@@ -53,13 +54,17 @@ func newGemini(c *Config) LLM {
Config: c,
})
+ rateLimiter := NewRateLimiter(c.RPM)
+
return &gemini{
Base: base,
text: &openaiText{
- Base: base,
- client: client,
+ Base: base,
+ client: client,
+ rateLimiter: rateLimiter,
},
hc: &http.Client{},
+ rateLimiter: rateLimiter,
embeddingSpliter: embeddingSpliter,
}
}
@@ -72,6 +77,11 @@ func (g *gemini) WAV(ctx context.Context, text string, speakers []Speaker) (r io
return nil, errors.New("tts model is not set")
}
+ // 应用限流
+ if err := g.rateLimiter.Wait(ctx); err != nil {
+ return nil, errors.Wrap(err, "rate limiter wait")
+ }
+
reqPayload, err := buildWAVRequestPayload(text, speakers)
if err != nil {
return nil, errors.Wrap(err, "build wav request payload")
diff --git a/pkg/llm/llm.go b/pkg/llm/llm.go
index 91fb21d..0a4b1fc 100644
--- a/pkg/llm/llm.go
+++ b/pkg/llm/llm.go
@@ -71,6 +71,7 @@ type Config struct {
APIKey string
Model, EmbeddingModel, TTSModel string
Temperature float32
+ RPM int // Requests per minute limit
}
type ProviderType string
@@ -200,6 +201,7 @@ func (c *FactoryConfig) From(app *config.App) {
EmbeddingModel: llm.EmbeddingModel,
TTSModel: llm.TTSModel,
Temperature: llm.Temperature,
+ RPM: llm.RPM,
})
}
}
diff --git a/pkg/llm/openai.go b/pkg/llm/openai.go
index 4ae3537..8876d0e 100644
--- a/pkg/llm/openai.go
+++ b/pkg/llm/openai.go
@@ -33,6 +33,7 @@ import (
type openai struct {
*component.Base[Config, struct{}]
text
+ rateLimiter RateLimiter
}
func newOpenAI(c *Config) LLM {
@@ -47,13 +48,17 @@ func newOpenAI(c *Config) LLM {
Config: c,
})
+ rateLimiter := NewRateLimiter(c.RPM)
+
return &openai{
Base: base,
text: &openaiText{
Base: base,
client: client,
embeddingSpliter: embeddingSpliter,
+ rateLimiter: rateLimiter,
},
+ rateLimiter: rateLimiter,
}
}
@@ -66,6 +71,7 @@ type openaiText struct {
client *oai.Client
embeddingSpliter embeddingSpliter
+ rateLimiter RateLimiter
}
func (o *openaiText) String(ctx context.Context, messages []string) (value string, err error) {
@@ -90,6 +96,11 @@ func (o *openaiText) String(ctx context.Context, messages []string) (value strin
Temperature: config.Temperature,
}
+ // 应用限流
+ if err := o.rateLimiter.Wait(ctx); err != nil {
+ return "", errors.Wrap(err, "rate limiter wait")
+ }
+
resp, err := o.client.CreateChatCompletion(ctx, req)
if err != nil {
return "", errors.Wrap(err, "create chat completion")
@@ -140,6 +151,11 @@ func (o *openaiText) Embedding(ctx context.Context, s string) (value []float32,
if config.EmbeddingModel == "" {
return nil, errors.New("embedding model is not set")
}
+
+ if err := o.rateLimiter.Wait(ctx); err != nil {
+ return nil, errors.Wrap(err, "rate limiter wait")
+ }
+
vec, err := o.client.CreateEmbeddings(ctx, oai.EmbeddingRequest{
Input: []string{s},
Model: oai.EmbeddingModel(config.EmbeddingModel),
diff --git a/pkg/llm/ratelimiter.go b/pkg/llm/ratelimiter.go
new file mode 100644
index 0000000..1ba7c72
--- /dev/null
+++ b/pkg/llm/ratelimiter.go
@@ -0,0 +1,148 @@
+// Copyright (C) 2025 wangyusong
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+package llm
+
+import (
+ "context"
+ "sync"
+ "time"
+
+ "github.com/pkg/errors"
+
+ "github.com/glidea/zenfeed/pkg/telemetry/log"
+)
+
+// RateLimiter 定义限流器接口
+type RateLimiter interface {
+ // Wait 等待直到可以执行请求,如果 context 被取消则返回错误
+ Wait(ctx context.Context) error
+ // TryAcquire 尝试获取一个令牌,如果成功返回 true,否则返回 false
+ TryAcquire() bool
+}
+
+// noopRateLimiter 是一个无操作的限流器,用于没有限流需求的场景
+type noopRateLimiter struct{}
+
+func newNoopRateLimiter() RateLimiter {
+ return &noopRateLimiter{}
+}
+
+func (n *noopRateLimiter) Wait(ctx context.Context) error {
+ return nil
+}
+
+func (n *noopRateLimiter) TryAcquire() bool {
+ return true
+}
+
+// tokenBucketRateLimiter 使用令牌桶算法实现的限流器
+type tokenBucketRateLimiter struct {
+ rpm int // 每分钟请求数限制
+ interval time.Duration // 每个令牌的生成间隔
+ tokens chan struct{} // 令牌桶
+ mu sync.Mutex
+ stopCh chan struct{}
+ stopped bool
+}
+
+// newTokenBucketRateLimiter 创建一个基于令牌桶算法的限流器
+// rpm: 每分钟请求数限制
+func newTokenBucketRateLimiter(rpm int) RateLimiter {
+ if rpm <= 0 {
+ return newNoopRateLimiter()
+ }
+
+ // 计算每个令牌的生成间隔
+ interval := time.Minute / time.Duration(rpm)
+
+ // 令牌桶容量设置为 rpm,允许突发流量
+ limiter := &tokenBucketRateLimiter{
+ rpm: rpm,
+ interval: interval,
+ tokens: make(chan struct{}, rpm),
+ stopCh: make(chan struct{}),
+ }
+
+ // 初始化令牌桶,填满令牌
+ for range rpm {
+ limiter.tokens <- struct{}{}
+ }
+
+ // 启动令牌生成器
+ go limiter.refillTokens()
+
+ log.Info(context.Background(), "rate limiter created", "rpm", rpm, "interval", interval)
+
+ return limiter
+}
+
+// refillTokens 定期向令牌桶中添加令牌
+func (t *tokenBucketRateLimiter) refillTokens() {
+ ticker := time.NewTicker(t.interval)
+ defer ticker.Stop()
+
+ for {
+ select {
+ case <-ticker.C:
+ // 尝试添加一个令牌,如果桶满了则丢弃
+ select {
+ case t.tokens <- struct{}{}:
+ default:
+ // 令牌桶已满,丢弃这个令牌
+ }
+ case <-t.stopCh:
+ return
+ }
+ }
+}
+
+// Wait 等待直到可以执行请求
+func (t *tokenBucketRateLimiter) Wait(ctx context.Context) error {
+ select {
+ case <-t.tokens:
+ return nil
+ case <-ctx.Done():
+ return errors.Wrap(ctx.Err(), "rate limiter wait cancelled")
+ case <-t.stopCh:
+ return errors.New("rate limiter stopped")
+ }
+}
+
+// TryAcquire 尝试获取一个令牌
+func (t *tokenBucketRateLimiter) TryAcquire() bool {
+ select {
+ case <-t.tokens:
+ return true
+ default:
+ return false
+ }
+}
+
+// Stop 停止限流器
+func (t *tokenBucketRateLimiter) Stop() {
+ t.mu.Lock()
+ defer t.mu.Unlock()
+
+ if !t.stopped {
+ close(t.stopCh)
+ t.stopped = true
+ }
+}
+
+// NewRateLimiter 根据 RPM 配置创建限流器
+func NewRateLimiter(rpm int) RateLimiter {
+ return newTokenBucketRateLimiter(rpm)
+}
diff --git a/pkg/llm/ratelimiter_test.go b/pkg/llm/ratelimiter_test.go
new file mode 100644
index 0000000..32a568d
--- /dev/null
+++ b/pkg/llm/ratelimiter_test.go
@@ -0,0 +1,149 @@
+// Copyright (C) 2025 wangyusong
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+package llm
+
+import (
+ "context"
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/assert"
+)
+
+func TestNoopRateLimiter(t *testing.T) {
+ limiter := newNoopRateLimiter()
+ ctx := context.Background()
+
+ // 无限制,应该立即返回
+ err := limiter.Wait(ctx)
+ assert.NoError(t, err)
+
+ // TryAcquire 应该总是返回 true
+ assert.True(t, limiter.TryAcquire())
+}
+
+func TestTokenBucketRateLimiter_Basic(t *testing.T) {
+ // 创建一个每分钟 60 个请求的限流器 (即每秒 1 个请求)
+ limiter := NewRateLimiter(60)
+ ctx := context.Background()
+
+ // 第一个请求应该立即成功
+ start := time.Now()
+ err := limiter.Wait(ctx)
+ assert.NoError(t, err)
+ assert.Less(t, time.Since(start), 100*time.Millisecond)
+}
+
+func TestTokenBucketRateLimiter_RateLimit(t *testing.T) {
+ // 创建一个每分钟 6 个请求的限流器 (即每 10 秒 1 个请求)
+ rpm := 6
+ limiter := NewRateLimiter(rpm)
+ ctx := context.Background()
+
+ // 快速消耗所有令牌
+ for i := 0; i < rpm; i++ {
+ assert.True(t, limiter.TryAcquire())
+ }
+
+ // 此时令牌桶应该为空,TryAcquire 应该返回 false
+ assert.False(t, limiter.TryAcquire())
+
+ // Wait 应该等待直到有新令牌
+ start := time.Now()
+ err := limiter.Wait(ctx)
+ assert.NoError(t, err)
+ elapsed := time.Since(start)
+
+ // 应该等待大约 10 秒 (60秒/6个请求)
+ expectedWait := time.Minute / time.Duration(rpm)
+ assert.Greater(t, elapsed, expectedWait-100*time.Millisecond)
+ assert.Less(t, elapsed, expectedWait+500*time.Millisecond)
+}
+
+func TestTokenBucketRateLimiter_ContextCancellation(t *testing.T) {
+ // 创建一个每分钟 1 个请求的限流器
+ limiter := NewRateLimiter(1)
+
+ // 消耗唯一的令牌
+ assert.True(t, limiter.TryAcquire())
+
+ // 创建一个会被取消的 context
+ ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+ defer cancel()
+
+ // Wait 应该在 context 被取消时返回错误
+ err := limiter.Wait(ctx)
+ assert.Error(t, err)
+ assert.Contains(t, err.Error(), "rate limiter wait cancelled")
+}
+
+func TestTokenBucketRateLimiter_ZeroRPM(t *testing.T) {
+ // RPM 为 0 应该返回 noop 限流器
+ limiter := NewRateLimiter(0)
+ ctx := context.Background()
+
+ // 应该没有限制
+ for i := 0; i < 100; i++ {
+ err := limiter.Wait(ctx)
+ assert.NoError(t, err)
+ }
+}
+
+func TestTokenBucketRateLimiter_NegativeRPM(t *testing.T) {
+ // 负数 RPM 应该返回 noop 限流器
+ limiter := NewRateLimiter(-10)
+ ctx := context.Background()
+
+ // 应该没有限制
+ for i := 0; i < 100; i++ {
+ err := limiter.Wait(ctx)
+ assert.NoError(t, err)
+ }
+}
+
+func TestTokenBucketRateLimiter_Concurrent(t *testing.T) {
+ // 创建一个每分钟 60 个请求的限流器
+ rpm := 60
+ limiter := NewRateLimiter(rpm)
+ ctx := context.Background()
+
+ // 并发请求
+ concurrency := 10
+ done := make(chan bool, concurrency)
+
+ start := time.Now()
+ for i := 0; i < concurrency; i++ {
+ go func() {
+ err := limiter.Wait(ctx)
+ assert.NoError(t, err)
+ done <- true
+ }()
+ }
+
+ // 等待所有请求完成
+ for i := 0; i < concurrency; i++ {
+ <-done
+ }
+
+ elapsed := time.Since(start)
+
+ // 前 rpm 个请求应该立即完成,剩余的需要等待
+ // 由于有 60 个初始令牌,前 60 个请求应该很快完成
+ // 但我们只发送了 10 个请求,所以应该很快完成
+ if concurrency <= rpm {
+ assert.Less(t, elapsed, 1*time.Second)
+ }
+}
diff --git a/pkg/model/model.go b/pkg/model/model.go
index 32a3079..02808ff 100644
--- a/pkg/model/model.go
+++ b/pkg/model/model.go
@@ -171,7 +171,11 @@ func (ls Labels) MarshalJSON() ([]byte, error) {
return nil, errors.Wrap(err, "write ending brace for Labels object")
}
- return buf.Bytes(), nil
+ // Make a copy of the bytes before returning the buffer to the pool
+ result := make([]byte, buf.Len())
+ copy(result, buf.Bytes())
+
+ return result, nil
}
func (ls *Labels) UnmarshalJSON(data []byte) error {