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

config/openim-api.yml

@@ -17,3 +17,13 @@ prometheus:
   ports:
   # This address can be accessed via a browser
   grafanaURL:
+
+ratelimiter:
+  # Whether to enable rate limiting
+  enable: false
+  # WindowSize defines time duration per window
+  window: 20s
+  # BucketNum defines bucket number for each window
+  bucket: 500
+  # CPU threshold; valid range 0–1000 (1000 = 100%)
+  cpuThreshold: 850

config/openim-msggateway.yml

@@ -26,3 +26,20 @@ longConnSvr:
   websocketMaxMsgLen: 4096
   # WebSocket connection handshake timeout in seconds
   websocketTimeout: 10
+
+ratelimiter:
+  # Whether to enable rate limiting
+  enable: false
+  # WindowSize defines time duration per window
+  window: 20s
+  # BucketNum defines bucket number for each window
+  bucket: 500
+  # CPU threshold; valid range 0–1000 (1000 = 100%)
+  cpuThreshold: 850
+
+circuitBreaker:
+  enable: false
+  window: 5s            # Time window size (seconds)
+  bucket: 100            # Number of buckets
+  success: 0.6          # Success rate threshold (0.6 means 60%)
+  request: 500 # Request threshold; circuit breaker evaluation occurs when reached

config/openim-msgtransfer.yml

@@ -6,3 +6,20 @@ prometheus:
   # List of ports that Prometheus listens on; each port corresponds to an instance of monitoring. Ensure these are managed accordingly
   # It will only take effect when autoSetPorts is set to false.
   ports:
+
+ratelimiter:
+  # Whether to enable rate limiting
+  enable: false
+  # WindowSize defines time duration per window
+  window: 20s
+  # BucketNum defines bucket number for each window
+  bucket: 500
+  # CPU threshold; valid range 0–1000 (1000 = 100%)
+  cpuThreshold: 850
+
+circuitBreaker:
+  enable: false
+  window: 5s            # Time window size (seconds)
+  bucket: 100            # Number of buckets
+  success: 0.6          # Success rate threshold (0.6 means 60%)
+  request: 500 # Request threshold; circuit breaker evaluation occurs when reached

config/openim-push.yml

@@ -10,10 +10,26 @@ rpc:
   # It will only take effect when autoSetPorts is set to false.
   ports:
 
+ratelimiter:
+  # Whether to enable rate limiting
+  enable: false
+  # WindowSize defines time duration per window
+  window: 20s
+  # BucketNum defines bucket number for each window
+  bucket: 500
+  # CPU threshold; valid range 0–1000 (1000 = 100%)
+  cpuThreshold: 850
+
+circuitBreaker:
+  enable: false
+  window: 5s            # Time window size (seconds)
+  bucket: 100            # Number of buckets
+  success: 0.6          # Success rate threshold (0.6 means 60%)
+  request: 500 # Request threshold; circuit breaker evaluation occurs when reached
 
 prometheus:
   # Enable or disable Prometheus monitoring
-  enable: true
+  enable: false
   # List of ports that Prometheus listens on; these must match the number of rpc.ports to ensure correct monitoring setup
   # It will only take effect when autoSetPorts is set to false.
   ports:

config/openim-rpc-auth.yml

@@ -20,3 +20,20 @@ prometheus:
 tokenPolicy:
   # Token validity period, in days
   expire: 90
+
+ratelimiter:
+  # Whether to enable rate limiting
+  enable: false
+  # WindowSize defines time duration per window
+  window: 20s
+  # BucketNum defines bucket number for each window
+  bucket: 500
+  # CPU threshold; valid range 0–1000 (1000 = 100%)
+  cpuThreshold: 850
+
+circuitBreaker:
+  enable: false
+  window: 5s            # Time window size (seconds)
+  bucket: 100            # Number of buckets
+  success: 0.6          # Success rate threshold (0.6 means 60%)
+  request: 500 # Request threshold; circuit breaker evaluation occurs when reached

config/openim-rpc-conversation.yml

@@ -16,3 +16,20 @@ prometheus:
   # List of ports that Prometheus listens on; these must match the number of rpc.ports to ensure correct monitoring setup
   # It will only take effect when autoSetPorts is set to false.
   ports:
