[feat][evaluation] add plain limiter factory (#297)

* feat: add evaluator-level rate limit to RunEvaluator

* feat: add evaluator-level rate limit to RunEvaluator

* fix trace

* feat: add evaluator-level rate limit to RunEvaluator

* fix: fix the authBuiltinManagement R/W auth issue

* feat: add plain limiter factory

* feat: add more condition protection for plain limiter factory

---------

Co-authored-by: tpfz <[email protected]>

Author: HearyShen

backend/api/api.go

@@ -10,9 +10,10 @@ import (
 	"context"
 
 	"github.com/cloudwego/kitex/pkg/endpoint"
+	"github.com/google/wire"
+
 	"github.com/coze-dev/coze-loop/backend/kitex_gen/coze/loop/evaluation/experimentservice"
 	task_processor "github.com/coze-dev/coze-loop/backend/modules/observability/domain/task/service/taskexe/processor"
-	"github.com/google/wire"
 
 	"github.com/coze-dev/coze-loop/backend/infra/ck"
 	"github.com/coze-dev/coze-loop/backend/infra/db"
@@ -169,6 +170,7 @@ func InitEvaluationHandler(
 	fileClient fileservice.Client,
 	tagClient tagservice.Client,
 	objectStorage fileserver.ObjectStorage,
+	plainLimiterFactory limiter.IPlainRateLimiterFactory,
 ) (*EvaluationHandler, error) {
 	wire.Build(
 		evaluationSet,

backend/api/handler/coze/loop/apis/wire.go

@@ -52,7 +52,7 @@ func main() {
 		panic(err)
 	}
 
-	handler, err := api.Init(ctx, c.idgen, c.db, c.redis, c.cfgFactory, c.mqFactory, c.objectStorage, c.batchObjectStorage, c.benefitSvc, c.auditClient, c.metric, c.limiterFactory, c.ckDb, c.translater)
+	handler, err := api.Init(ctx, c.idgen, c.db, c.redis, c.cfgFactory, c.mqFactory, c.objectStorage, c.batchObjectStorage, c.benefitSvc, c.auditClient, c.metric, c.limiterFactory, c.ckDb, c.translater, c.plainLimiterFactory)
 	if err != nil {
 		panic(err)
 	}
@@ -118,19 +118,20 @@ func getComponentConfig(configFactory conf.IConfigLoaderFactory) (*ComponentConf
 }
 
 type component struct {
-	idgen              idgen.IIDGenerator
-	db                 db.Provider
-	redis              redis.Cmdable
-	cfgFactory         conf.IConfigLoaderFactory
-	mqFactory          mq.IFactory
-	objectStorage      fileserver.ObjectStorage
-	batchObjectStorage fileserver.BatchObjectStorage
-	benefitSvc         benefit.IBenefitService
-	auditClient        audit.IAuditService
-	metric             metrics.Meter
-	limiterFactory     limiter.IRateLimiterFactory
-	ckDb               ck.Provider
-	translater         i18n.ITranslater
+	idgen               idgen.IIDGenerator
+	db                  db.Provider
+	redis               redis.Cmdable
+	cfgFactory          conf.IConfigLoaderFactory
+	mqFactory           mq.IFactory
+	objectStorage       fileserver.ObjectStorage
+	batchObjectStorage  fileserver.BatchObjectStorage
+	benefitSvc          benefit.IBenefitService
+	auditClient         audit.IAuditService
+	metric              metrics.Meter
+	limiterFactory      limiter.IRateLimiterFactory
+	ckDb                ck.Provider
+	translater          i18n.ITranslater
+	plainLimiterFactory limiter.IPlainRateLimiterFactory
 }
 
 func initTracer(handler *apis.APIHandler) error {
@@ -250,19 +251,20 @@ func newComponent(ctx context.Context) (*component, error) {
 	}
 
 	return &component{
-		idgen:              idgenerator,
-		db:                 db,
-		redis:              cmdable,
-		cfgFactory:         cfgFactory,
-		mqFactory:          rocketmq.NewFactory(),
-		objectStorage:      objectStorage,
-		batchObjectStorage: objectStorage,
-		benefitSvc:         benefit.NewNoopBenefitService(),
-		auditClient:        audit.NewNoopAuditService(),
-		metric:             metrics.GetMeter(),
-		limiterFactory:     dist.NewRateLimiterFactory(cmdable),
-		ckDb:               ckDb,
-		translater:         translater,
+		idgen:               idgenerator,
+		db:                  db,
+		redis:               cmdable,
+		cfgFactory:          cfgFactory,
+		mqFactory:           rocketmq.NewFactory(),
+		objectStorage:       objectStorage,
+		batchObjectStorage:  objectStorage,
+		benefitSvc:          benefit.NewNoopBenefitService(),
+		auditClient:         audit.NewNoopAuditService(),
+		metric:              metrics.GetMeter(),
+		limiterFactory:      dist.NewRateLimiterFactory(cmdable),
+		ckDb:                ckDb,
+		translater:          translater,
+		plainLimiterFactory: dist.NewPlainLimiterFactory(cmdable),
 	}, nil
 }
 

backend/api/handler/coze/loop/apis/wire_gen.go

@@ -36,3 +36,9 @@ func WithExprCacheSize(size int) FactoryOpt {
 type factoryOpt struct {
 	exprCacheSize int
 }
+
+func NewPlainLimiterFactory(cmdable redis.Cmdable) limiter.IPlainRateLimiterFactory {
+	return &plainLimiterFactory{
+		cmdable: cmdable,
+	}
+}

backend/cmd/main.go

@@ -0,0 +1,42 @@
+package dist
+
+import (
+	"fmt"
+
+	"github.com/go-redis/redis_rate/v10"
+
+	"github.com/coze-dev/coze-loop/backend/infra/limiter"
+	"github.com/coze-dev/coze-loop/backend/infra/redis"
+	"github.com/coze-dev/coze-loop/backend/pkg/logs"
+)
+
+type plainLimiterFactory struct {
+	cmdable redis.Cmdable
+}
+
+func (f *plainLimiterFactory) NewPlainRateLimiter(opts ...limiter.FactoryOptionFn) limiter.IPlainRateLimiter {
+	opt := &limiter.FactoryOption{}
+	for _, fn := range opts {
+		fn(opt)
+	}
+
+	rawRedis, ok := redis.Unwrap(f.cmdable)
+	if !ok {
+		panic(fmt.Errorf("redis cmdable must be unwrappable"))
+	}
+
+	rl := &rateLimiter{
+		rules:   make([]*rule, 0, len(opt.Rules)),
+		limiter: redis_rate.NewLimiter(rawRedis),
+	}
+
+	for _, r := range opt.Rules {
+		if rr, err := rl.newRule(r); err != nil {
+			logs.Error("rate limiter set rule failed, rule: %v, err: %v", r, err)
+		} else {
+			rl.addRule(rr)
+		}
+	}
+
+	return rl
+}

backend/infra/limiter/dist/dist.go

@@ -79,3 +79,13 @@ type Limit struct {
 	Burst  int           `json:"burst" yaml:"burst" mapstructure:"burst"`
 	Period time.Duration `json:"period" yaml:"period" mapstructure:"period"`
 }
+
+//go:generate mockgen -destination=mocks/rate_limiter_factory.go -package=mocks . IPlainRateLimiterFactory
+type IPlainRateLimiterFactory interface {
+	NewPlainRateLimiter(opts ...FactoryOptionFn) IPlainRateLimiter
+}
+
+//go:generate mockgen -destination=mocks/rate_limiter.go -package=mocks . IPlainRateLimiter
+type IPlainRateLimiter interface {
+	AllowN(ctx context.Context, key string, n int, opts ...LimitOptionFn) (*Result, error)
+}