fix test cases

Signed-off-by: Arthur Schreiber <[email protected]>

Author: arthurschreiber

go/pools/smartconnpool/pool.go

@@ -259,34 +259,23 @@ func (pool *ConnPool[C]) Open(connect Connector[C], refresh RefreshCheck) *ConnP
 // but calling ConnPool.Put is still allowed. This function will not return until all of the pool's
 // connections have been returned or the default PoolCloseTimeout has elapsed
 func (pool *ConnPool[C]) Close() {
-	ctx, cancel := context.WithTimeout(context.Background(), PoolCloseTimeout)
-	defer cancel()
-
-	if err := pool.CloseWithContext(ctx); err != nil {
-		log.Errorf("failed to close pool %q: %v", pool.Name, err)
-	}
-}
-
-// CloseWithContext behaves like Close but allows passing in a Context to time out the
-// pool closing operation
-func (pool *ConnPool[C]) CloseWithContext(ctx context.Context) error {
 	pool.capacityMu.Lock()
 	defer pool.capacityMu.Unlock()
 
 	if pool.close == nil || pool.capacity.Load() == 0 {
 		// already closed
-		return nil
+		return
 	}
 
 	// close all the connections in the pool; if we time out while waiting for
 	// users to return our connections, we still want to finish the shutdown
 	// for the pool
-	err := pool.setCapacity(ctx, 0)
+	pool.setCapacity(0)
 
 	close(pool.close)
 	pool.workers.Wait()
 	pool.close = nil
-	return err
+	return
 }
 
 func (pool *ConnPool[C]) reopen() {
@@ -298,19 +287,14 @@ func (pool *ConnPool[C]) reopen() {
 		return
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), PoolCloseTimeout)
-	defer cancel()
-
 	// to re-open the connection pool, first set the capacity to 0 so we close
 	// all the existing connections, as they're now connected to a stale MySQL
 	// instance.
-	if err := pool.setCapacity(ctx, 0); err != nil {
-		log.Errorf("failed to reopen pool %q: %v", pool.Name, err)
-	}
+	pool.setCapacity(0)
 
 	// the second call to setCapacity cannot fail because it's only increasing the number
 	// of connections and doesn't need to shut down any
-	_ = pool.setCapacity(ctx, capacity)
+	pool.setCapacity(capacity)
 }
 
 // IsOpen returns whether the pool is open
