IM/open-im-server
1
0
Fork 0
mirror of https://github.com/openimsdk/open-im-server.git synced 2026-05-11 12:36:00 +08:00
Code Issues Packages Projects Releases Wiki Activity

​Build: Implement rate limiting and circuit breaker for API and RPC services.​​ (#3572)

Browse Source 
* feat: implement ratelimit and circuitbreaker in middleware.

* ​Build: Implement rate limiting and circuit breaker for API and RPC services.​​

* revert change.

* update ratelimiter and circuitbreaker config.

* update tools to openimsdk tools
This commit is contained in:
Monet Lee
2025-11-05 15:12:06 +08:00
committed by GitHub
parent 3a1c8df3b9
commit 07badb162f
31 changed files with 574 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/common/startrpc/circuitbreaker.go
+107
View File
@@ -0,0 +1,107 @@
package startrpc
import (
"context"
"time"
"github.com/openimsdk/tools/log"
"github.com/openimsdk/tools/stability/circuitbreaker"
"github.com/openimsdk/tools/stability/circuitbreaker/sre"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
)
type CircuitBreaker struct {
Enable bool `yaml:"enable"`
Success float64 `yaml:"success"` // success rate threshold (0.0-1.0)
Request int64 `yaml:"request"` // request threshold
Bucket int `yaml:"bucket"` // number of buckets
Window time.Duration `yaml:"window"` // time window for statistics
}
func NewCircuitBreaker(config *CircuitBreaker) circuitbreaker.CircuitBreaker {
if !config.Enable {
return nil
}
return sre.NewSREBraker(
sre.WithWindow(config.Window),
sre.WithBucket(config.Bucket),
sre.WithSuccess(config.Success),
sre.WithRequest(config.Request),
)
}
func UnaryCircuitBreakerInterceptor(breaker circuitbreaker.CircuitBreaker) grpc.ServerOption {
if breaker == nil {
return grpc.ChainUnaryInterceptor(func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
return handler(ctx, req)
})
}
return grpc.ChainUnaryInterceptor(func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
if err := breaker.Allow(); err != nil {
log.ZWarn(ctx, "rpc circuit breaker open", err, "method", info.FullMethod)
return nil, status.Error(codes.Unavailable, "service unavailable due to circuit breaker")
}
resp, err = handler(ctx, req)
if err != nil {
if st, ok := status.FromError(err); ok {
switch st.Code() {
case codes.OK:
breaker.MarkSuccess()
case codes.InvalidArgument, codes.NotFound, codes.AlreadyExists, codes.PermissionDenied:
breaker.MarkSuccess()
default:
breaker.MarkFailed()
}
} else {
breaker.MarkFailed()
}
} else {
breaker.MarkSuccess()
}
return resp, err
})
}
func StreamCircuitBreakerInterceptor(breaker circuitbreaker.CircuitBreaker) grpc.ServerOption {
if breaker == nil {
return grpc.ChainStreamInterceptor(func(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
return handler(srv, ss)
})
}
return grpc.ChainStreamInterceptor(func(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
if err := breaker.Allow(); err != nil {
log.ZWarn(ss.Context(), "rpc circuit breaker open", err, "method", info.FullMethod)
return status.Error(codes.Unavailable, "service unavailable due to circuit breaker")
}
err := handler(srv, ss)
if err != nil {
if st, ok := status.FromError(err); ok {
switch st.Code() {
case codes.OK:
breaker.MarkSuccess()
case codes.InvalidArgument, codes.NotFound, codes.AlreadyExists, codes.PermissionDenied:
breaker.MarkSuccess()
default:
breaker.MarkFailed()
}
} else {
breaker.MarkFailed()
}
} else {
breaker.MarkSuccess()
}
return err
})
}
pkg/common/startrpc/ratelimit.go
+70
View File
@@ -0,0 +1,70 @@
package startrpc
import (
"context"
"time"
"github.com/openimsdk/tools/log"
"github.com/openimsdk/tools/stability/ratelimit"
"github.com/openimsdk/tools/stability/ratelimit/bbr"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
)
type RateLimiter struct {
Enable bool
Window time.Duration
Bucket int
CPUThreshold int64
}
func NewRateLimiter(config *RateLimiter) ratelimit.Limiter {
if !config.Enable {
return nil
}
return bbr.NewBBRLimiter(
bbr.WithWindow(config.Window),
bbr.WithBucket(config.Bucket),
bbr.WithCPUThreshold(config.CPUThreshold),
)
}
func UnaryRateLimitInterceptor(limiter ratelimit.Limiter) grpc.ServerOption {
if limiter == nil {
return grpc.ChainUnaryInterceptor(func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
return handler(ctx, req)
})
}
return grpc.ChainUnaryInterceptor(func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
done, err := limiter.Allow()
if err != nil {
log.ZWarn(ctx, "rpc rate limited", err, "method", info.FullMethod)
return nil, status.Errorf(codes.ResourceExhausted, "rpc request rate limit exceeded: %v, please try again later", err)
}
defer done(ratelimit.DoneInfo{})
return handler(ctx, req)
})
}
func StreamRateLimitInterceptor(limiter ratelimit.Limiter) grpc.ServerOption {
if limiter == nil {
return grpc.ChainStreamInterceptor(func(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
return handler(srv, ss)
})
}
return grpc.ChainStreamInterceptor(func(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
done, err := limiter.Allow()
if err != nil {
log.ZWarn(ss.Context(), "rpc rate limited", err, "method", info.FullMethod)
return status.Errorf(codes.ResourceExhausted, "rpc request rate limit exceeded: %v, please try again later", err)
}
defer done(ratelimit.DoneInfo{})
return handler(srv, ss)
})
}
pkg/common/startrpc/start.go
+41 -2
View File
@@ -47,7 +47,7 @@ func init() {
prommetrics.RegistryAll()
}
func Start[T any](ctx context.Context, disc *conf.Discovery, prometheusConfig *conf.Prometheus, listenIP,
func Start[T any](ctx context.Context, disc *conf.Discovery, circuitBreakerConfig *conf.CircuitBreaker, rateLimiterConfig *conf.RateLimiter, prometheusConfig *conf.Prometheus, listenIP,
registerIP string, autoSetPorts bool, rpcPorts []int, index int, rpcRegisterName string, notification *conf.Notification, config T,
watchConfigNames []string, watchServiceNames []string,
rpcFn func(ctx context.Context, config T, client discovery.SvcDiscoveryRegistry, server grpc.ServiceRegistrar) error,
@@ -84,6 +84,45 @@ func Start[T any](ctx context.Context, disc *conf.Discovery, prometheusConfig *c
}
}
if circuitBreakerConfig != nil && circuitBreakerConfig.Enable {
cb := &CircuitBreaker{
Enable: circuitBreakerConfig.Enable,
Success: circuitBreakerConfig.Success,
Request: circuitBreakerConfig.Request,
Bucket: circuitBreakerConfig.Bucket,
Window: circuitBreakerConfig.Window,
}
breaker := NewCircuitBreaker(cb)
options = append(options,
UnaryCircuitBreakerInterceptor(breaker),
StreamCircuitBreakerInterceptor(breaker),
)
log.ZInfo(ctx, "RPC circuit breaker enabled",
"service", rpcRegisterName,
"window", circuitBreakerConfig.Window,
"bucket", circuitBreakerConfig.Bucket,
"success", circuitBreakerConfig.Success,
"requestThreshold", circuitBreakerConfig.Request)
}
if rateLimiterConfig != nil && rateLimiterConfig.Enable {
limiter := NewRateLimiter((*RateLimiter)(rateLimiterConfig))
options = append(options,
UnaryRateLimitInterceptor(limiter),
StreamRateLimitInterceptor(limiter),
)
log.ZInfo(ctx, "RPC rate limiter enabled",
"service", rpcRegisterName,
"window", rateLimiterConfig.Window,
"bucket", rateLimiterConfig.Bucket,
"cpuThreshold", rateLimiterConfig.CPUThreshold)
}
registerIP, err := network.GetRpcRegisterIP(registerIP)
if err != nil {
return err
@@ -123,7 +162,7 @@ func Start[T any](ctx context.Context, disc *conf.Discovery, prometheusConfig *c
go func() {
sigs := make(chan os.Signal, 1)
signal.Notify(sigs, syscall.SIGTERM, syscall.SIGINT, syscall.SIGKILL)
signal.Notify(sigs, syscall.SIGTERM, syscall.SIGINT)
select {
case <-ctx.Done():
return