add rate limit

Author: xixingde

.gitignore

@@ -24,10 +24,10 @@
 .idea
 main
 main.exe
-.codebase_index/
 *.iml
 *.scip
 bin/*
 .roo
 docs/task.md
-testfiles
+testfiles
+cmd.sh

deploy/deployment.yaml

@@ -87,6 +87,9 @@ data:
       log_level: "info"
       max_concurrency: 5
       skip_patterns: []
+    TokenLimit:
+      max_running_tasks: 1  
+      enabled: true    
 ---
 apiVersion: apps/v1
 kind: Deployment

docs/charts/gettoken.md

@@ -0,0 +1,9 @@
+
+```mermaid
+graph TD
+    A[生成Token接口] --> B[读取限流配置文件]
+    B --> C[查询任务池正运行任务]
+    C --> D{到达限流配置?}
+    D -->|是| E[生成失败]
+    D -->|否| F[根据ClientId生成Token]
+```

etc/conf.yaml

@@ -89,3 +89,7 @@ Validation:
   log_level: "info"
   max_concurrency: 5
   skip_patterns: []
+
+TokenLimit:
+  max_running_tasks: 10
+  enabled: true

internal/config/config.go

@@ -17,6 +17,13 @@ type Config struct {
 	VectorStore VectorStoreConf
 	Cleaner     CleanerConf
 	Validation  ValidationConfig
+	TokenLimit  TokenLimitConf
+}
+
+// TokenLimitConf token限流配置
+type TokenLimitConf struct {
+	MaxRunningTasks int  `json:"max_running_tasks" yaml:"max_running_tasks"`
+	Enabled         bool `json:"enabled" yaml:"enabled"`
 }
 
 // Validate 实现 Validator 接口

internal/handler/token.go

@@ -1,6 +1,7 @@
 package handler
 
 import (
+	"errors"
 	"net/http"
 
 	"github.com/zgsm-ai/codebase-indexer/internal/logic"
@@ -33,11 +34,16 @@ func (h *tokenHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// 创建token逻辑
-	tokenLogic := logic.NewTokenLogic(r.Context())
+	tokenLogic := logic.NewTokenLogic(r.Context(), h.svcCtx)
 	tokenResp, err := tokenLogic.GenerateToken(&req)
 	if err != nil {
+		// 检查是否为限流错误
+		if errors.Is(err, types.ErrRateLimitReached) {
+			response.RateLimit(w, err)
+			return
+		}
 		response.Error(w, err)
 		return
 	}
 	response.Json(w, tokenResp)
-}
+}

internal/logic/token.go

@@ -5,18 +5,22 @@ import (
 	"errors"
 	"fmt"
 	"time"
+
+	"github.com/zgsm-ai/codebase-indexer/internal/svc"
 	"github.com/zgsm-ai/codebase-indexer/internal/types"
 )
 
 // TokenLogic token生成逻辑
 type TokenLogic struct {
 	ctx    context.Context
+	svcCtx *svc.ServiceContext
 }
 
 // NewTokenLogic 创建TokenLogic实例
-func NewTokenLogic(ctx context.Context) *TokenLogic {
+func NewTokenLogic(ctx context.Context, svcCtx *svc.ServiceContext) *TokenLogic {
 	return &TokenLogic{
 		ctx:    ctx,
+		svcCtx: svcCtx,
 	}
 }
 
@@ -26,6 +30,31 @@ func (l *TokenLogic) GenerateToken(req *types.TokenRequest) (*types.TokenRespons
 		return nil, fmt.Errorf("invalid request: %w", err)
 	}
 
+	// 1. 读取限流配置文件
+	tokenLimit := l.svcCtx.Config.TokenLimit
+	if !tokenLimit.Enabled {
+		// 限流未启用，直接生成token
+		return l.generateToken(req)
+	}
+
+	// 2. 查询任务池正运行任务
+	runningTasks, err := l.getRunningTasksCount()
+	if err != nil {
+		return nil, fmt.Errorf("查询运行中任务失败: %w", err)
+	}
+
+	// 3. 判断是否到达限流配置
+	if runningTasks >= tokenLimit.MaxRunningTasks {
+		// 4. 生成失败
+		return nil, types.ErrRateLimitReached
+	}
+
+	// 5. 根据ClientId生成Token
+	return l.generateToken(req)
+}
+
+// generateToken 生成token
+func (l *TokenLogic) generateToken(req *types.TokenRequest) (*types.TokenResponseData, error) {
 	// 使用clientId和codebasePath生成token
 	// 这里使用简单的哈希组合，实际生产环境应使用更安全的JWT实现
 	token := fmt.Sprintf("%s_%s_%s", req.ClientId, req.CodebasePath, l.generateRandomString(16))
@@ -37,6 +66,16 @@ func (l *TokenLogic) GenerateToken(req *types.TokenRequest) (*types.TokenRespons
 	}, nil
 }
 
+// getRunningTasksCount 获取运行中任务数量
+func (l *TokenLogic) getRunningTasksCount() (int, error) {
+	// 获取任务池中正在运行的任务数量
+	if l.svcCtx.TaskPool == nil {
+		return 0, fmt.Errorf("任务池未初始化")
+	}
+
+	return l.svcCtx.TaskPool.Running(), nil
+}
+
 // validateRequest 验证请求参数
 func (l *TokenLogic) validateRequest(req *types.TokenRequest) error {
 	if req.ClientId == "" {
@@ -51,27 +90,16 @@ func (l *TokenLogic) validateRequest(req *types.TokenRequest) error {
 	return nil
 }
 
-// generateJTI 生成唯一的JWT ID
-func (l *TokenLogic) generateJTI() string {
-	return fmt.Sprintf("%d-%s", time.Now().UnixNano(), l.generateRandomString(8))
-}
-
 // generateRandomString 生成随机字符串
 func (l *TokenLogic) generateRandomString(length int) string {
 	const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
 	result := make([]byte, length)
-	
+
 	// 使用更稳定的随机源
 	seed := time.Now().UnixNano()
 	for i := range result {
-		seed = (seed * 1103515245 + 12345) & 0x7fffffff
+		seed = (seed*1103515245 + 12345) & 0x7fffffff
 		result[i] = charset[seed%int64(len(charset))]
 	}
 	return string(result)
 }
-
-// getSecretKey 获取JWT签名密钥
-func (l *TokenLogic) getSecretKey() string {
-	// 默认密钥（仅用于开发环境）
-	return "default-secret-key-change-in-production"
-}

internal/response/code_msg.go

@@ -30,3 +30,8 @@ func NewAuthError(msg string) error {
 func NewPermissionError(msg string) error {
 	return &codeMsg{Code: 403, Message: msg}
 }
+
+// NewRateLimitError creates a new rate limit error.
+func NewRateLimitError(msg string) error {
+	return &codeMsg{Code: 429, Message: msg}
+}

internal/response/resp.go

@@ -2,9 +2,10 @@ package response
 
 import (
 	"context"
+	"net/http"
+
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/rest/httpx"
-	"net/http"
 )
 
 const (
@@ -31,6 +32,11 @@ func Error(w http.ResponseWriter, e error) {
 	httpx.WriteJson(w, http.StatusBadRequest, wrapResponse(e)) // TODO 500会触发go-zero breaker
 }
 
+func RateLimit(w http.ResponseWriter, e error) {
+	logx.WithCallerSkip(2).Errorf("rate limit error: %v", e)
+	httpx.WriteJson(w, http.StatusTooManyRequests, wrapResponse(e))
+}
+
 func Bytes(w http.ResponseWriter, v []byte) {
 	w.WriteHeader(http.StatusOK)
 	_, _ = w.Write(v)

internal/types/errors.go

@@ -1 +1,8 @@
 package types
+
+import "errors"
+
+var (
+	// ErrRateLimitReached 限流达到上限错误
+	ErrRateLimitReached = errors.New("The system is busy. Please try again later (maximum number of concurrent tasks reached).")
+)