vtorc: use RestartReplication RPC in restartDirectReplicas (#20719)

Signed-off-by: Tim Vaillancourt <tim@timvaillancourt.com>

Author: timvaillancourt

go/vt/vtorc/logic/topology_recovery.go

@@ -31,6 +31,8 @@ import (
 
 	"github.com/patrickmn/go-cache"
 	"golang.org/x/sync/errgroup"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 
 	"vitess.io/vitess/go/stats"
 	"vitess.io/vitess/go/vt/log"
@@ -571,16 +573,11 @@ func restartDirectReplicas(ctx context.Context, analysisEntry *inst.DetectionAna
 			logger.Info("Restarting replication on direct replica " + tabletAliasString)
 			_ = AuditTopologyRecovery(topologyRecovery, "Restarting replication on direct replica "+tabletAliasString)
 
-			if err := stopReplication(ctx, tablet); err != nil {
-				logger.Error(fmt.Sprintf("Failed to stop replication on %s: %v", tabletAliasString, err))
-				return err
-			}
-
 			// Determine if this replica should use semi-sync based on the durability policy
 			semiSync := policy.IsReplicaSemiSync(durabilityPolicy, primaryTablet, tablet)
 
-			if err := startReplication(ctx, tablet, semiSync); err != nil {
-				logger.Error(fmt.Sprintf("Failed to start replication on %s: %v", tabletAliasString, err))
+			if err := restartReplication(ctx, tablet, semiSync, logger); err != nil {
+				logger.Error(fmt.Sprintf("Failed to restart replication on %s: %v", tabletAliasString, err))
 				return err
 			}
 			logger.Info("Successfully restarted replication on " + tabletAliasString)
@@ -608,20 +605,39 @@ func getShardTablets(ctx context.Context, keyspace, shard string) ([]*topo.Table
 	return ts.GetTabletsByShard(ctx, keyspace, shard)
 }
 
-// stopReplication calls StopReplication RPC for the given tablet with a timeout.
-func stopReplication(ctx context.Context, tablet *topodatapb.Tablet) error {
-	ctx, cancel := context.WithTimeout(ctx, topo.RemoteOperationTimeout)
-	defer cancel()
+// restartReplication calls RestartReplication RPC for the given tablet with a timeout.
+func restartReplication(ctx context.Context, tablet *topodatapb.Tablet, semiSync bool, logger *log.PrefixedLogger) error {
+	restartCtx, restartCancel := context.WithTimeout(ctx, topo.RemoteOperationTimeout)
+	err := tmc.RestartReplication(restartCtx, tablet, semiSync)
+	restartCancel()
+	if err == nil {
+		return nil
+	}
 
-	return tmc.StopReplication(ctx, tablet)
-}
+	// Skip the fallback when the recovery context is done, since the detached
+	// fallback cannot be canceled and would only delay recovery teardown, or
+	// when the RPC never reached the tablet, in which case STOP cannot have run.
+	if ctx.Err() != nil || status.Code(err) == codes.Unavailable {
+		return err
+	}
 
-// startReplication calls StartReplication RPC for the given tablet with a timeout.
-func startReplication(ctx context.Context, tablet *topodatapb.Tablet, semiSync bool) error {
-	ctx, cancel := context.WithTimeout(ctx, topo.RemoteOperationTimeout)
-	defer cancel()
+	// TODO: Remove this StartReplication fallback in v26, when all supported
+	// vttablets include detached post-STOP cleanup in RestartReplication.
+	tabletAlias := topoproto.TabletAliasString(tablet.Alias)
+	logger.Warn("RestartReplication failed; attempting to ensure replication is started",
+		slog.String("tablet", tabletAlias),
+		slog.Any("error", err),
+	)
+
+	startCtx, startCancel := context.WithTimeout(context.WithoutCancel(ctx), topo.RemoteOperationTimeout)
+	defer startCancel()
+	if startErr := tmc.StartReplication(startCtx, tablet, semiSync); startErr != nil {
+		// Wrap with %w (not vterrors.Wrapf, which has no Unwrap) so callers can
+		// still match on the original RestartReplication error.
+		return fmt.Errorf("failed to ensure replication was started on tablet %s after RestartReplication error (%v): %w", tabletAlias, startErr, err)
+	}
 
-	return tmc.StartReplication(ctx, tablet, semiSync)
+	return err
 }
 
 // isERSEnabled returns true if ERS can be used globally or for the given keyspace.

go/vt/vtorc/logic/topology_recovery_test.go

@@ -34,6 +34,8 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 
 	"vitess.io/vitess/go/mysql/replication"
 	"vitess.io/vitess/go/vt/external/golib/sqlutils"
@@ -1411,7 +1413,117 @@ func TestReconcileStaleTopoPrimaryTopoTimeout(t *testing.T) {
 	})
 }
 
-// TestRestartDirectReplicasTimeout verifies that restartDirectReplicas does not block forever if an RPC hangs.
+func TestRestartReplicationFallback(t *testing.T) {
+	oldTMC := tmc
+	t.Cleanup(func() {
+		tmc = oldTMC
+	})
+
+	oldRemoteOpTimeout := topo.RemoteOperationTimeout
+	topo.RemoteOperationTimeout = 100 * time.Millisecond
+	t.Cleanup(func() {
+		topo.RemoteOperationTimeout = oldRemoteOpTimeout
+	})
+
+	tablet := &topodatapb.Tablet{Alias: &topodatapb.TabletAlias{Cell: "zone1", Uid: 101}}
+	logger := log.NewPrefixedLogger("test-restart-replication-fallback")
+
+	setup := func(t *testing.T) *tmcmock.MockTabletManagerClient {
+		mockController := gomock.NewController(t)
+		mockTMC := tmcmock.NewMockTabletManagerClient(mockController)
+		tmc = mockTMC
+		return mockTMC
+	}
+
+	// The server may have run STOP REPLICA before timing out, so the fallback
+	// must run while the recovery context is still alive.
+	t.Run("fallback starts replication after RPC timeout", func(t *testing.T) {
+		mockTMC := setup(t)
+		ctx, cancel := context.WithCancel(t.Context())
+		t.Cleanup(cancel)
+
+		gomock.InOrder(
+			mockTMC.EXPECT().
+				RestartReplication(gomock.Any(), tablet, true).
+				DoAndReturn(func(restartCtx context.Context, _ *topodatapb.Tablet, _ bool) error {
+					<-restartCtx.Done()
+					return restartCtx.Err()
+				}),
+			mockTMC.EXPECT().
+				StartReplication(gomock.Any(), tablet, true).
+				DoAndReturn(func(ctx context.Context, _ *topodatapb.Tablet, _ bool) error {
+					require.NoError(t, ctx.Err())
+					_, ok := ctx.Deadline()
+					require.True(t, ok, "cleanup context must have a deadline")
+					return nil
+				}),
+		)
+
+		err := restartReplication(ctx, tablet, true, logger)
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+	})
+
+	// When the fallback also fails, the returned error must still unwrap to
+	// the original RestartReplication error so callers can match on it.
+	t.Run("fallback failure preserves original error", func(t *testing.T) {
+		mockTMC := setup(t)
+
+		gomock.InOrder(
+			mockTMC.EXPECT().
+				RestartReplication(gomock.Any(), tablet, true).
+				DoAndReturn(func(restartCtx context.Context, _ *topodatapb.Tablet, _ bool) error {
+					<-restartCtx.Done()
+					return restartCtx.Err()
+				}),
+			mockTMC.EXPECT().
+				StartReplication(gomock.Any(), tablet, true).
+				Return(errors.New("start failed")),
+		)
+
+		err := restartReplication(t.Context(), tablet, true, logger)
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+		require.ErrorContains(t, err, "failed to ensure replication was started")
+	})
+
+	// A canceled recovery context means the recovery is being torn down; the
+	// detached fallback would only delay eg.Wait() and the next recovery.
+	t.Run("fallback skipped when recovery context is canceled", func(t *testing.T) {
+		mockTMC := setup(t)
+		ctx, cancel := context.WithCancel(t.Context())
+		t.Cleanup(cancel)
+
+		mockTMC.EXPECT().
+			RestartReplication(gomock.Any(), tablet, true).
+			DoAndReturn(func(context.Context, *topodatapb.Tablet, bool) error {
+				cancel()
+				return context.Canceled
+			})
+		mockTMC.EXPECT().
+			StartReplication(gomock.Any(), gomock.Any(), gomock.Any()).
+			Times(0)
+
+		err := restartReplication(ctx, tablet, true, logger)
+		require.ErrorIs(t, err, context.Canceled)
+	})
+
+	// On transport errors the RPC never reached the tablet, so STOP cannot
+	// have run and the fallback would waste a full RemoteOperationTimeout.
+	t.Run("fallback skipped when tablet is unreachable", func(t *testing.T) {
+		mockTMC := setup(t)
+
+		mockTMC.EXPECT().
+			RestartReplication(gomock.Any(), tablet, true).
+			Return(status.Error(codes.Unavailable, "connection refused"))
+		mockTMC.EXPECT().
+			StartReplication(gomock.Any(), gomock.Any(), gomock.Any()).
+			Times(0)
+
+		err := restartReplication(t.Context(), tablet, true, logger)
+		require.Equal(t, codes.Unavailable, status.Code(err))
+	})
+}
+
+// TestRestartDirectReplicasTimeout verifies that restartDirectReplicas does not block forever if replication RPCs hang.
 func TestRestartDirectReplicasTimeout(t *testing.T) {
 	orcDB, fromCache, err := db.OpenVTOrcWithCache()
 	require.NoError(t, err)
@@ -1499,8 +1611,15 @@ func TestRestartDirectReplicasTimeout(t *testing.T) {
 		// when the passed context is canceled.
 		mockTMC := tmcmock.NewMockTabletManagerClient(mockController)
 		mockTMC.EXPECT().
-			StopReplication(gomock.Any(), gomock.Any()).
-			DoAndReturn(func(ctx context.Context, _ *topodatapb.Tablet) error {
+			RestartReplication(gomock.Any(), gomock.Any(), gomock.Any()).
+			DoAndReturn(func(ctx context.Context, _ *topodatapb.Tablet, _ bool) error {
+				<-ctx.Done()
+				return ctx.Err()
+			}).
+			Times(1)
+		mockTMC.EXPECT().
+			StartReplication(gomock.Any(), gomock.Any(), gomock.Any()).
+			DoAndReturn(func(ctx context.Context, _ *topodatapb.Tablet, _ bool) error {
 				<-ctx.Done()
 				return ctx.Err()
 			}).
@@ -1541,10 +1660,13 @@ func TestRestartDirectReplicasTimeout(t *testing.T) {
 		}()
 
 		// Let the recovery goroutine reach a blocked state before advancing fake time (in this case,
-		// hanging on the StopReplication RPC).
+		// hanging on the RestartReplication RPC).
+		synctest.Wait()
+
+		// Move fake time beyond both RPC timeout boundaries.
+		time.Sleep(topo.RemoteOperationTimeout + time.Nanosecond)
 		synctest.Wait()
 
-		// Move fake time just beyond the expected RPC timeout boundary.
 		time.Sleep(topo.RemoteOperationTimeout + time.Nanosecond)
 		synctest.Wait()
 
@@ -1555,7 +1677,7 @@ func TestRestartDirectReplicasTimeout(t *testing.T) {
 			require.NotNil(t, result.topologyRecovery, "topology recovery record must be returned")
 			require.ErrorIs(t, result.err, context.DeadlineExceeded, "restartDirectReplicas must timeout and return when a replication RPC hangs indefinitely")
 		default:
-			require.FailNowf(t, "restartDirectReplicas did not return", "expected timeout after %s when a replication RPC hangs indefinitely", topo.RemoteOperationTimeout)
+			require.FailNowf(t, "restartDirectReplicas did not return", "expected timeout after %s when replication RPCs hang indefinitely", 2*topo.RemoteOperationTimeout)
 		}
 
 		activeRecoveries, err := ReadActiveClusterRecoveries(keyspace, shard)

go/vt/vttablet/tabletmanager/rpc_replication.go

@@ -340,17 +340,28 @@ func (tm *TabletManager) RestartReplication(ctx context.Context, semiSync bool)
 		return err
 	}
 
-	semiSyncAction, err := tm.convertBoolToSemiSyncAction(ctx, semiSync)
+	// Once STOP succeeds, run the post-STOP work with cancellation detached so
+	// caller cancellation cannot leave replication stopped. The detached work is
+	// bounded by the caller's remaining budget, capped at RemoteOperationTimeout,
+	// so the action lock cannot be held well beyond the caller's deadline.
+	postStopTimeout := topo.RemoteOperationTimeout
+	if deadline, ok := ctx.Deadline(); ok {
+		postStopTimeout = min(time.Until(deadline), postStopTimeout)
+	}
+	postStopCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), postStopTimeout)
+	defer cancel()
+
+	semiSyncAction, err := tm.convertBoolToSemiSyncAction(postStopCtx, semiSync)
 	if err != nil {
 		return err
 	}
 
-	if err := tm.fixSemiSync(ctx, tm.Tablet().Type, semiSyncAction); err != nil {
+	if err := tm.fixSemiSync(postStopCtx, tm.Tablet().Type, semiSyncAction); err != nil {
 		return err
 	}
 
 	// Start replication
-	return tm.startReplicationRecoverable(ctx)
+	return tm.startReplicationRecoverable(postStopCtx)
 }
 
 // StartReplicationUntilAfter will start the replication and let it catch up