19 Star 116 Fork 0

kelvins-io / kelvins

加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
克隆/下载
rpc_ratelimit.go 1.57 KB
一键复制 编辑 原始数据 按行查看 历史
雨化田 提交于 2021-10-01 03:07 . 限流
package middleware
import (
"context"
"gitee.com/kelvins-io/common/json"
"gitee.com/kelvins-io/kelvins/util/rpc_helper"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
)
type RPCRateLimitInterceptor struct {
limiter Limiter
}
func NewRPCRateLimitInterceptor(maxConcurrent int) *RPCRateLimitInterceptor {
return &RPCRateLimitInterceptor{
limiter: NewKelvinsRateLimit(maxConcurrent),
}
}
func (r *RPCRateLimitInterceptor) StreamServerInterceptor() grpc.StreamServerInterceptor {
return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
if r.limiter.Limit() {
requestMeta := rpc_helper.GetRequestMetadata(stream.Context())
return status.Errorf(codes.ResourceExhausted, "%s requestMeta:%v is rejected by grpc_ratelimit middleware, please retry later.", info.FullMethod, json.MarshalToStringNoError(requestMeta))
}
defer func() {
r.limiter.ReturnTicket()
}()
return handler(srv, stream)
}
}
func (r *RPCRateLimitInterceptor) UnaryServerInterceptor() grpc.UnaryServerInterceptor {
return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
if r.limiter.Limit() {
requestMeta := rpc_helper.GetRequestMetadata(ctx)
return nil, status.Errorf(codes.ResourceExhausted, "%s requestMeta:%v is rejected by grpc_ratelimit middleware, please retry later.", info.FullMethod, json.MarshalToStringNoError(requestMeta))
}
defer func() {
r.limiter.ReturnTicket()
}()
return handler(ctx, req)
}
}
Go
1
https://gitee.com/kelvins-io/kelvins.git
git@gitee.com:kelvins-io/kelvins.git
kelvins-io
kelvins
kelvins
v1.6.3

搜索帮助

14c37bed 8189591 565d56ea 8189591