kratos/contrib/polaris/limiter.go

package polaris

import (
	"time"

	"github.com/go-kratos/aegis/ratelimit"

	"github.com/polarismesh/polaris-go"
	"github.com/polarismesh/polaris-go/pkg/model"
)

type (
	// LimiterOption function for polaris limiter
	LimiterOption func(*limiterOptions)
)

type limiterOptions struct {
	// required, polaris limit namespace
	namespace string

	// required, polaris limit service name
	service string

	// optional, polaris limit request timeout
	// max value is (1+RetryCount) * Timeout
	timeout time.Duration

	// optional, polaris limit retryCount
	// init by polaris config
	retryCount int

	// optional, request limit quota
	token uint32
}

// WithLimiterNamespace with limiter namespace.
func WithLimiterNamespace(namespace string) LimiterOption {
	return func(o *limiterOptions) {
		o.namespace = namespace
	}
}

// WithLimiterService with limiter service.
func WithLimiterService(service string) LimiterOption {
	return func(o *limiterOptions) {
		o.service = service
	}
}

// WithLimiterTimeout with limiter arguments.
func WithLimiterTimeout(timeout time.Duration) LimiterOption {
	return func(o *limiterOptions) {
		o.timeout = timeout
	}
}

// WithLimiterRetryCount with limiter retryCount.
func WithLimiterRetryCount(retryCount int) LimiterOption {
	return func(o *limiterOptions) {
		o.retryCount = retryCount
	}
}

// WithLimiterToken with limiter token.
func WithLimiterToken(token uint32) LimiterOption {
	return func(o *limiterOptions) {
		o.token = token
	}
}

type Limiter struct {
	// polaris limit api
	limitAPI polaris.LimitAPI

	opts limiterOptions
}

// init quotaRequest
func buildRequest(opts limiterOptions) polaris.QuotaRequest {
	quotaRequest := polaris.NewQuotaRequest()
	quotaRequest.SetNamespace(opts.namespace)
	quotaRequest.SetRetryCount(opts.retryCount)
	quotaRequest.SetService(opts.service)
	quotaRequest.SetTimeout(opts.timeout)
	quotaRequest.SetToken(opts.token)
	return quotaRequest
}

// Allow interface impl
func (l *Limiter) Allow(method string, argument ...model.Argument) (ratelimit.DoneFunc, error) {
	request := buildRequest(l.opts)
	request.SetMethod(method)
	for _, arg := range argument {
		request.AddArgument(arg)
	}
	resp, err := l.limitAPI.GetQuota(request)
	if err != nil {
		// ignore err
		return func(ratelimit.DoneInfo) {}, nil
	}
	if resp.Get().Code == model.QuotaResultOk {
		return func(ratelimit.DoneInfo) {}, nil
	}
	return nil, ratelimit.ErrLimitExceed
}