(Experimental) VReplication Workflows: cache workflow metadata in topo and use to filter target shards for vreplication commands

go/vt/topo/keyspace.go

@@ -203,7 +203,7 @@ func (ts *Server) GetKeyspace(ctx context.Context, keyspace string) (*KeyspaceIn
 	if err = k.UnmarshalVT(data); err != nil {
 		return nil, vterrors.Wrap(err, "bad keyspace data")
 	}
-
+	log.Infof("GetKeyspace:%v:%v:%v:", keyspace, version, k)
 	return &KeyspaceInfo{
 		keyspace: keyspace,
 		version:  version,

go/vt/topo/keyspace_workflow.go

@@ -0,0 +1,95 @@
+/*
+Copyright 2023 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package topo
+
+import (
+	"context"
+
+	"vitess.io/vitess/go/vt/proto/topodata"
+)
+
+func (ts *Server) SaveWorkflowMetadata(ctx context.Context, targetKeyspace string, wm *topodata.WorkflowMetadata) (err error) {
+	ctx, unlock, lockErr := ts.LockKeyspace(ctx, targetKeyspace, "GetWorkflowMetadata")
+	if lockErr != nil {
+		err = lockErr
+		return err
+	}
+	defer unlock(&err)
+
+	ki, err := ts.GetKeyspace(ctx, targetKeyspace)
+	// we overwrite the workflow if it already exists
+	for i, w := range ki.Workflows {
+		if w.Name == wm.Name {
+			ki.Workflows = append(ki.Workflows[:i], ki.Workflows[i+1:]...)
+		}
+	}
+	ki.Workflows = append(ki.Workflows, &topodata.WorkflowMetadata{
+		Name:           wm.Name,
+		Type:           wm.Type,
+		TargetShards:   wm.TargetShards,
+		SourceKeyspace: wm.SourceKeyspace,
+		SourceShards:   wm.SourceShards,
+	})
+	if err := ts.UpdateKeyspace(ctx, ki); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (ts *Server) DeleteWorkflowMetadata(ctx context.Context, targetKeyspace, workflowName string) (err error) {
+	ctx, unlock, lockErr := ts.LockKeyspace(ctx, targetKeyspace, "GetWorkflowMetadata")
+	if lockErr != nil {
+		err = lockErr
+		return err
+	}
+	defer unlock(&err)
+
+	ki, err := ts.GetKeyspace(ctx, targetKeyspace)
+	if err != nil {
+		return err
+	}
+	for i, w := range ki.Workflows {
+		if w.Name == workflowName {
+			ki.Workflows = append(ki.Workflows[:i], ki.Workflows[i+1:]...)
+			if err := ts.UpdateKeyspace(ctx, ki); err != nil {
+				return err
+			}
+			return nil
+		}
+	}
+	return nil
+}
+
+func (ts *Server) GetWorkflowMetadata(ctx context.Context, targetKeyspace, workflowName string) (wm *topodata.WorkflowMetadata, err error) {
+	ctx, unlock, lockErr := ts.LockKeyspace(ctx, targetKeyspace, "GetWorkflowMetadata")
+	if lockErr != nil {
+		err = lockErr
+		return nil, err
+	}
+	defer unlock(&err)
+
+	ki, err := ts.GetKeyspace(ctx, targetKeyspace)
+	if err != nil {
+		return nil, err
+	}
+	for _, w := range ki.Workflows {
+		if w.Name == workflowName {
+			return w, nil
+		}
+	}
+	return nil, nil
+}

go/vt/vtctl/workflow/utils.go

@@ -647,6 +647,40 @@ func areTabletsAvailableToStreamFrom(ctx context.Context, req *vtctldatapb.Workf
 	return nil
 }
 
+func FilterShardsForWorkflow(ctx context.Context, ts *topo.Server, keyspace, workflow string, shards []string) ([]string, error) {
+	log.Infof("Filtering shards for workflow %s in keyspace %s", workflow, keyspace)
+	ks, err := ts.GetKeyspace(ctx, keyspace)
+	if ks.Keyspace == nil || err != nil {
+		return shards, nil // keyspace doesn't exist, so we can't filter shards, should only come here in tests.
+	}
+	wm, err := ts.GetWorkflowMetadata(ctx, keyspace, workflow)
+	if err != nil {
+		return nil, err
+	}
+	if wm == nil {
+		return shards, nil
+	}
+
+	var targetShards []string
+	for _, shard := range shards {
+		for _, s := range wm.TargetShards {
+			if shard == s {
+				targetShards = append(targetShards, shard)
+			}
+		}
+	}
+
+	if len(targetShards) != len(wm.TargetShards) {
+		return nil, fmt.Errorf("target shards %v for workflow %s does not match the shards in the metadata for keyspace %s",
+			wm.TargetShards, workflow, keyspace)
+	}
+
+	if len(shards) != len(targetShards) {
+		log.Infof("workflow has fewer shards: %d than the keyspace: %d", len(targetShards), len(shards))
+	}
+	return targetShards, nil
+}
+
 // LegacyBuildTargets collects MigrationTargets and other metadata (see TargetInfo)
 // from a workflow in the target keyspace. It uses VReplicationExec to get the workflow
 // details rather than the new TabletManager ReadVReplicationWorkflow RPC. This is
@@ -658,7 +692,12 @@ func areTabletsAvailableToStreamFrom(ctx context.Context, req *vtctldatapb.Workf
 //
 // It returns ErrNoStreams if there are no targets found for the workflow.
 func LegacyBuildTargets(ctx context.Context, ts *topo.Server, tmc tmclient.TabletManagerClient, targetKeyspace string, workflow string) (*TargetInfo, error) {
-	targetShards, err := ts.GetShardNames(ctx, targetKeyspace)
+	shards, err := ts.GetShardNames(ctx, targetKeyspace)
+	if err != nil {
+		return nil, err
+	}
+
+	targetShards, err := FilterShardsForWorkflow(ctx, ts, targetKeyspace, workflow, shards)
 	if err != nil {
 		return nil, err
 	}
@@ -687,6 +726,7 @@ func LegacyBuildTargets(ctx context.Context, ts *topo.Server, tmc tmclient.Table
 	// two target shards will have vreplication streams, and the other shards in
 	// the target keyspace will not.
 	for _, targetShard := range targetShards {
+		log.Infof("Checking shard %s for workflow %s", targetShard, workflow)
 		si, err := ts.GetShard(ctx, targetKeyspace, targetShard)
 		if err != nil {
 			return nil, err

go/vt/vtctld/api_test.go

@@ -243,7 +243,8 @@ func TestAPI(t *testing.T) {
 				"snapshot_time":null,
 				"durability_policy":"semi_sync",
 				"throttler_config": null,
-				"sidecar_db_name":"_vt_sidecar_ks1"
+				"sidecar_db_name":"_vt_sidecar_ks1",
+				"workflows": []
 			}`, http.StatusOK},
 		{"GET", "keyspaces/nonexistent", "", "404 page not found", http.StatusNotFound},
 		{"POST", "keyspaces/ks1?action=TestKeyspaceAction", "", `{
@@ -324,11 +325,11 @@ func TestAPI(t *testing.T) {
 		// vtctl RunCommand
 		{"POST", "vtctl/", `["GetKeyspace","ks1"]`, `{
 		   "Error": "",
-		   "Output": "{\n  \"served_froms\": [],\n  \"keyspace_type\": 0,\n  \"base_keyspace\": \"\",\n  \"snapshot_time\": null,\n  \"durability_policy\": \"semi_sync\",\n  \"throttler_config\": null,\n  \"sidecar_db_name\": \"_vt_sidecar_ks1\"\n}\n\n"
+		   "Output": "{\n  \"served_froms\": [],\n  \"keyspace_type\": 0,\n  \"base_keyspace\": \"\",\n  \"snapshot_time\": null,\n  \"durability_policy\": \"semi_sync\",\n  \"throttler_config\": null,\n  \"sidecar_db_name\": \"_vt_sidecar_ks1\",\n  \"workflows\": []\n}\n\n"
 		}`, http.StatusOK},
 		{"POST", "vtctl/", `["GetKeyspace","ks3"]`, `{
 		   "Error": "",
-		   "Output": "{\n  \"served_froms\": [],\n  \"keyspace_type\": 1,\n  \"base_keyspace\": \"ks1\",\n  \"snapshot_time\": {\n    \"seconds\": \"1136214245\",\n    \"nanoseconds\": 0\n  },\n  \"durability_policy\": \"none\",\n  \"throttler_config\": null,\n  \"sidecar_db_name\": \"_vt\"\n}\n\n"
+		   "Output": "{\n  \"served_froms\": [],\n  \"keyspace_type\": 1,\n  \"base_keyspace\": \"ks1\",\n  \"snapshot_time\": {\n    \"seconds\": \"1136214245\",\n    \"nanoseconds\": 0\n  },\n  \"durability_policy\": \"none\",\n  \"throttler_config\": null,\n  \"sidecar_db_name\": \"_vt\",\n  \"workflows\": []\n}\n\n"
 		}`, http.StatusOK},
 		{"POST", "vtctl/", `["GetVSchema","ks3"]`, `{
 		   "Error": "",

go/vt/wrangler/materializer.go

@@ -27,6 +27,8 @@ import (
 	"text/template"
 	"time"
 
+	"vitess.io/vitess/go/vt/proto/topodata"
+
 	"google.golang.org/protobuf/encoding/prototext"
 	"google.golang.org/protobuf/proto"
 
@@ -136,7 +138,6 @@ func (wr *Wrangler) MoveTables(ctx context.Context, workflow, sourceKeyspace, ta
 	//FIXME validate tableSpecs, allTables, excludeTables
 	var tables []string
 	var externalTopo *topo.Server
-
 	if externalCluster != "" { // when the source is an external mysql cluster mounted using the Mount command
 		externalTopo, err = wr.ts.OpenExternalVitessClusterServer(ctx, externalCluster)
 		if err != nil {
@@ -359,6 +360,23 @@ func (wr *Wrangler) MoveTables(ctx context.Context, workflow, sourceKeyspace, ta
 			return fmt.Errorf(msg)
 		}
 	}
+	var participatingShards []string
+	for _, si := range mz.targetShards {
+		participatingShards = append(participatingShards, si.ShardName())
+	}
+
+	wm := &topodata.WorkflowMetadata{
+		Name:           workflow,
+		Type:           "MoveTables",
+		TargetShards:   participatingShards,
+		SourceKeyspace: sourceKeyspace,
+		SourceShards:   participatingShards,
+	}
+	log.Infof("Going to save workflow metadata for workflow %s: %+v", workflow, wm)
+	if err := wr.ts.SaveWorkflowMetadata(ctx, targetKeyspace, wm); err != nil {
+		return err
+	}
+
 	if autoStart {
 		return mz.startStreams(ctx)
 	}

go/vt/wrangler/resharder.go

@@ -23,6 +23,8 @@ import (
 	"sync"
 	"time"
 
+	"vitess.io/vitess/go/vt/proto/topodata"
+
 	"google.golang.org/protobuf/encoding/prototext"
 
 	"vitess.io/vitess/go/vt/log"
@@ -95,6 +97,17 @@ func (wr *Wrangler) Reshard(ctx context.Context, keyspace, workflow string, sour
 		return vterrors.Wrap(err, "createStreams")
 	}
 
+	wm := &topodata.WorkflowMetadata{
+		Name:         workflow,
+		Type:         "Reshard",
+		TargetShards: targets,
+		SourceShards: sources,
+	}
+	log.Infof("Going to save workflow metadata for workflow %s: %+v", workflow, wm)
+	if err := wr.ts.SaveWorkflowMetadata(ctx, keyspace, wm); err != nil {
+		return err
+	}
+
 	if autoStart {
 		if err := rs.startStreams(ctx); err != nil {
 			return vterrors.Wrap(err, "startStreams")