Skip to content

Commit

Permalink
fix: idle connection count logic and added unit test
Browse files Browse the repository at this point in the history
Signed-off-by: Harshit Gangal <[email protected]>
  • Loading branch information
harshit-gangal committed Jan 6, 2025
1 parent e8eb29e commit ea13abf
Show file tree
Hide file tree
Showing 2 changed files with 52 additions and 10 deletions.
17 changes: 7 additions & 10 deletions go/pools/smartconnpool/pool.go
Original file line number Diff line number Diff line change
Expand Up @@ -413,14 +413,7 @@ func (pool *ConnPool[C]) put(conn *Pooled[C]) {
}

if !pool.wait.tryReturnConn(conn) {
// Option 1: do not care if more connections are closed than the allowed idle count
if pool.active.Load() > pool.idleCount.Load() {
conn.Close()
pool.closedConn()
return
}
// Option 2: precisely maintain the idle count
if pool.tryClose(conn) {
if pool.closeOnIdleLimitReached(conn) {
return
}

Expand All @@ -435,13 +428,17 @@ func (pool *ConnPool[C]) put(conn *Pooled[C]) {
}
}

func (pool *ConnPool[C]) tryClose(conn *Pooled[C]) bool {
// closeOnIdleLimitReached closes a connection if the number of idle connections (active - inuse) in the pool
// exceeds the idleCount limit. It returns true if the connection is closed, false otherwise.
func (pool *ConnPool[C]) closeOnIdleLimitReached(conn *Pooled[C]) bool {
for {
open := pool.active.Load()
if open <= pool.idleCount.Load() {
idle := open - pool.borrowed.Load()
if idle <= pool.idleCount.Load() {
return false
}
if pool.active.CompareAndSwap(open, open-1) {
pool.Metrics.idleClosed.Add(1)
conn.Close()
return true
}
Expand Down
45 changes: 45 additions & 0 deletions go/pools/smartconnpool/pool_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -746,6 +746,51 @@ func TestExtendedLifetimeTimeout(t *testing.T) {
}
}

// TestMaxIdleCount tests the MaxIdleCount setting, to check if the pool closes
// the idle connections when the number of idle connections exceeds the limit.
func TestMaxIdleCount(t *testing.T) {
testMaxIdleCount := func(t *testing.T, setting *Setting, maxIdleCount int64, expClosedConn int) {
var state TestState

ctx := context.Background()
p := NewPool(&Config[*TestConn]{
Capacity: 5,
MaxIdleCount: maxIdleCount,
LogWait: state.LogWait,
}).Open(newConnector(&state), nil)

defer p.Close()

var conns []*Pooled[*TestConn]
for i := 0; i < 5; i++ {
r, err := p.Get(ctx, setting)
require.NoError(t, err)
assert.EqualValues(t, i+1, state.open.Load())
assert.EqualValues(t, 0, p.Metrics.IdleClosed())

conns = append(conns, r)
}

for _, conn := range conns {
p.put(conn)
}

closedConn := 0
for _, conn := range conns {
if conn.Conn.IsClosed() {
closedConn++
}
}
assert.EqualValues(t, expClosedConn, closedConn)
assert.EqualValues(t, expClosedConn, p.Metrics.IdleClosed())
}

t.Run("WithoutSettings", func(t *testing.T) { testMaxIdleCount(t, nil, 2, 3) })
t.Run("WithSettings", func(t *testing.T) { testMaxIdleCount(t, sFoo, 2, 3) })
t.Run("WithoutSettings-MaxIdleCount-Zero", func(t *testing.T) { testMaxIdleCount(t, nil, 0, 0) })
t.Run("WithSettings-MaxIdleCount-Zero", func(t *testing.T) { testMaxIdleCount(t, sFoo, 0, 0) })
}

func TestCreateFail(t *testing.T) {
var state TestState
state.chaos.failConnect = true
Expand Down

0 comments on commit ea13abf

Please sign in to comment.