diff --git a/packages/shared/pkg/connlimit/limiter.go b/packages/shared/pkg/connlimit/limiter.go index a694afc742..ff4a30dd56 100644 --- a/packages/shared/pkg/connlimit/limiter.go +++ b/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 diff --git a/packages/shared/pkg/connlimit/limiter_test.go b/packages/shared/pkg/connlimit/limiter_test.go index 2478b2c4c8..45d1f582af 100644 --- a/packages/shared/pkg/connlimit/limiter_test.go +++ b/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()) + }) }