VTGate: fix warming reads timeout context (#19674)

Signed-off-by: Tim Vaillancourt <tim@timvaillancourt.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: timvaillancourt

go/vt/vtgate/engine/fake_vcursor_test.go

@@ -124,6 +124,10 @@ func (t *noopVCursor) CloneForReplicaWarming(ctx context.Context) VCursor {
 	panic("implement me")
 }
 
+func (t *noopVCursor) WarmingReadsContext(ctx context.Context) (context.Context, context.CancelFunc) {
+	panic("implement me")
+}
+
 func (t *noopVCursor) CloneForMirroring(ctx context.Context) VCursor {
 	panic("implement me")
 }
@@ -470,10 +474,11 @@ type loggingVCursor struct {
 
 	parser *sqlparser.Parser
 
-	onMirrorClonesFn       func(context.Context) VCursor
-	onExecuteMultiShardFn  func(context.Context, Primitive, []*srvtopo.ResolvedShard, []*querypb.BoundQuery, bool, bool)
-	onStreamExecuteMultiFn func(context.Context, Primitive, string, []*srvtopo.ResolvedShard, []map[string]*querypb.BindVariable, bool, bool, func(*sqltypes.Result) error)
-	onRecordMirrorStatsFn  func(time.Duration, time.Duration, error)
+	onMirrorClonesFn        func(context.Context) VCursor
+	onExecuteMultiShardFn   func(context.Context, Primitive, []*srvtopo.ResolvedShard, []*querypb.BoundQuery, bool, bool)
+	onStreamExecuteMultiFn  func(context.Context, Primitive, string, []*srvtopo.ResolvedShard, []map[string]*querypb.BindVariable, bool, bool, func(*sqltypes.Result) error)
+	onRecordMirrorStatsFn   func(time.Duration, time.Duration, error)
+	onResolveDestinationsFn func(context.Context)
 
 	metrics *Metrics
 }
@@ -599,6 +604,10 @@ func (f *loggingVCursor) CloneForReplicaWarming(ctx context.Context) VCursor {
 	return f
 }
 
+func (f *loggingVCursor) WarmingReadsContext(ctx context.Context) (context.Context, context.CancelFunc) {
+	return ctx, func() {}
+}
+
 func (f *loggingVCursor) CloneForMirroring(ctx context.Context) VCursor {
 	if f.onMirrorClonesFn != nil {
 		return f.onMirrorClonesFn(ctx)
@@ -662,6 +671,9 @@ func (f *loggingVCursor) StreamExecuteMulti(ctx context.Context, primitive Primi
 }
 
 func (f *loggingVCursor) ResolveDestinations(ctx context.Context, keyspace string, ids []*querypb.Value, destinations []key.ShardDestination) ([]*srvtopo.ResolvedShard, [][]*querypb.Value, error) {
+	if f.onResolveDestinationsFn != nil {
+		f.onResolveDestinationsFn(ctx)
+	}
 	f.log = append(f.log, fmt.Sprintf("ResolveDestinations %v %v %v", keyspace, ids, key.DestinationsString(destinations)))
 	if f.shardErr != nil {
 		return nil, nil, f.shardErr

go/vt/vtgate/engine/primitive.go

@@ -129,9 +129,16 @@ type (
 		// GetWarmingReadsChannel returns the channel for executing warming reads against replicas
 		GetWarmingReadsChannel() chan bool
 
-		// CloneForReplicaWarming clones the VCursor for re-use in warming queries to replicas
+		// CloneForReplicaWarming clones the VCursor for re-use in warming queries to replicas.
+		// The clone must be created before launching the warming goroutine to avoid
+		// concurrent access to the original VCursor. Use WarmingReadsContext on the
+		// clone to obtain a timeout-bounded context for the warming query.
 		CloneForReplicaWarming(ctx context.Context) VCursor
 
+		// WarmingReadsContext returns a timeout-bounded context for a warming query
+		// and a cancel function that must be called when the warming query completes.
+		WarmingReadsContext(ctx context.Context) (context.Context, context.CancelFunc)
+
 		// CloneForMirroring clones the VCursor for re-use in mirroring queries to other keyspaces
 		CloneForMirroring(ctx context.Context) VCursor
 		//

go/vt/vtgate/engine/route.go

@@ -525,22 +525,25 @@ func (route *Route) executeWarmingReplicaRead(ctx context.Context, vcursor VCurs
 		}
 	}
 
-	replicaVCursor := vcursor.CloneForReplicaWarming(ctx)
 	warmingReadsChannel := vcursor.GetWarmingReadsChannel()
 
 	select {
 	// if there's no more room in the channel, drop the warming read
 	case warmingReadsChannel <- true:
+		replicaVCursor := vcursor.CloneForReplicaWarming(ctx)
 		go func(replicaVCursor VCursor) {
+			warmingCtx, cancel := replicaVCursor.WarmingReadsContext(ctx)
+			// Defers run LIFO: channel slot is released first, then context is canceled.
+			defer cancel()
 			defer func() {
 				<-warmingReadsChannel
 			}()
-			rss, _, err := route.findRoute(ctx, replicaVCursor, bindVars)
+			rss, _, err := route.findRoute(warmingCtx, replicaVCursor, bindVars)
 			if err != nil {
 				return
 			}
 
-			_, errs := replicaVCursor.ExecuteMultiShard(ctx, route, rss, warmingQueries, false /*rollbackOnError*/, false /*canAutocommit*/, route.FetchLastInsertID)
+			_, errs := replicaVCursor.ExecuteMultiShard(warmingCtx, route, rss, warmingQueries, false /*rollbackOnError*/, false /*canAutocommit*/, route.FetchLastInsertID)
 			if len(errs) > 0 {
 				log.Warn(fmt.Sprintf("Failed to execute warming replica read: %v", errs))
 			} else {

go/vt/vtgate/engine/route_warming_reads_test.go

@@ -37,6 +37,7 @@ type warmingReadsVCursor struct {
 	*loggingVCursor
 	warmingReadsPercent     int
 	warmingReadsChannel     chan bool
+	warmingReadsTimeout     time.Duration
 	warmingReadsExecuteFunc func(context.Context, Primitive, []*srvtopo.ResolvedShard, []*querypb.BoundQuery, bool, bool)
 }
 
@@ -49,16 +50,26 @@ func (vc *warmingReadsVCursor) GetWarmingReadsChannel() chan bool {
 }
 
 func (vc *warmingReadsVCursor) CloneForReplicaWarming(ctx context.Context) VCursor {
+	clonedLogging := &loggingVCursor{
+		shards:                  vc.shards,
+		results:                 vc.results,
+		onResolveDestinationsFn: vc.onResolveDestinationsFn,
+	}
 	clone := &warmingReadsVCursor{
-		loggingVCursor:          vc.loggingVCursor,
+		loggingVCursor:          clonedLogging,
 		warmingReadsPercent:     vc.warmingReadsPercent,
 		warmingReadsChannel:     vc.warmingReadsChannel,
+		warmingReadsTimeout:     vc.warmingReadsTimeout,
 		warmingReadsExecuteFunc: vc.warmingReadsExecuteFunc,
 	}
 	clone.onExecuteMultiShardFn = vc.warmingReadsExecuteFunc
 	return clone
 }
 
+func (vc *warmingReadsVCursor) WarmingReadsContext(ctx context.Context) (context.Context, context.CancelFunc) {
+	return context.WithTimeout(context.Background(), vc.warmingReadsTimeout)
+}
+
 func TestWarmingReadsSkipsForUpdate(t *testing.T) {
 	vindex, _ := vindexes.CreateVindex("hash", "", nil)
 	testCases := []struct {
@@ -129,29 +140,160 @@ func TestWarmingReadsSkipsForUpdate(t *testing.T) {
 
 			var warmingReadExecuted atomic.Bool
 			var capturedQuery string
+			var capturedCtxHasDeadline atomic.Bool
+			var capturedCtxErr atomic.Pointer[error]
+			var resolveDestCtxHasDeadline atomic.Bool
+			// done is closed by the test to unblock the warming read goroutine
+			// after context assertions have been made.
+			done := make(chan struct{})
 			vc := &warmingReadsVCursor{
 				loggingVCursor: &loggingVCursor{
 					shards:  []string{"-20", "20-"},
 					results: []*sqltypes.Result{defaultSelectResult},
+					onResolveDestinationsFn: func(ctx context.Context) {
+						_, hasDeadline := ctx.Deadline()
+						resolveDestCtxHasDeadline.Store(hasDeadline)
+					},
 				},
 				warmingReadsPercent: 100,
 				warmingReadsChannel: make(chan bool, 1),
+				warmingReadsTimeout: 5 * time.Second,
 			}
 			vc.warmingReadsExecuteFunc = func(ctx context.Context, primitive Primitive, rss []*srvtopo.ResolvedShard, queries []*querypb.BoundQuery, rollbackOnError, canAutocommit bool) {
 				if len(queries) > 0 {
 					capturedQuery = queries[0].Sql
 				}
+				_, hasDeadline := ctx.Deadline()
+				capturedCtxHasDeadline.Store(hasDeadline)
+				ctxErr := ctx.Err()
+				capturedCtxErr.Store(&ctxErr)
 				warmingReadExecuted.Store(true)
+				// Block until the test has checked our context assertions,
+				// preventing defer cancel() from running.
+				select {
+				case <-done:
+				case <-t.Context().Done():
+				}
 			}
 
-			_, err := route.TryExecute(t.Context(), vc, map[string]*querypb.BindVariable{}, false)
+			// Use a cancelable parent context to verify the warming read
+			// context is independent of the parent request context.
+			parentCtx, parentCancel := context.WithCancel(t.Context())
+			_, err := route.TryExecute(parentCtx, vc, map[string]*querypb.BindVariable{}, false)
 			require.NoError(t, err)
 
+			// Cancel the parent context to simulate the primary request completing.
+			parentCancel()
+
 			require.Eventually(t, func() bool {
 				return warmingReadExecuted.Load()
 			}, time.Second, 10*time.Millisecond, "warming read should be executed")
 
 			require.Equal(t, tc.expectedWarmingQuery, capturedQuery, "warming read query should match expected")
+			require.True(t, capturedCtxHasDeadline.Load(), "warming read context should have a deadline from the timeout")
+
+			// The warming read context should still be active even though the
+			// parent request context was canceled.
+			require.NoError(t, *capturedCtxErr.Load(), "warming read context should not be canceled when parent context is canceled")
+
+			// Verify findRoute received the warming context (with deadline), not the parent context.
+			require.True(t, resolveDestCtxHasDeadline.Load(), "ResolveDestinations should receive a context with deadline from the warming timeout")
+
+			// Unblock the warming read goroutine.
+			close(done)
 		})
 	}
 }
+
+func TestWarmingReadsDroppedWhenChannelFull(t *testing.T) {
+	vindex, _ := vindexes.CreateVindex("hash", "", nil)
+	route := NewRoute(
+		EqualUnique,
+		&vindexes.Keyspace{
+			Name:    "ks",
+			Sharded: true,
+		},
+		"SELECT * FROM users WHERE id = 1",
+		"dummy_select_field",
+	)
+	parser, _ := sqlparser.NewTestParser().Parse("SELECT * FROM users WHERE id = 1")
+	route.QueryStatement = parser
+	route.Vindex = vindex.(vindexes.SingleColumn)
+	route.Values = []evalengine.Expr{
+		evalengine.NewLiteralInt(1),
+	}
+
+	var warmingReadExecuted atomic.Bool
+	vc := &warmingReadsVCursor{
+		loggingVCursor: &loggingVCursor{
+			shards:  []string{"-20", "20-"},
+			results: []*sqltypes.Result{defaultSelectResult},
+		},
+		warmingReadsPercent: 100,
+		warmingReadsChannel: make(chan bool, 1),
+		warmingReadsTimeout: 5 * time.Second,
+	}
+	vc.warmingReadsExecuteFunc = func(ctx context.Context, primitive Primitive, rss []*srvtopo.ResolvedShard, queries []*querypb.BoundQuery, rollbackOnError, canAutocommit bool) {
+		warmingReadExecuted.Store(true)
+	}
+
+	// Pre-fill the channel to simulate a full pool.
+	vc.warmingReadsChannel <- true
+
+	_, err := route.TryExecute(t.Context(), vc, map[string]*querypb.BindVariable{}, false)
+	require.NoError(t, err)
+
+	// Verify over a short window that no warming read is executed while the channel is full.
+	require.Never(t, func() bool {
+		return warmingReadExecuted.Load()
+	}, 100*time.Millisecond, 5*time.Millisecond, "warming read should not execute when the channel is full")
+	// Drain the channel.
+	<-vc.warmingReadsChannel
+}
+
+func TestWarmingReadsContextTimeout(t *testing.T) {
+	vindex, _ := vindexes.CreateVindex("hash", "", nil)
+	route := NewRoute(
+		EqualUnique,
+		&vindexes.Keyspace{
+			Name:    "ks",
+			Sharded: true,
+		},
+		"SELECT * FROM users WHERE id = 1",
+		"dummy_select_field",
+	)
+	parser, _ := sqlparser.NewTestParser().Parse("SELECT * FROM users WHERE id = 1")
+	route.QueryStatement = parser
+	route.Vindex = vindex.(vindexes.SingleColumn)
+	route.Values = []evalengine.Expr{
+		evalengine.NewLiteralInt(1),
+	}
+
+	var capturedCtxErr atomic.Pointer[error]
+	var warmingReadExecuted atomic.Bool
+	vc := &warmingReadsVCursor{
+		loggingVCursor: &loggingVCursor{
+			shards:  []string{"-20", "20-"},
+			results: []*sqltypes.Result{defaultSelectResult},
+		},
+		warmingReadsPercent: 100,
+		warmingReadsChannel: make(chan bool, 1),
+		warmingReadsTimeout: 1 * time.Millisecond,
+	}
+	vc.warmingReadsExecuteFunc = func(ctx context.Context, primitive Primitive, rss []*srvtopo.ResolvedShard, queries []*querypb.BoundQuery, rollbackOnError, canAutocommit bool) {
+		// Block until the warming context times out.
+		<-ctx.Done()
+		ctxErr := ctx.Err()
+		capturedCtxErr.Store(&ctxErr)
+		warmingReadExecuted.Store(true)
+	}
+
+	_, err := route.TryExecute(t.Context(), vc, map[string]*querypb.BindVariable{}, false)
+	require.NoError(t, err)
+
+	require.Eventually(t, func() bool {
+		return warmingReadExecuted.Load()
+	}, time.Second, 10*time.Millisecond, "warming read should have been executed and timed out")
+
+	require.ErrorIs(t, *capturedCtxErr.Load(), context.DeadlineExceeded, "warming read context should have timed out")
+}

go/vt/vtgate/executorcontext/vcursor_impl.go

@@ -299,12 +299,6 @@ func (vc *VCursorImpl) CloneForMirroring(ctx context.Context) engine.VCursor {
 }
 
 func (vc *VCursorImpl) CloneForReplicaWarming(ctx context.Context) engine.VCursor {
-	callerId := callerid.EffectiveCallerIDFromContext(ctx)
-	immediateCallerId := callerid.ImmediateCallerIDFromContext(ctx)
-
-	timedCtx, _ := context.WithTimeout(context.Background(), vc.config.WarmingReadsTimeout) //nolint
-	clonedCtx := callerid.NewContext(timedCtx, callerId, immediateCallerId)
-
 	v := &VCursorImpl{
 		config:         vc.config,
 		SafeSession:    NewAutocommitSession(vc.SafeSession.Session),
@@ -315,7 +309,7 @@ func (vc *VCursorImpl) CloneForReplicaWarming(ctx context.Context) engine.VCurso
 		executor:       vc.executor,
 		resolver:       vc.resolver,
 		topoServer:     vc.topoServer,
-		logStats:       &logstats.LogStats{Ctx: clonedCtx},
+		logStats:       &logstats.LogStats{},
 		metrics:        vc.metrics,
 
 		ignoreMaxMemoryRows: vc.ignoreMaxMemoryRows,
@@ -332,6 +326,19 @@ func (vc *VCursorImpl) CloneForReplicaWarming(ctx context.Context) engine.VCurso
 	return v
 }
 
+func (vc *VCursorImpl) WarmingReadsContext(ctx context.Context) (context.Context, context.CancelFunc) {
+	callerId := callerid.EffectiveCallerIDFromContext(ctx)
+	immediateCallerId := callerid.ImmediateCallerIDFromContext(ctx)
+
+	baseCtx := context.WithoutCancel(ctx)
+	timedCtx, cancel := context.WithTimeout(baseCtx, vc.config.WarmingReadsTimeout)
+	clonedCtx := callerid.NewContext(timedCtx, callerId, immediateCallerId)
+
+	vc.logStats = &logstats.LogStats{Ctx: clonedCtx}
+
+	return clonedCtx, cancel
+}
+
 func (vc *VCursorImpl) cloneWithAutocommitSession() *VCursorImpl {
 	safeSession := vc.SafeSession.NewAutocommitSession()
 	return &VCursorImpl{