fix(middleware): make max request limiting atomic (#237)

Author: acefun29

internal/middleware/cocrrent.go

@@ -27,6 +27,25 @@ type MaxRequestIface struct {
 	lock    *sync.RWMutex
 }
 
+func (m *MaxRequestIface) tryAcquire(max int) bool {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+
+	if m.current >= max {
+		return false
+	}
+
+	m.current++
+	return true
+}
+
+func (m *MaxRequestIface) release() {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+
+	m.current--
+}
+
 func MaxRequest(max int) gin.HandlerFunc {
 	slog.Info("setting max requests", "max", max)
 	m := &MaxRequestIface{
@@ -35,20 +54,13 @@ func MaxRequest(max int) gin.HandlerFunc {
 	}
 
 	return func(c *gin.Context) {
-		m.lock.RLock()
-		if m.current >= max {
-			m.lock.RUnlock()
+		if !m.tryAcquire(max) {
 			c.JSON(http.StatusServiceUnavailable, types.ErrorResponse(-503, "Too many requests"))
 			c.Abort()
 			return
 		}
-		m.lock.RUnlock()
-		m.lock.Lock()
-		m.current++
-		m.lock.Unlock()
+		defer m.release()
+
 		c.Next()
-		m.lock.Lock()
-		m.current--
-		m.lock.Unlock()
 	}
 }

internal/middleware/cocrrent_test.go

@@ -0,0 +1,78 @@
+package middleware
+
+import (
+	"sync"
+	"testing"
+)
+
+func TestMaxRequestTryAcquireRespectsLimitUnderConcurrency(t *testing.T) {
+	m := &MaxRequestIface{
+		current: 0,
+		lock:    &sync.RWMutex{},
+	}
+
+	const (
+		max        = 1
+		goroutines = 32
+	)
+
+	start := make(chan struct{})
+	results := make(chan bool, goroutines)
+
+	var wg sync.WaitGroup
+	for i := 0; i < goroutines; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			<-start
+			results <- m.tryAcquire(max)
+		}()
+	}
+
+	close(start)
+	wg.Wait()
+	close(results)
+
+	successes := 0
+	failures := 0
+	for acquired := range results {
+		if acquired {
+			successes++
+			continue
+		}
+		failures++
+	}
+
+	if successes != max {
+		t.Fatalf("expected %d successful acquisition, got %d", max, successes)
+	}
+
+	if failures != goroutines-max {
+		t.Fatalf("expected %d failed acquisitions, got %d", goroutines-max, failures)
+	}
+
+	if m.current != max {
+		t.Fatalf("expected current to remain %d, got %d", max, m.current)
+	}
+}
+
+func TestMaxRequestReleaseRestoresCapacity(t *testing.T) {
+	m := &MaxRequestIface{
+		current: 0,
+		lock:    &sync.RWMutex{},
+	}
+
+	if !m.tryAcquire(1) {
+		t.Fatal("expected first acquisition to succeed")
+	}
+
+	if m.tryAcquire(1) {
+		t.Fatal("expected second acquisition to fail while limit is reached")
+	}
+
+	m.release()
+
+	if !m.tryAcquire(1) {
+		t.Fatal("expected acquisition to succeed after release")
+	}
+}