@@ -713,23 +697,24 @@ func (pool *ConnPool[C]) getWithSetting(ctx context.Context, setting *Setting) (
 // that means waiting for clients to return connections to the pool.
 // If the given context times out before we've managed to close enough connections
 // an error will be returned.
-func (pool *ConnPool[C]) SetCapacity(ctx context.Context, newcap int64) error {
+func (pool *ConnPool[C]) SetCapacity(newcap int64) {
 	pool.capacityMu.Lock()
 	defer pool.capacityMu.Unlock()
-	return pool.setCapacity(ctx, newcap)
+	pool.setCapacity(newcap)
 }
 
 // setCapacity is the internal implementation for SetCapacity; it must be called
 // with pool.capacityMu being held
-func (pool *ConnPool[C]) setCapacity(ctx context.Context, newcap int64) error {
+func (pool *ConnPool[C]) setCapacity(newcap int64) {
 	if newcap < 0 {
 		panic("negative capacity")
 	}
 
 	oldcap := pool.capacity.Swap(newcap)
 	if oldcap == newcap {
-		return nil
+		return
 	}
+
 	// update the idle count to match the new capacity if necessary
 	// wait for connections to be returned to the pool if we're reducing the capacity.
 	defer pool.setIdleCount()
@@ -775,8 +760,6 @@ func (pool *ConnPool[C]) setCapacity(ctx context.Context, newcap int64) error {
 			pool.closedConn()
 		}
 	}
-
-	return nil
 }
 
 func (pool *ConnPool[C]) closeIdleResources(now time.Time) {

go/pools/smartconnpool/pool_test.go

@@ -209,15 +209,15 @@ func TestOpen(t *testing.T) {
 	assert.EqualValues(t, 6, state.lastID.Load())
 
 	// SetCapacity
-	err = p.SetCapacity(ctx, 3)
-	require.NoError(t, err)
+	p.SetCapacity(3)
+
 	assert.EqualValues(t, 3, state.open.Load())
 	assert.EqualValues(t, 6, state.lastID.Load())
 	assert.EqualValues(t, 3, p.Capacity())
 	assert.EqualValues(t, 3, p.Available())
 
-	err = p.SetCapacity(ctx, 6)
-	require.NoError(t, err)
+	p.SetCapacity(6)
+
 	assert.EqualValues(t, 6, p.Capacity())
 	assert.EqualValues(t, 6, p.Available())
 
@@ -266,13 +266,8 @@ func TestShrinking(t *testing.T) {
 		require.NoError(t, err)
 		resources[i] = r
 	}
-	done := make(chan bool)
-	go func() {
-		err := p.SetCapacity(ctx, 3)
-		require.NoError(t, err)
 
-		done <- true
-	}()
+	p.SetCapacity(3)
 	expected := map[string]any{
 		"Capacity":          3,
 		"Available":         -1, // negative because we've borrowed past our capacity
@@ -294,27 +289,59 @@ func TestShrinking(t *testing.T) {
 			assert.Equal(t, expected, stats)
 		}
 	}
-	// There are already 2 resources available in the pool.
-	// So, returning one should be enough for SetCapacity to complete.
+
+	p.put(resources[0], resources[0].generation)
+	assert.Equal(t, map[string]any{
+		"Capacity":          3,
+		"Available":         0,
+		"Active":            4,
+		"InUse":             3,
+		"WaitCount":         0,
+		"WaitTime":          time.Duration(0),
+		"IdleTimeout":       1 * time.Second,
+		"IdleClosed":        0,
+		"MaxLifetimeClosed": 0,
+	}, p.StatsJSON())
+
+	p.put(resources[1], resources[1].generation)
+	assert.Equal(t, map[string]any{
+		"Capacity":          3,
+		"Available":         1,
+		"Active":            4,
+		"InUse":             2,
+		"WaitCount":         0,
+		"WaitTime":          time.Duration(0),
+		"IdleTimeout":       1 * time.Second,
+		"IdleClosed":        0,
+		"MaxLifetimeClosed": 0,
+	}, p.StatsJSON())
+
+	p.put(resources[2], resources[2].generation)
+	assert.Equal(t, map[string]any{
+		"Capacity":          3,
+		"Available":         2,
+		"Active":            4,
+		"InUse":             1,
+		"WaitCount":         0,
+		"WaitTime":          time.Duration(0),
+		"IdleTimeout":       1 * time.Second,
+		"IdleClosed":        0,
+		"MaxLifetimeClosed": 0,
+	}, p.StatsJSON())
+
 	p.put(resources[3], resources[3].generation)
-	<-done
-	// Return the rest of the resources
-	for i := 0; i < 3; i++ {
-		p.put(resources[i], resources[i].generation)
-	}
-	stats := p.StatsJSON()
-	expected = map[string]any{
+	assert.Equal(t, map[string]any{
 		"Capacity":          3,
 		"Available":         3,
 		"Active":            3,
 		"InUse":             0,
 		"WaitCount":         0,
 		"WaitTime":          time.Duration(0),
 		"IdleTimeout":       1 * time.Second,
-		"IdleClosed":        0,
+		"IdleClosed":        1,
 		"MaxLifetimeClosed": 0,
-	}
-	assert.Equal(t, expected, stats)
+	}, p.StatsJSON())
+
 	assert.EqualValues(t, 3, state.open.Load())
 
 	// Ensure no deadlock if SetCapacity is called after we start
@@ -330,36 +357,36 @@ func TestShrinking(t *testing.T) {
 		require.NoError(t, err)
 		resources[i] = r
 	}
+
+	wg := sync.WaitGroup{}
 	// This will wait because pool is empty
+	wg.Add(1)
 	go func() {
 		r, err := p.Get(ctx, nil)
 		require.NoError(t, err)
 		p.put(r, r.generation)
-		done <- true
+		wg.Done()
 	}()
 
-	// This will also wait
-	go func() {
-		err := p.SetCapacity(ctx, 2)
-		require.NoError(t, err)
-		done <- true
-	}()
+	p.SetCapacity(2)
+
 	time.Sleep(10 * time.Millisecond)
 
 	// This should not hang
 	for i := 0; i < 3; i++ {
 		p.put(resources[i], resources[i].generation)
 	}
-	<-done
-	<-done
+	wg.Wait()
+
 	assert.EqualValues(t, 2, p.Capacity())
 	assert.EqualValues(t, 2, p.Available())
 	assert.EqualValues(t, 1, p.Metrics.WaitCount())
 	assert.EqualValues(t, p.Metrics.WaitCount(), len(state.waits))
 	assert.EqualValues(t, 2, state.open.Load())
 
 	// Test race condition of SetCapacity with itself
-	err = p.SetCapacity(ctx, 3)
+	p.SetCapacity(3)
+
 	require.NoError(t, err)
 	for i := 0; i < 3; i++ {
 		var r *Pooled[*TestConn]
@@ -372,35 +399,34 @@ func TestShrinking(t *testing.T) {
 		require.NoError(t, err)
 		resources[i] = r
 	}
+
 	// This will wait because pool is empty
+	wg.Add(1)
 	go func() {
 		r, err := p.Get(ctx, nil)
 		require.NoError(t, err)
 		p.put(r, r.generation)
-		done <- true
+		wg.Done()
 	}()
 	time.Sleep(10 * time.Millisecond)
 
 	// This will wait till we Put
-	go func() {
-		err := p.SetCapacity(ctx, 2)
-		require.NoError(t, err)
-	}()
+	p.SetCapacity(2)
+
 	time.Sleep(10 * time.Millisecond)
-	go func() {
-		err := p.SetCapacity(ctx, 4)
-		require.NoError(t, err)
-	}()
+
+	p.SetCapacity(4)
+
 	time.Sleep(10 * time.Millisecond)
 
 	// This should not hang
 	for i := 0; i < 3; i++ {
 		p.put(resources[i], resources[i].generation)
 	}
-	<-done
+	wg.Wait()
 
 	assert.Panics(t, func() {
-		_ = p.SetCapacity(ctx, -1)
+		p.SetCapacity(-1)
 	})
 
 	assert.EqualValues(t, 4, p.Capacity())