+
+ratelimiter:
+  # Whether to enable rate limiting
+  enable: false
+  # WindowSize defines time duration per window
+  window: 20s
+  # BucketNum defines bucket number for each window
+  bucket: 500
+  # CPU threshold; valid range 0–1000 (1000 = 100%)
+  cpuThreshold: 850
+
+circuitBreaker:
+  enable: false
+  window: 5s            # Time window size (seconds)
+  bucket: 100            # Number of buckets
+  success: 0.6          # Success rate threshold (0.6 means 60%)
+  request: 500 # Request threshold; circuit breaker evaluation occurs when reached

config/openim-rpc-friend.yml

@@ -16,3 +16,20 @@ prometheus:
   # List of ports that Prometheus listens on; these must match the number of rpc.ports to ensure correct monitoring setup
   # It will only take effect when autoSetPorts is set to false.
   ports:
+
+ratelimiter:
+  # Whether to enable rate limiting
+  enable: false
+  # WindowSize defines time duration per window
+  window: 20s
+  # BucketNum defines bucket number for each window
+  bucket: 500
+  # CPU threshold; valid range 0–1000 (1000 = 100%)
+  cpuThreshold: 850
+
+circuitBreaker:
+  enable: false
+  window: 5s            # Time window size (seconds)
+  bucket: 100            # Number of buckets
+  success: 0.6          # Success rate threshold (0.6 means 60%)
+  request: 500 # Request threshold; circuit breaker evaluation occurs when reached

config/openim-rpc-group.yml

@@ -19,3 +19,20 @@ prometheus:
 
 
 enableHistoryForNewMembers: true
+
+ratelimiter:
+  # Whether to enable rate limiting
+  enable: false
+  # WindowSize defines time duration per window
+  window: 20s
+  # BucketNum defines bucket number for each window
+  bucket: 500
+  # CPU threshold; valid range 0–1000 (1000 = 100%)
+  cpuThreshold: 850
+
+circuitBreaker:
+  enable: false
+  window: 5s            # Time window size (seconds)
+  bucket: 100            # Number of buckets
+  success: 0.6          # Success rate threshold (0.6 means 60%)
+  request: 500 # Request threshold; circuit breaker evaluation occurs when reached

config/openim-rpc-msg.yml

@@ -20,3 +20,20 @@ prometheus:
 
 # Does sending messages require friend verification
 friendVerify: false
+
+ratelimiter:
+  # Whether to enable rate limiting
+  enable: false
+  # WindowSize defines time duration per window
+  window: 20s
+  # BucketNum defines bucket number for each window
+  bucket: 500
+  # CPU threshold; valid range 0–1000 (1000 = 100%)
+  cpuThreshold: 850
+
+circuitBreaker:
+  enable: false
+  window: 5s            # Time window size (seconds)
+  bucket: 100            # Number of buckets
+  success: 0.6          # Success rate threshold (0.6 means 60%)
+  request: 500 # Request threshold; circuit breaker evaluation occurs when reached

config/openim-rpc-third.yml

@@ -17,6 +17,22 @@ prometheus:
   # It will only take effect when autoSetPorts is set to false.
   ports:
 
+ratelimiter:
+  # Whether to enable rate limiting
+  enable: false
+  # WindowSize defines time duration per window
+  window: 20s
+  # BucketNum defines bucket number for each window
+  bucket: 500
+  # CPU threshold; valid range 0–1000 (1000 = 100%)
+  cpuThreshold: 850
+
+circuitBreaker:
+  enable: false
+  window: 5s            # Time window size (seconds)
+  bucket: 100            # Number of buckets
+  success: 0.6          # Success rate threshold (0.6 means 60%)
+  request: 500 # 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

@@ -16,3 +16,20 @@ prometheus:
   # Prometheus listening ports, must be consistent with the number of rpc.ports
   # It will only take effect when autoSetPorts is set to false.
   ports:
+
+ratelimiter:
+  # Whether to enable rate limiting
+  enable: false
+  # WindowSize defines time duration per window
+  window: 20s
+  # BucketNum defines bucket number for each window
+  bucket: 500
+  # CPU threshold; valid range 0–1000 (1000 = 100%)
+  cpuThreshold: 850
+
+circuitBreaker:
+  enable: false
+  window: 5s            # Time window size (seconds)
+  bucket: 100            # Number of buckets
+  success: 0.6          # Success rate threshold (0.6 means 60%)
+  request: 500 # Request threshold; circuit breaker evaluation occurs when reached

