mirror of
https://github.com/hanzoai/traces.git
synced 2026-10-30 18:20:41 +00:00
Christian HaudumandGitHub
4bbb89047b
Signed-off-by: Christian Haudum <[email protected]> Co-authored-by: Susana Ferreira <[email protected]> Fix duplicate metrics registration Signed-off-by: Christian Haudum <[email protected]> Add debug logs to hedge request test Signed-off-by: Christian Haudum <[email protected]> Fix e2e test metric names Remove prometheus/common replace in go.mod Signed-off-by: Christian Haudum <[email protected]>
27 lines
770 B
Go
27 lines
770 B
Go
package grpcclient
|
|
|
|
import (
|
|
"context"
|
|
|
|
"golang.org/x/time/rate"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/codes"
|
|
"google.golang.org/grpc/status"
|
|
)
|
|
|
|
// NewRateLimiter creates a UnaryClientInterceptor for client side rate limiting.
|
|
func NewRateLimiter(cfg *Config) grpc.UnaryClientInterceptor {
|
|
burst := cfg.RateLimitBurst
|
|
if burst == 0 {
|
|
burst = int(cfg.RateLimit)
|
|
}
|
|
limiter := rate.NewLimiter(rate.Limit(cfg.RateLimit), burst)
|
|
return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
|
|
err := limiter.Wait(ctx)
|
|
if err != nil {
|
|
return status.Error(codes.ResourceExhausted, err.Error())
|
|
}
|
|
return invoker(ctx, method, req, reply, cc, opts...)
|
|
}
|
|
}
|