feat: 实现消息提交防重配额授权与最小分发

internal/app/message/rate.go

@@ -0,0 +1,59 @@
+package message
+
+import (
+	"sync"
+	"time"
+)
+
+// rateLimiter 是每端一个令牌桶：速率 rps、容量 burst。
+// rps<=0 表示不限速。
+type rateLimiter struct {
+	rps   float64
+	burst float64
+
+	mu sync.Mutex
+	m  map[string]*tokenBucket
+}
+
+type tokenBucket struct {
+	tokens float64
+	last   time.Time
+}
+
+func newRateLimiter(rps float64, burst int) *rateLimiter {
+	if burst <= 0 {
+		burst = defaultRequestBurst
+	}
+	return &rateLimiter{
+		rps:   rps,
+		burst: float64(burst),
+		m:     make(map[string]*tokenBucket),
+	}
+}
+
+// allow 消耗 1 个令牌；允许则 true。
+func (r *rateLimiter) allow(endpointID string, now time.Time) bool {
+	if r == nil || r.rps <= 0 {
+		return true
+	}
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	b := r.m[endpointID]
+	if b == nil {
+		b = &tokenBucket{tokens: r.burst, last: now}
+		r.m[endpointID] = b
+	}
+	elapsed := now.Sub(b.last).Seconds()
+	if elapsed > 0 {
+		b.tokens += elapsed * r.rps
+		if b.tokens > r.burst {
+			b.tokens = r.burst
+		}
+		b.last = now
+	}
+	if b.tokens < 1 {
+		return false
+	}
+	b.tokens--
+	return true
+}