VReplication: Properly support cancel and delete for multi-tenant Mov…

…eTables (#16906)

Signed-off-by: Matt Lord <[email protected]>

go/cmd/vtctldclient/command/vreplication/common/cancel.go

@@ -21,6 +21,8 @@ import (
 	"sort"
 
 	"github.com/spf13/cobra"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 
 	"vitess.io/vitess/go/cmd/vtctldclient/cli"
 
@@ -31,6 +33,7 @@ var CancelOptions = struct {
 	KeepData         bool
 	KeepRoutingRules bool
 	Shards           []string
+	DeleteBatchSize  int64
 }{}
 
 func GetCancelCommand(opts *SubCommandsOpts) *cobra.Command {
@@ -60,9 +63,13 @@ func commandCancel(cmd *cobra.Command, args []string) error {
 		KeepData:         CancelOptions.KeepData,
 		KeepRoutingRules: CancelOptions.KeepRoutingRules,
 		Shards:           CancelOptions.Shards,
+		DeleteBatchSize:  CancelOptions.DeleteBatchSize,
 	}
 	resp, err := GetClient().WorkflowDelete(GetCommandCtx(), req)
 	if err != nil {
+		if grpcerr, ok := status.FromError(err); ok && (grpcerr.Code() == codes.DeadlineExceeded) {
+			return fmt.Errorf("Cancel action timed out. Please try again and the work will pick back up where it left off. Note that you can control the timeout using the --action_timeout flag and the delete batch size with --delete-batch-size.")
+		}
 		return err
 	}
 

go/cmd/vtctldclient/command/vreplication/movetables/movetables.go

@@ -28,6 +28,10 @@ import (
 	vtctldatapb "vitess.io/vitess/go/vt/proto/vtctldata"
 )
 
+// The default batch size to use when deleting tenant data
+// if a multi-tenant migration is canceled or deleted.
+const DefaultDeleteBatchSize = 1000
+
 var (
 	// base is the base command for all actions related to MoveTables.
 	base = &cobra.Command{
@@ -108,6 +112,7 @@ func registerCommands(root *cobra.Command) {
 	cancel := common.GetCancelCommand(opts)
 	cancel.Flags().BoolVar(&common.CancelOptions.KeepData, "keep-data", false, "Keep the partially copied table data from the MoveTables workflow in the target keyspace.")
 	cancel.Flags().BoolVar(&common.CancelOptions.KeepRoutingRules, "keep-routing-rules", false, "Keep the routing rules created for the MoveTables workflow.")
+	cancel.Flags().Int64Var(&common.CancelOptions.DeleteBatchSize, "delete-batch-size", DefaultDeleteBatchSize, "When cleaning up the migrated data in tables moved as part of a multi-tenant workflow, delete the records in batches of this size.")
 	common.AddShardSubsetFlag(cancel, &common.CancelOptions.Shards)
 	base.AddCommand(cancel)
 }

go/cmd/vtctldclient/command/vreplication/workflow/delete.go

@@ -32,6 +32,7 @@ var (
 	deleteOptions = struct {
 		KeepData         bool
 		KeepRoutingRules bool
+		DeleteBatchSize  int64
 	}{}
 
 	// delete makes a WorkflowDelete gRPC call to a vtctld.
@@ -55,6 +56,7 @@ func commandDelete(cmd *cobra.Command, args []string) error {
 		KeepData:         deleteOptions.KeepData,
 		KeepRoutingRules: deleteOptions.KeepRoutingRules,
 		Shards:           baseOptions.Shards,
+		DeleteBatchSize:  deleteOptions.DeleteBatchSize,
 	}
 	resp, err := common.GetClient().WorkflowDelete(common.GetCommandCtx(), req)
 	if err != nil {

go/cmd/vtctldclient/command/vreplication/workflow/workflow.go

@@ -20,6 +20,7 @@ import (
 	"github.com/spf13/cobra"
 
 	"vitess.io/vitess/go/cmd/vtctldclient/command/vreplication/common"
+	"vitess.io/vitess/go/cmd/vtctldclient/command/vreplication/movetables"
 	"vitess.io/vitess/go/vt/topo/topoproto"
 )
 
@@ -60,6 +61,7 @@ func registerCommands(root *cobra.Command) {
 	delete.MarkFlagRequired("workflow")
 	delete.Flags().BoolVar(&deleteOptions.KeepData, "keep-data", false, "Keep the partially copied table data from the workflow in the target keyspace.")
 	delete.Flags().BoolVar(&deleteOptions.KeepRoutingRules, "keep-routing-rules", false, "Keep the routing rules created for the workflow.")
+	delete.Flags().Int64Var(&deleteOptions.DeleteBatchSize, "delete-batch-size", movetables.DefaultDeleteBatchSize, "When cleaning up the migrated data in tables moved as part of a multi-tenant MoveTables workflow, delete the records in batches of this size.")
 	common.AddShardSubsetFlag(delete, &baseOptions.Shards)
 	base.AddCommand(delete)
 

go/test/endtoend/vreplication/multi_tenant_test.go

@@ -196,12 +196,26 @@ func TestMultiTenantSimple(t *testing.T) {
 
 	require.Zero(t, len(getKeyspaceRoutingRules(t, vc).Rules))
 
-	mt.Create()
-	confirmKeyspacesRoutedTo(t, sourceKeyspace, "s1", "t1", nil)
-	validateKeyspaceRoutingRules(t, vc, initialRules)
+	createFunc := func() {
+		mt.Create()
+		confirmKeyspacesRoutedTo(t, sourceKeyspace, "s1", "t1", nil)
+		validateKeyspaceRoutingRules(t, vc, initialRules)
 
-	lastIndex = insertRows(lastIndex, sourceKeyspace)
-	waitForWorkflowState(t, vc, fmt.Sprintf("%s.%s", targetKeyspace, mt.workflowName), binlogdatapb.VReplicationWorkflowState_Running.String())
+		lastIndex = insertRows(lastIndex, sourceKeyspace)
+		waitForWorkflowState(t, vc, fmt.Sprintf("%s.%s", targetKeyspace, mt.workflowName), binlogdatapb.VReplicationWorkflowState_Running.String())
+	}
+
+	t.Run("cancel", func(t *testing.T) {
+		// First let's test canceling the workflow to ensure that it properly
+		// cleans up all of the data.
+		createFunc()
+		mt.Cancel()
+		rowCount := getRowCount(t, vtgateConn, fmt.Sprintf("%s.%s", targetKeyspace, "t1"))
+		require.Zero(t, rowCount)
+	})
+
+	// Create again and run it to completion.
+	createFunc()
 
 	vdiff(t, targetKeyspace, workflowName, defaultCellName, false, true, nil)
 	mt.SwitchReads()

go/test/endtoend/vreplication/wrappers_test.go

@@ -298,7 +298,8 @@ func (v VtctldMoveTables) ReverseWrites() {
 }
 
 func (v VtctldMoveTables) Cancel() {
-	v.exec("Cancel")
+	args := []string{"Cancel", "--delete-batch-size=500"}
+	v.exec(args...)
 }
 
 func (v VtctldMoveTables) Complete() {