fix(shared): clean up zero-count entries in ConnectionLimiter to prevent memory leak

packages/shared/pkg/connlimit/limiter.go

@@ -20,7 +20,7 @@ func NewConnectionLimiter() *ConnectionLimiter {
 
 func (l *ConnectionLimiter) getCounter(key string) *atomic.Int64 {
 	return l.connections.Upsert(key, &atomic.Int64{}, func(exists bool, valueInMap, newValue *atomic.Int64) *atomic.Int64 {
-		if exists {
+		if exists && valueInMap.Load() >= 0 {
 			return valueInMap
 		}
 
@@ -32,9 +32,12 @@ func (l *ConnectionLimiter) getCounter(key string) *atomic.Int64 {
 // Returns (current count after increment, true) if successful, or (current count, false) if limit exceeded.
 // If maxLimit is negative, no limit is enforced. If maxLimit is 0, all connections are blocked.
 func (l *ConnectionLimiter) TryAcquire(key string, maxLimit int) (int64, bool) {
-	counter := l.getCounter(key)
 	for {
+		counter := l.getCounter(key)
 		current := counter.Load()
+		if current < 0 {
+			continue
+		}
 		if maxLimit >= 0 && current >= int64(maxLimit) {
 			return current, false
 		}
@@ -57,6 +60,14 @@ func (l *ConnectionLimiter) Release(key string) {
 			return
 		}
 		if counter.CompareAndSwap(current, current-1) {
+			if current-1 == 0 {
+				if counter.CompareAndSwap(0, -1) {
+					l.connections.RemoveCb(key, func(k string, v *atomic.Int64, exists bool) bool {
+						return exists && v == counter
+					})
+				}
+			}
+
 			return
 		}
 	}
@@ -70,7 +81,9 @@ func (l *ConnectionLimiter) Remove(key string) {
 // Count returns the current connection count for a key.
 func (l *ConnectionLimiter) Count(key string) int64 {
 	if counter, ok := l.connections.Get(key); ok {
-		return counter.Load()
+		if count := counter.Load(); count > 0 {
+			return count
+		}
 	}
 
 	return 0

packages/shared/pkg/connlimit/limiter_test.go

@@ -101,6 +101,33 @@ func TestConnectionLimiter_Release(t *testing.T) {
 		limiter.Release("unknown")
 	})
 
+	t.Run("release removes map entry when count reaches zero", func(t *testing.T) {
+		t.Parallel()
+		limiter := NewConnectionLimiter()
+
+		_, ok := limiter.TryAcquire("sandbox1", 5)
+		assert.True(t, ok)
+		assert.Equal(t, 1, limiter.connections.Count())
+
+		limiter.Release("sandbox1")
+		assert.Equal(t, int64(0), limiter.Count("sandbox1"))
+		assert.Equal(t, 0, limiter.connections.Count())
+	})
+
+	t.Run("acquiring after zero count cleanup starts fresh", func(t *testing.T) {
+		t.Parallel()
+		limiter := NewConnectionLimiter()
+
+		limiter.TryAcquire("sandbox1", 5)
+		limiter.Release("sandbox1")
+		assert.Equal(t, 0, limiter.connections.Count())
+
+		count, ok := limiter.TryAcquire("sandbox1", 5)
+		assert.True(t, ok)
+		assert.Equal(t, int64(1), count)
+		assert.Equal(t, 1, limiter.connections.Count())
+	})
+
 	t.Run("double release does not go negative", func(t *testing.T) {
 		t.Parallel()
 		limiter := NewConnectionLimiter()
@@ -224,4 +251,28 @@ func TestConnectionLimiter_Concurrent(t *testing.T) {
 
 		wg.Wait()
 	})
+
+	t.Run("concurrent acquire during release does not orphan connections", func(t *testing.T) {
+		t.Parallel()
+		limiter := NewConnectionLimiter()
+		const limit = 100
+		const goroutines = 50
+		const iterations = 500
+
+		var wg sync.WaitGroup
+		for range goroutines {
+			wg.Go(func() {
+				for range iterations {
+					if _, ok := limiter.TryAcquire("sandbox-race", limit); ok {
+						limiter.Release("sandbox-race")
+					}
+				}
+			})
+		}
+
+		wg.Wait()
+
+		assert.Equal(t, int64(0), limiter.Count("sandbox-race"))
+		assert.Equal(t, 0, limiter.connections.Count())
+	})
 }