From ea13abf3344dc8467b3fda1bab1c36b49f8a2d38 Mon Sep 17 00:00:00 2001
From: Harshit Gangal <harshit@planetscale.com>
Date: Mon, 6 Jan 2025 17:29:04 +0530
Subject: [PATCH] fix: idle connection count logic and added unit test

Signed-off-by: Harshit Gangal <harshit@planetscale.com>
---
 go/pools/smartconnpool/pool.go      | 17 +++++------
 go/pools/smartconnpool/pool_test.go | 45 +++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 10 deletions(-)

diff --git a/go/pools/smartconnpool/pool.go b/go/pools/smartconnpool/pool.go
index 91b90c3772f..33b91d8068c 100644
--- a/go/pools/smartconnpool/pool.go
+++ b/go/pools/smartconnpool/pool.go
@@ -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
 		}
 
@@ -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
 		}
diff --git a/go/pools/smartconnpool/pool_test.go b/go/pools/smartconnpool/pool_test.go
index 701327005ad..44bd431d189 100644
--- a/go/pools/smartconnpool/pool_test.go
+++ b/go/pools/smartconnpool/pool_test.go
@@ -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