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

config/openim-msggateway.yml

@@ -42,4 +42,4 @@ circuitBreaker:
   window: 3s            # Time window size (seconds)
   bucket: 10            # Number of buckets
   success: 0.6          # Success rate threshold (0.6 means 60%)
-  requestThreshold: 100 # Request threshold; circuit breaker evaluation occurs when reached
+  request: 100 # Request threshold; circuit breaker evaluation occurs when reached

config/openim-msgtransfer.yml

@@ -22,4 +22,4 @@ circuitBreaker:
   window: 3s            # Time window size (seconds)
   bucket: 10            # Number of buckets
   success: 0.6          # Success rate threshold (0.6 means 60%)
-  requestThreshold: 100 # Request threshold; circuit breaker evaluation occurs when reached
+  request: 100 # Request threshold; circuit breaker evaluation occurs when reached

config/openim-push.yml

@@ -25,7 +25,7 @@ circuitBreaker:
   window: 3s            # Time window size (seconds)
   bucket: 10            # Number of buckets
   success: 0.6          # Success rate threshold (0.6 means 60%)
-  requestThreshold: 100 # Request threshold; circuit breaker evaluation occurs when reached
+  request: 100 # Request threshold; circuit breaker evaluation occurs when reached
 
 prometheus:
   # Enable or disable Prometheus monitoring

config/openim-rpc-auth.yml

@@ -36,4 +36,4 @@ circuitBreaker:
   window: 3s            # Time window size (seconds)
   bucket: 10            # Number of buckets
   success: 0.6          # Success rate threshold (0.6 means 60%)
-  requestThreshold: 100 # Request threshold; circuit breaker evaluation occurs when reached
+  request: 100 # Request threshold; circuit breaker evaluation occurs when reached

config/openim-rpc-conversation.yml

@@ -32,4 +32,4 @@ circuitBreaker:
   window: 3s            # Time window size (seconds)
   bucket: 10            # Number of buckets
   success: 0.6          # Success rate threshold (0.6 means 60%)
-  requestThreshold: 100 # Request threshold; circuit breaker evaluation occurs when reached
+  request: 100 # Request threshold; circuit breaker evaluation occurs when reached

config/openim-rpc-friend.yml

@@ -32,4 +32,4 @@ circuitBreaker:
   window: 3s            # Time window size (seconds)
   bucket: 10            # Number of buckets
   success: 0.6          # Success rate threshold (0.6 means 60%)
-  requestThreshold: 100 # Request threshold; circuit breaker evaluation occurs when reached
+  request: 100 # Request threshold; circuit breaker evaluation occurs when reached

config/openim-rpc-group.yml

@@ -35,4 +35,4 @@ circuitBreaker:
   window: 3s            # Time window size (seconds)
   bucket: 10            # Number of buckets
   success: 0.6          # Success rate threshold (0.6 means 60%)
-  requestThreshold: 100 # Request threshold; circuit breaker evaluation occurs when reached
+  request: 100 # Request threshold; circuit breaker evaluation occurs when reached

config/openim-rpc-third.yml

@@ -32,7 +32,7 @@ circuitBreaker:
   window: 3s            # Time window size (seconds)
   bucket: 10            # Number of buckets
   success: 0.6          # Success rate threshold (0.6 means 60%)
-  requestThreshold: 100 # Request threshold; circuit breaker evaluation occurs when reached
+  request: 100 # Request threshold; circuit breaker evaluation occurs when reached
 
 object:
   # Use MinIO as object storage, or set to "cos", "oss", "kodo", "aws", while also configuring the corresponding settings

config/openim-rpc-user.yml

@@ -32,4 +32,4 @@ circuitBreaker:
   window: 3s            # Time window size (seconds)
   bucket: 10            # Number of buckets
   success: 0.6          # Success rate threshold (0.6 means 60%)
-  requestThreshold: 100 # Request threshold; circuit breaker evaluation occurs when reached
+  request: 100 # Request threshold; circuit breaker evaluation occurs when reached

internal/api/ratelimit.go

@@ -1,7 +1,10 @@
 package api
 
 import (
+	"fmt"
+	"math"
 	"net/http"
+	"strconv"
 	"time"
 
 	"github.com/gin-gonic/gin"
@@ -12,7 +15,6 @@ import (
 	"github.com/openimsdk/tools/log"
 )
 
-
 type RateLimiter struct {
 	Enable       bool          `yaml:"enable"`
 	Window       time.Duration `yaml:"window"`       // time duration per window
@@ -34,8 +36,23 @@ func RateLimitMiddleware(config *RateLimiter) gin.HandlerFunc {
 	)
 
 	return func(c *gin.Context) {
+		status := limiter.Stat()
+
+		c.Header("X-BBR-CPU", strconv.FormatInt(status.CPU, 10))
+		c.Header("X-BBR-MinRT", strconv.FormatInt(status.MinRt, 10))
+		c.Header("X-BBR-MaxPass", strconv.FormatInt(status.MaxPass, 10))
+		c.Header("X-BBR-MaxInFlight", strconv.FormatInt(status.MaxInFlight, 10))
+		c.Header("X-BBR-InFlight", strconv.FormatInt(status.InFlight, 10))
+
 		done, err := limiter.Allow()
 		if err != nil {
+
+			c.Header("X-RateLimit-Policy", "BBR")
+			c.Header("Retry-After", calculateBBRRetryAfter(status))
+			c.Header("X-RateLimit-Limit", strconv.FormatInt(status.MaxInFlight, 10))
+			c.Header("X-RateLimit-Remaining", "0") // There is no concept of remaining quota in BBR.
+
+			fmt.Println("rate limited:", err, "path:", c.Request.URL.Path)
 			log.ZWarn(c, "rate limited", err, "path", c.Request.URL.Path)
 			c.AbortWithStatus(http.StatusTooManyRequests)
 			apiresp.GinError(c, errs.NewCodeError(http.StatusTooManyRequests, "too many requests, please try again later"))
@@ -47,3 +64,20 @@ func RateLimitMiddleware(config *RateLimiter) gin.HandlerFunc {
 		done(ratelimit.DoneInfo{})
 	}
 }
+
+func calculateBBRRetryAfter(status bbr.Stat) string {
+	loadRatio := float64(status.CPU) / float64(status.CPU)
+
+	if loadRatio < 0.8 {
+		return "1"
+	}
+	if loadRatio < 0.95 {
+		return "2"
+	}
+
+	backoff := 1 + int64(math.Pow(loadRatio-0.95, 2)*50)
+	if backoff > 5 {
+		backoff = 5
+	}
+	return strconv.FormatInt(backoff, 10)
+}

pkg/common/startrpc/circuitbreaker.go

@@ -49,10 +49,17 @@ func UnaryCircuitBreakerInterceptor(breaker circuitbreaker.CircuitBreaker) grpc.
 		resp, err = handler(ctx, req)
 
 		if err != nil {
-			if st, ok := status.FromError(err); ok && st.Code() == codes.Internal {
-				breaker.MarkFailed()
+			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.MarkSuccess()
+				breaker.MarkFailed()
 			}
 		} else {
 			breaker.MarkSuccess()
@@ -79,10 +86,17 @@ func StreamCircuitBreakerInterceptor(breaker circuitbreaker.CircuitBreaker) grpc
 		err := handler(srv, ss)
 
 		if err != nil {
-			if st, ok := status.FromError(err); ok && st.Code() == codes.Internal {
-				breaker.MarkFailed()
+			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.MarkSuccess()
+				breaker.MarkFailed()
 			}
 		} else {
 			breaker.MarkSuccess()