feat: support polaris ratelimit ability (#2586)

pull/2662/head
YuanXin Hu 2 years ago committed by GitHub
parent 7def38acde
commit 0a076443cb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
  1. 2
      contrib/polaris/go.mod
  2. 105
      contrib/polaris/limiter.go
  3. 13
      contrib/polaris/polaris.go
  4. 53
      contrib/polaris/ratelimit.go

@ -3,6 +3,7 @@ module github.com/go-kratos/kratos/contrib/polaris/v2
go 1.18
require (
github.com/go-kratos/aegis v0.1.4
github.com/go-kratos/kratos/v2 v2.5.3
github.com/google/uuid v1.3.0
github.com/polarismesh/polaris-go v1.3.0
@ -13,7 +14,6 @@ require (
github.com/beorn7/perks v1.0.1 // indirect
github.com/cespare/xxhash/v2 v2.1.2 // indirect
github.com/dlclark/regexp2 v1.7.0 // indirect
github.com/go-kratos/aegis v0.1.4 // indirect
github.com/go-playground/form/v4 v4.2.0 // indirect
github.com/golang/protobuf v1.5.2 // indirect
github.com/gorilla/mux v1.8.0 // indirect

@ -0,0 +1,105 @@
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
}

@ -77,3 +77,16 @@ func (p *Polaris) Registry(opts ...RegistryOption) (r *Registry) {
consumer: p.discovery,
}
}
func (p *Polaris) Limiter(opts ...LimiterOption) (r *Limiter) {
op := limiterOptions{
namespace: p.namespace,
}
for _, option := range opts {
option(&op)
}
return &Limiter{
limitAPI: p.limit,
opts: op,
}
}

@ -0,0 +1,53 @@
package polaris
import (
"context"
"strings"
"github.com/go-kratos/aegis/ratelimit"
"github.com/go-kratos/kratos/v2/errors"
"github.com/go-kratos/kratos/v2/middleware"
"github.com/go-kratos/kratos/v2/transport"
"github.com/go-kratos/kratos/v2/transport/http"
"github.com/polarismesh/polaris-go/pkg/model"
)
// ErrLimitExceed is service unavailable due to rate limit exceeded.
var (
ErrLimitExceed = errors.New(429, "RATELIMIT", "service unavailable due to rate limit exceeded")
)
// Ratelimit ratelimiter middleware
func Ratelimit(l Limiter) middleware.Middleware {
return func(handler middleware.Handler) middleware.Handler {
return func(ctx context.Context, req interface{}) (reply interface{}, err error) {
if tr, ok := transport.FromServerContext(ctx); ok {
var args []model.Argument
headers := tr.RequestHeader()
// handle header
for _, header := range headers.Keys() {
args = append(args, model.BuildHeaderArgument(header, headers.Get(header)))
}
// handle http
if ht, ok := tr.(*http.Transport); ok {
// url query
for key, values := range ht.Request().URL.Query() {
args = append(args, model.BuildQueryArgument(key, strings.Join(values, ",")))
}
}
done, e := l.Allow(tr.Operation(), args...)
if e != nil {
// rejected
return nil, ErrLimitExceed
}
// allowed
reply, err = handler(ctx, req)
done(ratelimit.DoneInfo{Err: err})
return
}
return nil, errors.New(400, "Error with transport.FromServerContext", "Error with transport.FromServerContext")
}
}
}
Loading…
Cancel
Save