Cherry-pick fd9662b with conflicts

go/vt/discovery/tablet_picker_test.go

@@ -194,6 +194,7 @@ func TestPickMultiCell(t *testing.T) {
 	assert.True(t, proto.Equal(want, tablet), "Pick: %v, want %v", tablet, want)
 }
 
+<<<<<<< HEAD
 func TestPickPrimary(t *testing.T) {
 	te := newPickerTestEnv(t, []string{"cell", "otherCell"})
 	want := addTablet(te, 100, topodatapb.TabletType_PRIMARY, "cell", true, true)
@@ -204,6 +205,49 @@ func TestPickPrimary(t *testing.T) {
 		si.PrimaryAlias = want.Alias
 		return nil
 	})
+=======
+// TestPickUsingCellAsAlias confirms that when the tablet picker is
+// given a cell name that is an alias, it will choose a tablet that
+// exists within a cell that is part of the alias.
+func TestPickUsingCellAsAlias(t *testing.T) {
+	ctx := utils.LeakCheckContext(t)
+
+	// The test env puts all cells into an alias called "cella".
+	// We're also going to specify an optional extraCell that is NOT
+	// added to the alias.
+	te := newPickerTestEnv(t, ctx, []string{"cell1", "cell2", "cell3"}, "xtracell")
+	// Specify the alias as the cell.
+	tp, err := NewTabletPicker(ctx, te.topoServ, []string{"cella"}, "cell1", te.keyspace, te.shard, "replica", TabletPickerOptions{})
+	require.NoError(t, err)
+
+	// Create a tablet in one of the main cells, it should be
+	// picked as it is part of the cella alias. This tablet is
+	// NOT part of the talbet picker's local cell (cell1) so it
+	// will not be given local preference.
+	want := addTablet(ctx, te, 101, topodatapb.TabletType_REPLICA, "cell2", true, true)
+	defer deleteTablet(t, te, want)
+	// Create a tablet in an extra cell which is thus NOT part of
+	// the cella alias so it should NOT be picked.
+	noWant := addTablet(ctx, te, 102, topodatapb.TabletType_REPLICA, "xtracell", true, true)
+	defer deleteTablet(t, te, noWant)
+	// Try it many times to be sure we don't ever pick the wrong one.
+	for i := 0; i < 100; i++ {
+		tablet, err := tp.PickForStreaming(ctx)
+		require.NoError(t, err)
+		assert.True(t, proto.Equal(want, tablet), "Pick: %v, want %v", tablet, want)
+	}
+}
+
+func TestPickUsingCellAliasOnlySpecified(t *testing.T) {
+	ctx := utils.LeakCheckContextTimeout(t, 200*time.Millisecond)
+
+	// test env puts all cells into an alias called "cella"
+	te := newPickerTestEnv(t, ctx, []string{"cell", "otherCell"})
+	want1 := addTablet(ctx, te, 100, topodatapb.TabletType_REPLICA, "cell", true, true)
+	defer deleteTablet(t, te, want1)
+
+	tp, err := NewTabletPicker(ctx, te.topoServ, []string{"cella"}, "cell", te.keyspace, te.shard, "replica", TabletPickerOptions{CellPreference: "OnlySpecified"})
+>>>>>>> fd9662b1d9 (VDiff: correct handling of default source and target cells (#13969))
 	require.NoError(t, err)
 
 	tp, err := NewTabletPicker(te.topoServ, []string{"otherCell"}, te.keyspace, te.shard, "primary")
@@ -428,17 +472,30 @@ type pickerTestEnv struct {
 	topoServ *topo.Server
 }
 
+<<<<<<< HEAD
 func newPickerTestEnv(t *testing.T, cells []string) *pickerTestEnv {
 	ctx := context.Background()
 
+=======
+// newPickerTestEnv creates a test environment for TabletPicker tests.
+// It creates a cell alias called 'cella' which contains all of the
+// provided cells. However, if any optional extraCells are provided, those
+// are NOT added to the cell alias.
+func newPickerTestEnv(t *testing.T, ctx context.Context, cells []string, extraCells ...string) *pickerTestEnv {
+	allCells := append(cells, extraCells...)
+>>>>>>> fd9662b1d9 (VDiff: correct handling of default source and target cells (#13969))
 	te := &pickerTestEnv{
 		t:        t,
 		keyspace: "ks",
 		shard:    "0",
 		cells:    cells,
+<<<<<<< HEAD
 		topoServ: memorytopo.NewServer(cells...),
+=======
+		topoServ: memorytopo.NewServer(ctx, allCells...),
+>>>>>>> fd9662b1d9 (VDiff: correct handling of default source and target cells (#13969))
 	}
-	// create cell alias
+	// Create cell alias containing the cells (but NOT the extraCells).
 	err := te.topoServ.CreateCellsAlias(ctx, "cella", &topodatapb.CellsAlias{
 		Cells: cells,
 	})

go/vt/vttablet/tabletmanager/vdiff/action.go

@@ -20,6 +20,8 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"sort"
+	"strings"
 
 	"github.com/google/uuid"
 
@@ -108,6 +110,9 @@ func (vde *Engine) getVDiffSummary(vdiffID int64, dbClient binlogplayer.DBClient
 // Validate vdiff options. Also setup defaults where applicable.
 func (vde *Engine) fixupOptions(options *tabletmanagerdatapb.VDiffOptions) (*tabletmanagerdatapb.VDiffOptions, error) {
 	// Assign defaults to sourceCell and targetCell if not specified.
+	if options == nil {
+		options = &tabletmanagerdatapb.VDiffOptions{}
+	}
 	sourceCell := options.PickerOptions.SourceCell
 	targetCell := options.PickerOptions.TargetCell
 	var defaultCell string
@@ -118,10 +123,10 @@ func (vde *Engine) fixupOptions(options *tabletmanagerdatapb.VDiffOptions) (*tab
 			return nil, err
 		}
 	}
-	if sourceCell == "" {
+	if sourceCell == "" { // Default is all cells
 		sourceCell = defaultCell
 	}
-	if targetCell == "" {
+	if targetCell == "" { // Default is all cells
 		targetCell = defaultCell
 	}
 	options.PickerOptions.SourceCell = sourceCell
@@ -130,6 +135,8 @@ func (vde *Engine) fixupOptions(options *tabletmanagerdatapb.VDiffOptions) (*tab
 	return options, nil
 }
 
+// getDefaultCell returns all of the cells in the topo as a comma
+// separated string as the default value is all available cells.
 func (vde *Engine) getDefaultCell() (string, error) {
 	cells, err := vde.ts.GetCellInfoNames(vde.ctx)
 	if err != nil {
@@ -139,7 +146,8 @@ func (vde *Engine) getDefaultCell() (string, error) {
 		// Unreachable
 		return "", fmt.Errorf("there are no cells in the topo")
 	}
-	return cells[0], nil
+	sort.Strings(cells) // Ensure that the resulting value is deterministic
+	return strings.Join(cells, ","), nil
 }
 
 func (vde *Engine) handleCreateResumeAction(ctx context.Context, dbClient binlogplayer.DBClient, action VDiffAction, req *tabletmanagerdatapb.VDiffRequest, resp *tabletmanagerdatapb.VDiffResponse) error {

go/vt/vttablet/tabletmanager/vdiff/action_test.go

@@ -24,11 +24,14 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
 
 	"vitess.io/vitess/go/sqltypes"
+	"vitess.io/vitess/go/vt/vterrors"
+
 	tabletmanagerdatapb "vitess.io/vitess/go/vt/proto/tabletmanagerdata"
+	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
 	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
-	"vitess.io/vitess/go/vt/vterrors"
 )
 
 func TestPerformVDiffAction(t *testing.T) {
@@ -43,6 +46,8 @@ func TestPerformVDiffAction(t *testing.T) {
 		name          string
 		vde           *Engine
 		req           *tabletmanagerdatapb.VDiffRequest
+		preFunc       func() error
+		postFunc      func() error
 		want          *tabletmanagerdatapb.VDiffResponse
 		expectQueries []string
 		wantErr       error
@@ -52,6 +57,62 @@ func TestPerformVDiffAction(t *testing.T) {
 			vde:     &Engine{isOpen: false},
 			wantErr: vterrors.New(vtrpcpb.Code_UNAVAILABLE, "vdiff engine is closed"),
 		},
+		{
+			name: "create with defaults",
+			req: &tabletmanagerdatapb.VDiffRequest{
+				Action:    string(CreateAction),
+				VdiffUuid: uuid,
+				Options: &tabletmanagerdatapb.VDiffOptions{
+					PickerOptions: &tabletmanagerdatapb.VDiffPickerOptions{},
+				},
+			},
+			// Add a second cell. The default for source_cell and target_cell is all
+			// available cells, so this additional cell should then show up in the
+			// created vdiff record.
+			preFunc: func() error {
+				return tstenv.TopoServ.CreateCellInfo(ctx, "zone100_test", &topodatapb.CellInfo{})
+			},
+			expectQueries: []string{
+				fmt.Sprintf("select id as id from _vt.vdiff where vdiff_uuid = %s", encodeString(uuid)),
+				fmt.Sprintf(`insert into _vt.vdiff(keyspace, workflow, state, options, shard, db_name, vdiff_uuid) values('', '', 'pending', '{\"picker_options\":{\"source_cell\":\"cell1,zone100_test\",\"target_cell\":\"cell1,zone100_test\"}}', '0', 'vt_vttest', %s)`, encodeString(uuid)),
+			},
+			postFunc: func() error {
+				return tstenv.TopoServ.DeleteCellInfo(ctx, "zone100_test", true)
+			},
+		},
+		{
+			name: "create with cell alias",
+			req: &tabletmanagerdatapb.VDiffRequest{
+				Action:    string(CreateAction),
+				VdiffUuid: uuid,
+				Options: &tabletmanagerdatapb.VDiffOptions{
+					PickerOptions: &tabletmanagerdatapb.VDiffPickerOptions{
+						SourceCell: "all",
+						TargetCell: "all",
+					},
+				},
+			},
+			// Add a second cell and create an cell alias that contains it.
+			preFunc: func() error {
+				if err := tstenv.TopoServ.CreateCellInfo(ctx, "zone100_test", &topodatapb.CellInfo{}); err != nil {
+					return err
+				}
+				cells := append(tstenv.Cells, "zone100_test")
+				return tstenv.TopoServ.CreateCellsAlias(ctx, "all", &topodatapb.CellsAlias{
+					Cells: cells,
+				})
+			},
+			expectQueries: []string{
+				fmt.Sprintf("select id as id from _vt.vdiff where vdiff_uuid = %s", encodeString(uuid)),
+				fmt.Sprintf(`insert into _vt.vdiff(keyspace, workflow, state, options, shard, db_name, vdiff_uuid) values('', '', 'pending', '{\"picker_options\":{\"source_cell\":\"all\",\"target_cell\":\"all\"}}', '0', 'vt_vttest', %s)`, encodeString(uuid)),
+			},
+			postFunc: func() error {
+				if err := tstenv.TopoServ.DeleteCellInfo(ctx, "zone100_test", true); err != nil {
+					return err
+				}
+				return tstenv.TopoServ.DeleteCellsAlias(ctx, "all")
+			},
+		},
 		{
 			name: "delete by uuid",
 			req: &tabletmanagerdatapb.VDiffRequest{
@@ -80,6 +141,10 @@ func TestPerformVDiffAction(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
+			if tt.preFunc != nil {
+				err := tt.preFunc()
+				require.NoError(t, err, "pre function failed: %v", err)
+			}
 			if tt.vde == nil {
 				tt.vde = vdiffenv.vde
 			}
@@ -94,6 +159,10 @@ func TestPerformVDiffAction(t *testing.T) {
 			if tt.want != nil && !reflect.DeepEqual(got, tt.want) {
 				t.Errorf("Engine.PerformVDiffAction() = %v, want %v", got, tt.want)
 			}
+			if tt.postFunc != nil {
+				err := tt.postFunc()
+				require.NoError(t, err, "post function failed: %v", err)
+			}
 		})
 	}
 }