go.mod

@@ -135,6 +135,7 @@ require (
 	github.com/leodido/go-urn v1.4.0 // indirect
 	github.com/lestrrat-go/strftime v1.0.6 // indirect
 	github.com/lithammer/shortuuid v3.0.0+incompatible // indirect
+	github.com/lufia/plan9stats v0.0.0-20230110061619-bbe2e5e100de // indirect
 	github.com/magefile/mage v1.15.0 // indirect
 	github.com/magiconair/properties v1.8.7 // indirect
 	github.com/mailru/easyjson v0.7.7 // indirect
@@ -151,6 +152,7 @@ require (
 	github.com/pelletier/go-toml/v2 v2.2.2 // indirect
 	github.com/pierrec/lz4/v4 v4.1.21 // indirect
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
+	github.com/power-devops/perfstat v0.0.0-20221212215047-62379fc7944b // indirect
 	github.com/prometheus/client_model v0.5.0 // indirect
 	github.com/prometheus/common v0.45.0 // indirect
 	github.com/prometheus/procfs v0.12.0 // indirect
@@ -160,6 +162,8 @@ require (
 	github.com/sagikazarmark/locafero v0.4.0 // indirect
 	github.com/sagikazarmark/slog-shim v0.1.0 // indirect
 	github.com/sercand/kuberesolver/v6 v6.0.1 // indirect
+	github.com/shirou/gopsutil/v3 v3.24.5 // indirect
+	github.com/shoenig/go-m1cpu v0.1.6 // indirect
 	github.com/sourcegraph/conc v0.3.0 // indirect
 	github.com/spf13/afero v1.11.0 // indirect
 	github.com/spf13/cast v1.6.0 // indirect

go.sum

@@ -303,6 +303,8 @@ github.com/likexian/gokit v0.25.13 h1:p2Uw3+6fGG53CwdU2Dz0T6bOycdb2+bAFAa3ymwWVk
 github.com/likexian/gokit v0.25.13/go.mod h1:qQhEWFBEfqLCO3/vOEo2EDKd+EycekVtUK4tex+l2H4=
 github.com/lithammer/shortuuid v3.0.0+incompatible h1:NcD0xWW/MZYXEHa6ITy6kaXN5nwm/V115vj2YXfhS0w=
 github.com/lithammer/shortuuid v3.0.0+incompatible/go.mod h1:FR74pbAuElzOUuenUHTK2Tciko1/vKuIKS9dSkDrA4w=
+github.com/lufia/plan9stats v0.0.0-20230110061619-bbe2e5e100de h1:V53FWzU6KAZVi1tPp5UIsMoUWJ2/PNwYIDXnu7QuBCE=
+github.com/lufia/plan9stats v0.0.0-20230110061619-bbe2e5e100de/go.mod h1:JKx41uQRwqlTZabZc+kILPrO/3jlKnQ2Z8b7YiVw5cE=
 github.com/magefile/mage v1.15.0 h1:BvGheCMAsG3bWUDbZ8AyXXpCNwU9u5CB6sM+HNb9HYg=
 github.com/magefile/mage v1.15.0/go.mod h1:z5UZb/iS3GoOSn0JgWuiw7dxlurVYTu+/jHXqQg881A=
 github.com/magiconair/properties v1.8.7 h1:IeQXZAiQcpL9mgcAe1Nu6cX9LLw6ExEHKjN0VQdvPDY=
@@ -361,6 +363,8 @@ github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINE
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=
 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/power-devops/perfstat v0.0.0-20221212215047-62379fc7944b h1:0LFwY6Q3gMACTjAbMZBjXAqTOzOwFaj2Ld6cjeQ7Rig=
+github.com/power-devops/perfstat v0.0.0-20221212215047-62379fc7944b/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE=
 github.com/prashantv/gostub v1.1.0 h1:BTyx3RfQjRHnUWaGF9oQos79AlQ5k8WNktv7VGvVH4g=
 github.com/prashantv/gostub v1.1.0/go.mod h1:A5zLQHz7ieHGG7is6LLXLz7I8+3LZzsrV0P1IAHhP5U=
 github.com/prometheus/client_golang v1.18.0 h1:HzFfmkOzH5Q8L8G+kSJKUx5dtG87sewO+FoDDqP5Tbk=
@@ -397,6 +401,12 @@ github.com/sercand/kuberesolver/v6 v6.0.1 h1:XZUTA0gy/lgDYp/UhEwv7Js24F1j8NJ833Q
 github.com/sercand/kuberesolver/v6 v6.0.1/go.mod h1:C0tsTuRMONSY+Xf7pv7RMW1/JlewY1+wS8SZE+1lf1s=
 github.com/shirou/gopsutil v3.21.11+incompatible h1:+1+c1VGhc88SSonWP6foOcLhvnKlUeu/erjjvaPEYiI=
 github.com/shirou/gopsutil v3.21.11+incompatible/go.mod h1:5b4v6he4MtMOwMlS0TUMTu2PcXUg8+E1lC7eC3UO/RA=
+github.com/shirou/gopsutil/v3 v3.24.5 h1:i0t8kL+kQTvpAYToeuiVk3TgDeKOFioZO3Ztz/iZ9pI=
+github.com/shirou/gopsutil/v3 v3.24.5/go.mod h1:bsoOS1aStSs9ErQ1WWfxllSeS1K5D+U30r2NfcubMVk=
+github.com/shoenig/go-m1cpu v0.1.6 h1:nxdKQNcEB6vzgA2E2bvzKIYRuNj7XNJ4S/aRSwKzFtM=
+github.com/shoenig/go-m1cpu v0.1.6/go.mod h1:1JJMcUBvfNwpq05QDQVAnx3gUHr9IYF7GNg9SUEw2VQ=
+github.com/shoenig/test v0.6.4 h1:kVTaSd7WLz5WZ2IaoM0RSzRsUD+m8wRR+5qvntpn4LU=
+github.com/shoenig/test v0.6.4/go.mod h1:byHiCGXqrVaflBLAMq/srcZIHynQPQgeyvkvXnjqq0k=
 github.com/sourcegraph/conc v0.3.0 h1:OQTbbt6P72L20UqAkXXuLOj79LfEanQ+YQFNpLA9ySo=
 github.com/sourcegraph/conc v0.3.0/go.mod h1:Sdozi7LEKbFPqYX2/J+iBAM6HpqSLTASQIKqDmF7Mt0=
 github.com/spf13/afero v1.11.0 h1:WJQKhtpdm3v2IzqG8VMqrr6Rf3UYpEF239Jy9wNepM8=
@@ -548,6 +558,7 @@ golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20201204225414-ed752295db88/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=

internal/api/ratelimit.go

@@ -0,0 +1,83 @@
+package api
+
+import (
+	"fmt"
+	"math"
+	"net/http"
+	"strconv"
+	"time"
+
+	"github.com/gin-gonic/gin"
+
+	"github.com/openimsdk/tools/apiresp"
+	"github.com/openimsdk/tools/errs"
+	"github.com/openimsdk/tools/log"
+	"github.com/openimsdk/tools/stability/ratelimit"
+	"github.com/openimsdk/tools/stability/ratelimit/bbr"
+)
+
+type RateLimiter struct {
+	Enable       bool          `yaml:"enable"`
+	Window       time.Duration `yaml:"window"`       // time duration per window
+	Bucket       int           `yaml:"bucket"`       // bucket number for each window
+	CPUThreshold int64         `yaml:"cpuThreshold"` // CPU threshold; valid range 0–1000 (1000 = 100%)
+}
+
+func RateLimitMiddleware(config *RateLimiter) gin.HandlerFunc {
+	if !config.Enable {
+		return func(c *gin.Context) {
+			c.Next()
+		}
+	}
+
+	limiter := bbr.NewBBRLimiter(
+		bbr.WithWindow(config.Window),
+		bbr.WithBucket(config.Bucket),
+		bbr.WithCPUThreshold(config.CPUThreshold),
+	)
+
+	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"))
+			return
+		}
+
+		c.Next()
+		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)
+}