VStream: Allow for automatic resume after Reshard across VStreams (vitessio#15393) (#627)

Signed-off-by: Tanjin Xu <tanjin.xu@slack-corp.com>
Co-authored-by: Matt Lord <mattalord@gmail.com>

Author: 2 people

go/test/endtoend/vreplication/vreplication_test.go

@@ -1431,7 +1431,7 @@ func reshardAction(t *testing.T, action, workflow, keyspaceName, sourceShards, t
 			action, workflow, output)
 	}
 	if err != nil {
-		t.Fatalf("Reshard %s command failed with %+v\n", action, err)
+		t.Fatalf("Reshard %s command failed with %+v\nOutput: %s", action, err, output)
 	}
 }
 

go/test/endtoend/vreplication/vstream_test.go

@@ -226,7 +226,7 @@ func insertRow(keyspace, table string, id int) {
 	vtgateConn.ExecuteFetch("begin", 1000, false)
 	_, err := vtgateConn.ExecuteFetch(fmt.Sprintf("insert into %s (name) values ('%s%d')", table, table, id), 1000, false)
 	if err != nil {
-		log.Infof("error inserting row %d: %v", id, err)
+		log.Errorf("error inserting row %d: %v", id, err)
 	}
 	vtgateConn.ExecuteFetch("commit", 1000, false)
 }
@@ -390,13 +390,15 @@ func testVStreamCopyMultiKeyspaceReshard(t *testing.T, baseTabletID int) numEven
 	defer vc.TearDown()
 
 	defaultCell := vc.Cells[vc.CellNames[0]]
-	vc.AddKeyspace(t, []*Cell{defaultCell}, "unsharded", "0", vschemaUnsharded, schemaUnsharded, defaultReplicas, defaultRdonly, baseTabletID+100, nil)
+	_, err := vc.AddKeyspace(t, []*Cell{defaultCell}, "unsharded", "0", vschemaUnsharded, schemaUnsharded, defaultReplicas, defaultRdonly, baseTabletID+100, nil)
+	require.NoError(t, err)
 
 	vtgateConn := getConnection(t, vc.ClusterConfig.hostname, vc.ClusterConfig.vtgateMySQLPort)
 	defer vtgateConn.Close()
 	verifyClusterHealth(t, vc)
 
-	vc.AddKeyspace(t, []*Cell{defaultCell}, "sharded", "-80,80-", vschemaSharded, schemaSharded, defaultReplicas, defaultRdonly, baseTabletID+200, nil)
+	_, err = vc.AddKeyspace(t, []*Cell{defaultCell}, "sharded", "-80,80-", vschemaSharded, schemaSharded, defaultReplicas, defaultRdonly, baseTabletID+200, nil)
+	require.NoError(t, err)
 
 	ctx := context.Background()
 	vstreamConn, err := vtgateconn.Dial(ctx, fmt.Sprintf("%s:%d", vc.ClusterConfig.hostname, vc.ClusterConfig.vtgateGrpcPort))

go/vt/topo/faketopo/faketopo.go

@@ -21,12 +21,11 @@ import (
 	"strings"
 	"sync"
 
+	"vitess.io/vitess/go/vt/log"
+	"vitess.io/vitess/go/vt/topo"
 	"vitess.io/vitess/go/vt/topo/memorytopo"
 
-	"vitess.io/vitess/go/vt/log"
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
-
-	"vitess.io/vitess/go/vt/topo"
 )
 
 // FakeFactory implements the Factory interface. This is supposed to be used only for testing

go/vt/vtgate/vstream_manager.go

