github.com/machinefi/w3bstream@v1.6.5-rc9.0.20240426031326-b8c7c4876e72/pkg/depends/conf/rate_limit/rate_limit.go (about)

     1  package rate_limit
     2  
     3  import (
     4  	"context"
     5  
     6  	"go.uber.org/ratelimit"
     7  
     8  	"github.com/machinefi/w3bstream/pkg/depends/base/types"
     9  	"github.com/machinefi/w3bstream/pkg/depends/x/contextx"
    10  	"github.com/machinefi/w3bstream/pkg/depends/x/misc/must"
    11  )
    12  
    13  type RateLimit struct {
    14  	Count    int            `env:""`
    15  	Duration types.Duration `env:""`
    16  
    17  	Limiter ratelimit.Limiter `env:"-"`
    18  }
    19  
    20  func (r *RateLimit) SetDefault() {}
    21  
    22  func (r *RateLimit) Init() {
    23  	if r.Count > 0 && r.Duration > 0 {
    24  		// TODO check Duration
    25  		r.Limiter = ratelimit.New(r.Count, ratelimit.Per(r.Duration.Duration()))
    26  	} else {
    27  		r.Limiter = ratelimit.NewUnlimited()
    28  	}
    29  }
    30  
    31  type rateLimitKey struct{}
    32  
    33  func WithRateLimitKeyContext(rateLimit *RateLimit) func(context.Context) context.Context {
    34  	return func(ctx context.Context) context.Context {
    35  		return contextx.WithValue(ctx, rateLimitKey{}, rateLimit)
    36  	}
    37  }
    38  
    39  func RateLimitKeyFromContext(ctx context.Context) (*RateLimit, bool) {
    40  	j, ok := ctx.Value(rateLimitKey{}).(*RateLimit)
    41  	return j, ok
    42  }
    43  
    44  func MustRateLimitKeyFromContext(ctx context.Context) *RateLimit {
    45  	j, ok := ctx.Value(rateLimitKey{}).(*RateLimit)
    46  	must.BeTrue(ok)
    47  	return j
    48  }