github.com/songzhibin97/gkit@v1.2.13/overload/bbr/middle.go (about)

     1  package bbr
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/songzhibin97/gkit/middleware"
     7  	"github.com/songzhibin97/gkit/options"
     8  	"github.com/songzhibin97/gkit/overload"
     9  )
    10  
    11  const (
    12  	LimitKey = "LimitKey"
    13  	LimitOp  = "LimitLoad"
    14  )
    15  
    16  func NewLimiter(options ...options.Option) middleware.MiddleWare {
    17  	g := NewGroup(options...)
    18  	return func(next middleware.Endpoint) middleware.Endpoint {
    19  		return func(ctx context.Context, i interface{}) (interface{}, error) {
    20  			// 通过ctx 获取 g中的限制器
    21  			defaultKey := "default"
    22  			defaultOp := overload.Success
    23  			if v := ctx.Value(LimitKey); v != nil {
    24  				defaultKey = v.(string)
    25  			}
    26  			if v := ctx.Value(LimitOp); v != nil {
    27  				defaultOp = v.(overload.Op)
    28  			}
    29  			limiter := g.Get(defaultKey)
    30  			if f, err := limiter.Allow(ctx); err != nil {
    31  				return nil, err
    32  			} else {
    33  				f(overload.DoneInfo{Op: defaultOp})
    34  				return next(ctx, i)
    35  			}
    36  		}
    37  	}
    38  }