@@ -26,6 +26,8 @@ import (
 	"sync"
 	"time"
 
+	"golang.org/x/exp/maps"
+
 	"vitess.io/vitess/go/stats"
 	"vitess.io/vitess/go/vt/discovery"
 	"vitess.io/vitess/go/vt/key"
@@ -546,25 +548,42 @@ func (vs *vstream) streamFromTablet(ctx context.Context, sgtid *binlogdatapb.Sha
 		var err error
 		cells := vs.getCells()
 
-		tp, err := discovery.NewTabletPicker(ctx, vs.ts, cells, vs.vsm.cell, sgtid.Keyspace, sgtid.Shard, vs.tabletType.String(), vs.tabletPickerOptions, ignoreTablets...)
+		tpo := vs.tabletPickerOptions
+		resharded, err := vs.keyspaceHasBeenResharded(ctx, sgtid.Keyspace)
 		if err != nil {
-			log.Errorf(err.Error())
-			return err
+			return vterrors.Wrapf(err, "failed to determine if keyspace %s has been resharded", sgtid.Keyspace)
+		}
+		if resharded {
+			// The non-serving tablet in the old / non-serving shard will contain all of
+			// the GTIDs that we need before transitioning to the new shards along with
+			// the journal event that will then allow us to automatically transition to
+			// the new shards (provided the stop_on_reshard option is not set).
+			tpo.IncludeNonServingTablets = true
 		}
 
+		tabletPickerErr := func(err error) error {
+			tperr := vterrors.Wrapf(err, "failed to find a %s tablet for VStream in %s/%s within the %s cell(s)",
+				vs.tabletType.String(), sgtid.GetKeyspace(), sgtid.GetShard(), strings.Join(cells, ","))
+			log.Errorf("%v", tperr)
+			return tperr
+		}
+		tp, err := discovery.NewTabletPicker(ctx, vs.ts, cells, vs.vsm.cell, sgtid.GetKeyspace(), sgtid.GetShard(), vs.tabletType.String(), tpo, ignoreTablets...)
+		if err != nil {
+			return tabletPickerErr(err)
+		}
 		// Create a child context with a stricter timeout when picking a tablet.
 		// This will prevent hanging in the case no tablets are found.
 		tpCtx, tpCancel := context.WithTimeout(ctx, tabletPickerContextTimeout)
 		defer tpCancel()
-
 		tablet, err := tp.PickForStreaming(tpCtx)
 		if err != nil {
-			log.Errorf(err.Error())
-			return err
+			return tabletPickerErr(err)
 		}
+
 		tabletAliasString := topoproto.TabletAliasString(tablet.Alias)
-		log.Infof("Picked %s tablet %s for VStream in %s/%s within the %s cell(s)",
-			vs.tabletType.String(), tabletAliasString, sgtid.GetKeyspace(), sgtid.GetShard(), strings.Join(cells, ","))
+
+		log.Infof("Picked a %s tablet for VStream in %s/%s within the %s cell(s)",
+			vs.tabletType.String(), sgtid.GetKeyspace(), sgtid.GetShard(), strings.Join(cells, ","))
 
 		target := &querypb.Target{
 			Keyspace:   sgtid.Keyspace,
@@ -828,7 +847,7 @@ func (vs *vstream) sendAll(ctx context.Context, sgtid *binlogdatapb.ShardGtid, e
 		if err := vs.getError(); err != nil {
 			return err
 		}
-		// convert all gtids to vgtids. This should be done here while holding the lock.
+		// Convert all gtids to vgtids. This should be done here while holding the lock.
 		for j, event := range events {
 			if event.Type == binlogdatapb.VEventType_GTID {
 				// Update the VGtid and send that instead.
@@ -1054,16 +1073,7 @@ func (vs *vstream) keyspaceHasBeenResharded(ctx context.Context, keyspace string
 	// Now that we know there MAY have been an applicable reshard, let's make a
 	// definitive determination by looking at the shard keyranges.
 	// All we care about are the shard info records now.
-
-	// Added in backport of d916e81 on 1/16/2025
-	sis := func(shardMap map[string]*topo.ShardInfo) []*topo.ShardInfo {
-		ret := []*topo.ShardInfo{}
-		for _, shardInfo := range shardMap {
-			ret = append(ret, shardInfo)
-		}
-
-		return ret
-	}(shards)
+	sis := maps.Values(shards)
 
 	for i := range sis {
 		for j := range sis {