Cherry-pick afbce6a with conflicts

go/vt/vttablet/endtoend/misc_test.go

@@ -265,8 +265,10 @@ func TestSidecarTables(t *testing.T) {
 }
 
 func TestConsolidation(t *testing.T) {
-	defer framework.Server.SetPoolSize(framework.Server.PoolSize())
-	framework.Server.SetPoolSize(1)
+	defer framework.Server.SetPoolSize(context.Background(), framework.Server.PoolSize())
+
+	err := framework.Server.SetPoolSize(context.Background(), 1)
+	require.NoError(t, err)
 
 	const tag = "Waits/Histograms/Consolidations/Count"
 

go/vt/vttablet/endtoend/stream_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package endtoend
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"reflect"
@@ -98,11 +99,13 @@ func TestStreamConsolidation(t *testing.T) {
 
 	defaultPoolSize := framework.Server.StreamPoolSize()
 
-	framework.Server.SetStreamPoolSize(4)
+	err = framework.Server.SetStreamPoolSize(context.Background(), 4)
+	require.NoError(t, err)
+
 	framework.Server.SetStreamConsolidationBlocking(true)
 
 	defer func() {
-		framework.Server.SetStreamPoolSize(defaultPoolSize)
+		_ = framework.Server.SetStreamPoolSize(context.Background(), defaultPoolSize)
 		framework.Server.SetStreamConsolidationBlocking(false)
 	}()
 

go/vt/vttablet/tabletserver/connpool/pool.go

@@ -34,15 +34,9 @@ import (
 	"vitess.io/vitess/go/vt/log"
 	"vitess.io/vitess/go/vt/mysqlctl"
 	"vitess.io/vitess/go/vt/servenv"
-	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"
-
-	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
 )
 
-// ErrConnPoolClosed is returned when the connection pool is closed.
-var ErrConnPoolClosed = vterrors.New(vtrpcpb.Code_INTERNAL, "internal error: unexpected: conn pool is closed")
-
 const (
 	getWithoutS = "GetWithoutSettings"
 	getWithS    = "GetWithSettings"

go/vt/vttablet/tabletserver/connpool/pool_test.go

@@ -69,7 +69,7 @@ func TestConnPoolTimeout(t *testing.T) {
 	require.NoError(t, err)
 	defer dbConn.Recycle()
 	_, err = connPool.Get(context.Background(), nil)
-	assert.EqualError(t, err, "resource pool timed out")
+	assert.EqualError(t, err, "connection pool timed out")
 }
 
 func TestConnPoolMaxWaiters(t *testing.T) {
@@ -181,6 +181,7 @@ func TestConnPoolSetCapacity(t *testing.T) {
 	connPool := newPool()
 	connPool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
 	defer connPool.Close()
+<<<<<<< HEAD
 	err := connPool.SetCapacity(-10)
 	if err == nil {
 		t.Fatalf("set capacity should return error for negative capacity")
@@ -189,6 +190,14 @@ func TestConnPoolSetCapacity(t *testing.T) {
 	if err != nil {
 		t.Fatalf("set capacity should succeed")
 	}
+=======
+
+	assert.Panics(t, func() {
+		_ = connPool.SetCapacity(context.Background(), -10)
+	})
+	err := connPool.SetCapacity(context.Background(), 10)
+	assert.NoError(t, err)
+>>>>>>> afbce6aa87 (connpool: Allow time out during shutdown (#15979))
 	if connPool.Capacity() != 10 {
 		t.Fatalf("capacity should be 10")
 	}

go/vt/vttablet/tabletserver/debugenv.go

@@ -17,6 +17,7 @@ limitations under the License.
 package tabletserver
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"html"
@@ -82,6 +83,17 @@ func debugEnvHandler(tsv *TabletServer, w http.ResponseWriter, r *http.Request)
 			f(ival)
 			msg = fmt.Sprintf("Setting %v to: %v", varname, value)
 		}
+		setIntValCtx := func(f func(context.Context, int) error) {
+			ival, err := strconv.Atoi(value)
+			if err == nil {
+				err = f(r.Context(), ival)
+				if err == nil {
+					msg = fmt.Sprintf("Setting %v to: %v", varname, value)
+					return
+				}
+			}
+			msg = fmt.Sprintf("Failed setting value for %v: %v", varname, err)
+		}
 		setInt64Val := func(f func(int64)) {
 			ival, err := strconv.ParseInt(value, 10, 64)
 			if err != nil {
@@ -111,11 +123,11 @@ func debugEnvHandler(tsv *TabletServer, w http.ResponseWriter, r *http.Request)
 		}
 		switch varname {
 		case "PoolSize":
-			setIntVal(tsv.SetPoolSize)
+			setIntValCtx(tsv.SetPoolSize)
 		case "StreamPoolSize":
-			setIntVal(tsv.SetStreamPoolSize)
+			setIntValCtx(tsv.SetStreamPoolSize)
 		case "TxPoolSize":
-			setIntVal(tsv.SetTxPoolSize)
+			setIntValCtx(tsv.SetTxPoolSize)
 		case "MaxResultSize":
 			setIntVal(tsv.SetMaxResultSize)
 		case "WarnResultSize":

go/vt/vttablet/tabletserver/query_executor.go

@@ -772,6 +772,7 @@ func (qre *QueryExecutor) getConn() (*connpool.DBConn, error) {
 	span, ctx := trace.NewSpan(qre.ctx, "QueryExecutor.getConn")
 	defer span.Finish()
 
+<<<<<<< HEAD
 	start := time.Now()
 	conn, err := qre.tsv.qe.conns.Get(ctx, qre.setting)
 
@@ -783,12 +784,19 @@ func (qre *QueryExecutor) getConn() (*connpool.DBConn, error) {
 		return nil, err
 	}
 	return nil, err
+=======
+	defer func(start time.Time) {
+		qre.logStats.WaitingForConnection += time.Since(start)
+	}(time.Now())
+	return qre.tsv.qe.conns.Get(ctx, qre.setting)
+>>>>>>> afbce6aa87 (connpool: Allow time out during shutdown (#15979))
 }
 
 func (qre *QueryExecutor) getStreamConn() (*connpool.DBConn, error) {
 	span, ctx := trace.NewSpan(qre.ctx, "QueryExecutor.getStreamConn")
 	defer span.Finish()
 
+<<<<<<< HEAD
 	start := time.Now()
 	conn, err := qre.tsv.qe.streamConns.Get(ctx, qre.setting)
 	switch err {
@@ -799,6 +807,12 @@ func (qre *QueryExecutor) getStreamConn() (*connpool.DBConn, error) {
 		return nil, err
 	}
 	return nil, err
+=======
+	defer func(start time.Time) {
+		qre.logStats.WaitingForConnection += time.Since(start)
+	}(time.Now())
+	return qre.tsv.qe.streamConns.Get(ctx, qre.setting)
+>>>>>>> afbce6aa87 (connpool: Allow time out during shutdown (#15979))
 }
 
 // txFetch fetches from a TxConnection.

go/vt/vttablet/tabletserver/tabletserver.go

@@ -1979,11 +1979,15 @@ func (tsv *TabletServer) EnableHistorian(enabled bool) {
 }
 
 // SetPoolSize changes the pool size to the specified value.
-func (tsv *TabletServer) SetPoolSize(val int) {
+func (tsv *TabletServer) SetPoolSize(ctx context.Context, val int) error {
 	if val <= 0 {
-		return
+		return nil
 	}
+<<<<<<< HEAD
 	tsv.qe.conns.SetCapacity(val)
+=======
+	return tsv.qe.conns.SetCapacity(ctx, int64(val))
+>>>>>>> afbce6aa87 (connpool: Allow time out during shutdown (#15979))
 }
 
 // PoolSize returns the pool size.
@@ -1992,8 +1996,13 @@ func (tsv *TabletServer) PoolSize() int {
 }
 
 // SetStreamPoolSize changes the pool size to the specified value.
+<<<<<<< HEAD
 func (tsv *TabletServer) SetStreamPoolSize(val int) {
 	tsv.qe.streamConns.SetCapacity(val)
+=======
+func (tsv *TabletServer) SetStreamPoolSize(ctx context.Context, val int) error {
+	return tsv.qe.streamConns.SetCapacity(ctx, int64(val))
+>>>>>>> afbce6aa87 (connpool: Allow time out during shutdown (#15979))
 }
 
 // SetStreamConsolidationBlocking sets whether the stream consolidator should wait for slow clients
@@ -2007,8 +2016,13 @@ func (tsv *TabletServer) StreamPoolSize() int {
 }
 
 // SetTxPoolSize changes the tx pool size to the specified value.
+<<<<<<< HEAD
 func (tsv *TabletServer) SetTxPoolSize(val int) {
 	tsv.te.txPool.scp.conns.SetCapacity(val)
+=======
+func (tsv *TabletServer) SetTxPoolSize(ctx context.Context, val int) error {
+	return tsv.te.txPool.scp.conns.SetCapacity(ctx, int64(val))
+>>>>>>> afbce6aa87 (connpool: Allow time out during shutdown (#15979))
 }
 
 // TxPoolSize returns the tx pool size.

go/vt/vttablet/tabletserver/tabletserver_test.go

@@ -2046,23 +2046,29 @@ func TestConfigChanges(t *testing.T) {
 	newSize := 10
 	newDuration := time.Duration(10 * time.Millisecond)
 
-	tsv.SetPoolSize(newSize)
+	err := tsv.SetPoolSize(context.Background(), newSize)
+	require.NoError(t, err)
+
 	if val := tsv.PoolSize(); val != newSize {
 		t.Errorf("PoolSize: %d, want %d", val, newSize)
 	}
 	if val := int(tsv.qe.conns.Capacity()); val != newSize {
 		t.Errorf("tsv.qe.connPool.Capacity: %d, want %d", val, newSize)
 	}
 
-	tsv.SetStreamPoolSize(newSize)
+	err = tsv.SetStreamPoolSize(context.Background(), newSize)
+	require.NoError(t, err)
+
 	if val := tsv.StreamPoolSize(); val != newSize {
 		t.Errorf("StreamPoolSize: %d, want %d", val, newSize)
 	}
 	if val := int(tsv.qe.streamConns.Capacity()); val != newSize {
 		t.Errorf("tsv.qe.streamConnPool.Capacity: %d, want %d", val, newSize)
 	}
 
-	tsv.SetTxPoolSize(newSize)
+	err = tsv.SetTxPoolSize(context.Background(), newSize)
+	require.NoError(t, err)
+
 	if val := tsv.TxPoolSize(); val != newSize {
 		t.Errorf("TxPoolSize: %d, want %d", val, newSize)
 	}

go/vt/vttablet/tabletserver/tx_pool_test.go

@@ -214,8 +214,15 @@ func primeTxPoolWithConnection(t *testing.T, ctx context.Context) (*fakesqldb.DB
 	db := fakesqldb.New(t)
 	txPool, _ := newTxPool()
 	// Set the capacity to 1 to ensure that the db connection is reused.
+<<<<<<< HEAD
 	txPool.scp.conns.SetCapacity(1)
 	txPool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
+=======
+	err := txPool.scp.conns.SetCapacity(context.Background(), 1)
+	require.NoError(t, err)
+	params := dbconfigs.New(db.ConnParams())
+	txPool.Open(params, params, params)
+>>>>>>> afbce6aa87 (connpool: Allow time out during shutdown (#15979))
 
 	// Run a query to trigger a database connection. That connection will be
 	// reused by subsequent transactions.