github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/kv/kvserver/reports/reporter_test.go (about)

     1  // Copyright 2019 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package reports
    12  
    13  import (
    14  	"context"
    15  	gosql "database/sql"
    16  	"fmt"
    17  	"testing"
    18  
    19  	"github.com/cockroachdb/cockroach/pkg/base"
    20  	"github.com/cockroachdb/cockroach/pkg/config"
    21  	"github.com/cockroachdb/cockroach/pkg/config/zonepb"
    22  	"github.com/cockroachdb/cockroach/pkg/gossip"
    23  	"github.com/cockroachdb/cockroach/pkg/kv/kvserver"
    24  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    25  	"github.com/cockroachdb/cockroach/pkg/testutils"
    26  	"github.com/cockroachdb/cockroach/pkg/testutils/keysutils"
    27  	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
    28  	"github.com/cockroachdb/cockroach/pkg/util"
    29  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    30  	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
    31  	"github.com/cockroachdb/cockroach/pkg/util/uuid"
    32  	"github.com/cockroachdb/errors"
    33  	"github.com/stretchr/testify/require"
    34  )
    35  
    36  // Test the constraint conformance report in a real cluster.
    37  func TestConstraintConformanceReportIntegration(t *testing.T) {
    38  	defer leaktest.AfterTest(t)()
    39  	if testing.Short() {
    40  		// This test takes seconds because of replication vagaries.
    41  		t.Skip("short flag")
    42  	}
    43  	if testutils.NightlyStress() && util.RaceEnabled {
    44  		// Under stressrace, replication changes seem to hit 1m deadline errors and
    45  		// don't make progress.
    46  		t.Skip("test too slow for stressrace")
    47  	}
    48  
    49  	ctx := context.Background()
    50  	tc := serverutils.StartTestCluster(t, 5, base.TestClusterArgs{
    51  		ServerArgsPerNode: map[int]base.TestServerArgs{
    52  			0: {Locality: roachpb.Locality{Tiers: []roachpb.Tier{{Key: "region", Value: "r1"}}}},
    53  			1: {Locality: roachpb.Locality{Tiers: []roachpb.Tier{{Key: "region", Value: "r1"}}}},
    54  			2: {Locality: roachpb.Locality{Tiers: []roachpb.Tier{{Key: "region", Value: "r2"}}}},
    55  			3: {Locality: roachpb.Locality{Tiers: []roachpb.Tier{{Key: "region", Value: "r2"}}}},
    56  			4: {Locality: roachpb.Locality{Tiers: []roachpb.Tier{{Key: "region", Value: "r2"}}}},
    57  		},
    58  	})
    59  	defer tc.Stopper().Stop(ctx)
    60  
    61  	db := tc.ServerConn(0)
    62  	// Speed up the generation of the
    63  	_, err := db.Exec("set cluster setting kv.replication_reports.interval = '1ms'")
    64  	require.NoError(t, err)
    65  
    66  	// Create a table and a zone config for it.
    67  	// The zone will be configured with a constraints that can't be satisfied
    68  	// because there are not enough nodes in the requested region.
    69  	_, err = db.Exec("create table t(x int primary key); " +
    70  		"alter table t configure zone using constraints='[+region=r1]'")
    71  	require.NoError(t, err)
    72  
    73  	// Get the id of the newly created zone.
    74  	r := db.QueryRow("select zone_id from crdb_internal.zones where table_name = 't'")
    75  	var zoneID int
    76  	require.NoError(t, r.Scan(&zoneID))
    77  
    78  	// Wait for the violation to be detected.
    79  	testutils.SucceedsSoon(t, func() error {
    80  		r := db.QueryRow(
    81  			"select violating_ranges from system.replication_constraint_stats where zone_id = $1",
    82  			zoneID)
    83  		var numViolations int
    84  		if err := r.Scan(&numViolations); err != nil {
    85  			return err
    86  		}
    87  		if numViolations == 0 {
    88  			return fmt.Errorf("violation not detected yet")
    89  		}
    90  		return nil
    91  	})
    92  
    93  	// Now change the constraint asking for t to be placed in r2. This time it can be satisfied.
    94  	_, err = db.Exec("alter table t configure zone using constraints='[+region=r2]'")
    95  	require.NoError(t, err)
    96  
    97  	// Wait for the violation to clear.
    98  	testutils.SucceedsSoon(t, func() error {
    99  		// Kick the replication queues, given that our rebalancing is finicky.
   100  		for i := 0; i < tc.NumServers(); i++ {
   101  			if err := tc.Server(i).GetStores().(*kvserver.Stores).VisitStores(func(s *kvserver.Store) error {
   102  				return s.ForceReplicationScanAndProcess()
   103  			}); err != nil {
   104  				t.Fatal(err)
   105  			}
   106  		}
   107  		r := db.QueryRow(
   108  			"select violating_ranges from system.replication_constraint_stats where zone_id = $1",
   109  			zoneID)
   110  		var numViolations int
   111  		if err := r.Scan(&numViolations); err != nil {
   112  			return err
   113  		}
   114  		if numViolations > 0 {
   115  			return fmt.Errorf("still reporting violations")
   116  		}
   117  		return nil
   118  	})
   119  }
   120  
   121  // Test the critical localities report in a real cluster.
   122  func TestCriticalLocalitiesReportIntegration(t *testing.T) {
   123  	defer leaktest.AfterTest(t)()
   124  	ctx := context.Background()
   125  	// 2 regions, 3 dcs per region.
   126  	tc := serverutils.StartTestCluster(t, 6, base.TestClusterArgs{
   127  		// We're going to do our own replication.
   128  		// All the system ranges will start with a single replica on node 1.
   129  		ReplicationMode: base.ReplicationManual,
   130  		ServerArgsPerNode: map[int]base.TestServerArgs{
   131  			0: {
   132  				Locality: roachpb.Locality{Tiers: []roachpb.Tier{
   133  					{Key: "region", Value: "r1"}, {Key: "dc", Value: "dc1"}},
   134  				},
   135  			},
   136  			1: {
   137  				Locality: roachpb.Locality{Tiers: []roachpb.Tier{
   138  					{Key: "region", Value: "r1"}, {Key: "dc", Value: "dc2"}},
   139  				},
   140  			},
   141  			2: {
   142  				Locality: roachpb.Locality{Tiers: []roachpb.Tier{
   143  					{Key: "region", Value: "r1"}, {Key: "dc", Value: "dc3"}},
   144  				},
   145  			},
   146  			3: {
   147  				Locality: roachpb.Locality{Tiers: []roachpb.Tier{
   148  					{Key: "region", Value: "r2"}, {Key: "dc", Value: "dc4"}},
   149  				},
   150  			},
   151  			4: {
   152  				Locality: roachpb.Locality{Tiers: []roachpb.Tier{
   153  					{Key: "region", Value: "r2"}, {Key: "dc", Value: "dc5"}},
   154  				},
   155  			},
   156  			5: {
   157  				Locality: roachpb.Locality{Tiers: []roachpb.Tier{
   158  					{Key: "region", Value: "r2"}, {Key: "dc", Value: "dc6"}},
   159  				},
   160  			},
   161  		},
   162  	})
   163  	defer tc.Stopper().Stop(ctx)
   164  
   165  	db := tc.ServerConn(0)
   166  	// Speed up the generation of the reports.
   167  	_, err := db.Exec("set cluster setting kv.replication_reports.interval = '1ms'")
   168  	require.NoError(t, err)
   169  
   170  	// Since we're using ReplicationManual, all the ranges will start with a
   171  	// single replica on node 1. So, the node's dc and the node's region are
   172  	// critical. Let's verify that.
   173  
   174  	// Collect all the zones that exist at cluster bootstrap.
   175  	systemZoneIDs := make([]int, 0, 10)
   176  	systemZones := make([]zonepb.ZoneConfig, 0, 10)
   177  	{
   178  		rows, err := db.Query("select id, config from system.zones")
   179  		require.NoError(t, err)
   180  		for rows.Next() {
   181  			var zoneID int
   182  			var buf []byte
   183  			cfg := zonepb.ZoneConfig{}
   184  			require.NoError(t, rows.Scan(&zoneID, &buf))
   185  			require.NoError(t, protoutil.Unmarshal(buf, &cfg))
   186  			systemZoneIDs = append(systemZoneIDs, zoneID)
   187  			systemZones = append(systemZones, cfg)
   188  		}
   189  		require.NoError(t, rows.Err())
   190  	}
   191  	require.Greater(t, len(systemZoneIDs), 0, "expected some system zones, got none")
   192  	// Remove the entries in systemZoneIDs that don't get critical locality reports.
   193  	i := 0
   194  	for j, zid := range systemZoneIDs {
   195  		if zoneChangesReplication(&systemZones[j]) {
   196  			systemZoneIDs[i] = zid
   197  			i++
   198  		}
   199  	}
   200  	systemZoneIDs = systemZoneIDs[:i]
   201  
   202  	expCritLoc := []string{"region=r1", "region=r1,dc=dc1"}
   203  
   204  	// Wait for the report to be generated.
   205  	{
   206  		var rowCount int
   207  		testutils.SucceedsSoon(t, func() error {
   208  			r := db.QueryRow("select count(1) from system.replication_critical_localities")
   209  			require.NoError(t, r.Scan(&rowCount))
   210  			if rowCount == 0 {
   211  				return fmt.Errorf("no report yet")
   212  			}
   213  			return nil
   214  		})
   215  		require.Equal(t, 2*len(systemZoneIDs), rowCount)
   216  	}
   217  
   218  	// Check that we have all the expected rows.
   219  	for _, zid := range systemZoneIDs {
   220  		for _, s := range expCritLoc {
   221  			r := db.QueryRow(
   222  				"select at_risk_ranges from system.replication_critical_localities "+
   223  					"where zone_id=$1 and locality=$2",
   224  				zid, s)
   225  			var numRanges int
   226  			msg := fmt.Sprintf("zone_id: %d, locality: %s", zid, s)
   227  			require.NoError(t, r.Scan(&numRanges), msg)
   228  			require.NotEqual(t, 0, numRanges, msg)
   229  		}
   230  	}
   231  
   232  	// Now create a table and a zone for it. At first n1 should be critical for it.
   233  	// Then we'll upreplicate it in different ways.
   234  
   235  	// Create a table with a dummy zone config. Configuring the zone is useful
   236  	// only for creating the zone; we don't actually care about the configuration.
   237  	// Also do a split by hand. With manual replication, we're not getting the
   238  	// split for the table automatically.
   239  	_, err = db.Exec("create table t(x int primary key); " +
   240  		"alter table t configure zone using num_replicas=3; " +
   241  		"alter table t split at values (0);")
   242  	require.NoError(t, err)
   243  	// Get the id of the newly created zone.
   244  	r := db.QueryRow("select zone_id from crdb_internal.zones where table_name = 't'")
   245  	var zoneID int
   246  	require.NoError(t, r.Scan(&zoneID))
   247  
   248  	// Check initial conditions.
   249  	require.NoError(t, checkCritical(db, zoneID, "region=r1", "region=r1,dc=dc1"))
   250  
   251  	// Upreplicate to 2 dcs. Now they're both critical.
   252  	_, err = db.Exec("ALTER TABLE t EXPERIMENTAL_RELOCATE VALUES (ARRAY[1,2], 1)")
   253  	require.NoError(t, err)
   254  	require.NoError(t, checkCritical(db, zoneID, "region=r1", "region=r1,dc=dc1", "region=r1,dc=dc2"))
   255  
   256  	// Upreplicate to one more dc. Now no dc is critical, only the region.
   257  	_, err = db.Exec("ALTER TABLE t EXPERIMENTAL_RELOCATE VALUES (ARRAY[1,2,3], 1)")
   258  	require.NoError(t, err)
   259  	require.NoError(t, checkCritical(db, zoneID, "region=r1"))
   260  
   261  	// Move two replicas to the other region. Now that region is critical.
   262  	_, err = db.Exec("ALTER TABLE t EXPERIMENTAL_RELOCATE VALUES (ARRAY[1,4,5], 1)")
   263  	require.NoError(t, err)
   264  	require.NoError(t, checkCritical(db, zoneID, "region=r2"))
   265  }
   266  
   267  func checkCritical(db *gosql.DB, zoneID int, locs ...string) error {
   268  	return testutils.SucceedsSoonError(func() error {
   269  		rows, err := db.Query(
   270  			"select locality, at_risk_ranges from system.replication_critical_localities "+
   271  				"where zone_id=$1", zoneID)
   272  		if err != nil {
   273  			return err
   274  		}
   275  		critical := make(map[string]struct{})
   276  		for rows.Next() {
   277  			var numRanges int
   278  			var loc string
   279  			err := rows.Scan(&loc, &numRanges)
   280  			if err != nil {
   281  				return err
   282  			}
   283  			if numRanges == 0 {
   284  				return fmt.Errorf("expected ranges_at_risk for %s", loc)
   285  			}
   286  			critical[loc] = struct{}{}
   287  		}
   288  		if err := rows.Err(); err != nil {
   289  			return err
   290  		}
   291  		if len(locs) != len(critical) {
   292  			return fmt.Errorf("expected critical: %s, got: %s", locs, critical)
   293  		}
   294  		for _, l := range locs {
   295  			if _, ok := critical[l]; !ok {
   296  				return fmt.Errorf("missing critical locality: %s", l)
   297  			}
   298  		}
   299  		return nil
   300  	})
   301  }
   302  
   303  // Test the replication status report in a real cluster.
   304  func TestReplicationStatusReportIntegration(t *testing.T) {
   305  	defer leaktest.AfterTest(t)()
   306  	ctx := context.Background()
   307  	tc := serverutils.StartTestCluster(t, 4, base.TestClusterArgs{
   308  		// We're going to do our own replication.
   309  		// All the system ranges will start with a single replica on node 1.
   310  		ReplicationMode: base.ReplicationManual,
   311  	})
   312  	defer tc.Stopper().Stop(ctx)
   313  
   314  	db := tc.ServerConn(0)
   315  	// Speed up the generation of the
   316  	_, err := db.Exec("set cluster setting kv.replication_reports.interval = '1ms'")
   317  	require.NoError(t, err)
   318  
   319  	// Create a table with a dummy zone config. Configuring the zone is useful
   320  	// only for creating the zone; we don't actually care about the configuration.
   321  	// Also do a split by hand. With manual replication, we're not getting the
   322  	// split for the table automatically.
   323  	_, err = db.Exec("create table t(x int primary key); " +
   324  		"alter table t configure zone using num_replicas=3; " +
   325  		"alter table t split at values (0);")
   326  	require.NoError(t, err)
   327  	// Get the id of the newly created zone.
   328  	r := db.QueryRow("select zone_id from crdb_internal.zones where table_name = 't'")
   329  	var zoneID int
   330  	require.NoError(t, r.Scan(&zoneID))
   331  
   332  	// Upreplicate the range.
   333  	_, err = db.Exec("ALTER TABLE t EXPERIMENTAL_RELOCATE VALUES (ARRAY[1,2,3], 1)")
   334  	require.NoError(t, err)
   335  	require.NoError(t, checkZoneReplication(db, zoneID, 1, 0, 0, 0))
   336  
   337  	// Over-replicate.
   338  	_, err = db.Exec("ALTER TABLE t EXPERIMENTAL_RELOCATE VALUES (ARRAY[1,2,3,4], 1)")
   339  	require.NoError(t, err)
   340  	require.NoError(t, checkZoneReplication(db, zoneID, 1, 0, 1, 0))
   341  
   342  	// TODO(andrei): I'd like to downreplicate to one replica and then stop that
   343  	// node and check that the range is counter us "unavailable", but stopping a
   344  	// node makes the report generation simply block sometimes trying to scan
   345  	// Meta2. I believe I believe it's due to #40529.
   346  	// Once stopping a node works, next thing is to start it up again.
   347  	// Take inspiration from replica_learner_test.go.
   348  
   349  	//// Down-replicate to one node and then kill the node. Check that the range becomes unavailable.
   350  	//_, err = db.Exec("ALTER TABLE t EXPERIMENTAL_RELOCATE VALUES (ARRAY[4], 1)")
   351  	//require.NoError(t, err)
   352  	//tc.StopServer(3)
   353  	//require.NoError(t, checkZoneReplication(db, zoneID, 1, 1, 0, 1))
   354  }
   355  
   356  func checkZoneReplication(db *gosql.DB, zoneID, total, under, over, unavailable int) error {
   357  	return testutils.SucceedsSoonError(func() error {
   358  		r := db.QueryRow(
   359  			"select total_ranges, under_replicated_ranges, over_replicated_ranges, "+
   360  				"unavailable_ranges from system.replication_stats where zone_id=$1",
   361  			zoneID)
   362  		var gotTotal, gotUnder, gotOver, gotUnavailable int
   363  		if err := r.Scan(&gotTotal, &gotUnder, &gotOver, &gotUnavailable); err != nil {
   364  			return err
   365  		}
   366  		if total != gotTotal {
   367  			return fmt.Errorf("expected total: %d, got: %d", total, gotTotal)
   368  		}
   369  		if under != gotUnder {
   370  			return fmt.Errorf("expected under: %d, got: %d", total, gotUnder)
   371  		}
   372  		if over != gotOver {
   373  			return fmt.Errorf("expected over: %d, got: %d", over, gotOver)
   374  		}
   375  		if unavailable != gotUnavailable {
   376  			return fmt.Errorf("expected unavailable: %d, got: %d", unavailable, gotUnavailable)
   377  		}
   378  		return nil
   379  	})
   380  }
   381  
   382  func TestMeta2RangeIter(t *testing.T) {
   383  	defer leaktest.AfterTest(t)()
   384  	ctx := context.Background()
   385  	s, _, db := serverutils.StartServer(t, base.TestServerArgs{})
   386  	defer s.Stopper().Stop(ctx)
   387  
   388  	// First make an interator with a large page size and use it to determine the numner of ranges.
   389  	iter := makeMeta2RangeIter(db, 10000 /* batchSize */)
   390  	numRanges := 0
   391  	for {
   392  		rd, err := iter.Next(ctx)
   393  		require.NoError(t, err)
   394  		if rd.RangeID == 0 {
   395  			break
   396  		}
   397  		numRanges++
   398  	}
   399  	require.Greater(t, numRanges, 20, "expected over 20 ranges, got: %d", numRanges)
   400  
   401  	// Now make an interator with a small page size and check that we get just as many ranges.
   402  	iter = makeMeta2RangeIter(db, 2 /* batch size */)
   403  	numRangesPaginated := 0
   404  	for {
   405  		rd, err := iter.Next(ctx)
   406  		require.NoError(t, err)
   407  		if rd.RangeID == 0 {
   408  			break
   409  		}
   410  		numRangesPaginated++
   411  	}
   412  	require.Equal(t, numRanges, numRangesPaginated)
   413  }
   414  
   415  // Test that a retriable error returned from the range iterator is properly
   416  // handled by resetting the report.
   417  func TestRetriableErrorWhenGenerationReport(t *testing.T) {
   418  	defer leaktest.AfterTest(t)()
   419  	ctx := context.Background()
   420  	s, _, db := serverutils.StartServer(t, base.TestServerArgs{})
   421  	defer s.Stopper().Stop(ctx)
   422  
   423  	cfg := s.GossipI().(*gossip.Gossip).GetSystemConfig()
   424  	dummyNodeChecker := func(id roachpb.NodeID) bool { return true }
   425  
   426  	v := makeReplicationStatsVisitor(ctx, cfg, dummyNodeChecker)
   427  	realIter := makeMeta2RangeIter(db, 10000 /* batchSize */)
   428  	require.NoError(t, visitRanges(ctx, &realIter, cfg, &v))
   429  	expReport := v.Report()
   430  	require.Greater(t, len(expReport), 0, "unexpected empty report")
   431  
   432  	realIter = makeMeta2RangeIter(db, 10000 /* batchSize */)
   433  	errorIter := erroryRangeIterator{
   434  		iter:           realIter,
   435  		injectErrAfter: 3,
   436  	}
   437  	v = makeReplicationStatsVisitor(ctx, cfg, func(id roachpb.NodeID) bool { return true })
   438  	require.NoError(t, visitRanges(ctx, &errorIter, cfg, &v))
   439  	require.Greater(t, len(v.report), 0, "unexpected empty report")
   440  	require.Equal(t, expReport, v.report)
   441  }
   442  
   443  type erroryRangeIterator struct {
   444  	iter           meta2RangeIter
   445  	rangesReturned int
   446  	injectErrAfter int
   447  }
   448  
   449  var _ RangeIterator = &erroryRangeIterator{}
   450  
   451  func (it *erroryRangeIterator) Next(ctx context.Context) (roachpb.RangeDescriptor, error) {
   452  	if it.rangesReturned == it.injectErrAfter {
   453  		// Don't inject any more errors.
   454  		it.injectErrAfter = -1
   455  
   456  		var err error
   457  		err = roachpb.NewTransactionRetryWithProtoRefreshError(
   458  			"injected err", uuid.Nil, roachpb.Transaction{})
   459  		// Let's wrap the error to check the unwrapping.
   460  		err = errors.Wrap(err, "dummy wrapper")
   461  		// Feed the error to the underlying iterator to reset it.
   462  		it.iter.handleErr(ctx, err)
   463  		return roachpb.RangeDescriptor{}, err
   464  	}
   465  	it.rangesReturned++
   466  	rd, err := it.iter.Next(ctx)
   467  	return rd, err
   468  }
   469  
   470  func (it *erroryRangeIterator) Close(ctx context.Context) {
   471  	it.iter.Close(ctx)
   472  }
   473  
   474  func TestZoneChecker(t *testing.T) {
   475  	defer leaktest.AfterTest(t)()
   476  	ctx := context.Background()
   477  
   478  	type tc struct {
   479  		split          string
   480  		newZone        bool
   481  		newRootZoneCfg *zonepb.ZoneConfig
   482  		newZoneKey     ZoneKey
   483  	}
   484  	// NB: IDs need to be beyond MaxSystemConfigDescID, otherwise special logic
   485  	// kicks in for mapping keys to zones.
   486  	dbID := 50
   487  	t1ID := 51
   488  	t1 := table{name: "t1",
   489  		partitions: []partition{
   490  			{
   491  				name:  "p1",
   492  				start: []int{100},
   493  				end:   []int{200},
   494  				zone:  &zone{constraints: "[+p1]"},
   495  			},
   496  			{
   497  				name:  "p2",
   498  				start: []int{300},
   499  				end:   []int{400},
   500  				zone:  &zone{constraints: "[+p2]"},
   501  			},
   502  		},
   503  	}
   504  	t1.addPKIdx()
   505  	// Create a table descriptor to be used for creating the zone config.
   506  	t1Desc, err := makeTableDesc(t1, t1ID, dbID)
   507  	require.NoError(t, err)
   508  	t1Zone, err := generateTableZone(t1, t1Desc)
   509  	require.NoError(t, err)
   510  	p1SubzoneIndex := 0
   511  	p2SubzoneIndex := 1
   512  	require.Equal(t, "p1", t1Zone.Subzones[p1SubzoneIndex].PartitionName)
   513  	require.Equal(t, "p2", t1Zone.Subzones[p2SubzoneIndex].PartitionName)
   514  	t1ZoneKey := MakeZoneKey(uint32(t1ID), NoSubzone)
   515  	p1ZoneKey := MakeZoneKey(uint32(t1ID), base.SubzoneIDFromIndex(p1SubzoneIndex))
   516  	p2ZoneKey := MakeZoneKey(uint32(t1ID), base.SubzoneIDFromIndex(p2SubzoneIndex))
   517  
   518  	ranges := []tc{
   519  		{
   520  			split:          "/Table/t1/pk/1",
   521  			newZone:        true,
   522  			newZoneKey:     t1ZoneKey,
   523  			newRootZoneCfg: t1Zone,
   524  		},
   525  		{
   526  			split:   "/Table/t1/pk/2",
   527  			newZone: false,
   528  		},
   529  		{
   530  			// p1's zone
   531  			split:          "/Table/t1/pk/100",
   532  			newZone:        true,
   533  			newZoneKey:     p1ZoneKey,
   534  			newRootZoneCfg: t1Zone,
   535  		},
   536  		{
   537  			split:   "/Table/t1/pk/101",
   538  			newZone: false,
   539  		},
   540  		{
   541  			// Back to t1's zone
   542  			split:          "/Table/t1/pk/200",
   543  			newZone:        true,
   544  			newZoneKey:     t1ZoneKey,
   545  			newRootZoneCfg: t1Zone,
   546  		},
   547  		{
   548  			// p2's zone
   549  			split:          "/Table/t1/pk/305",
   550  			newZone:        true,
   551  			newZoneKey:     p2ZoneKey,
   552  			newRootZoneCfg: t1Zone,
   553  		},
   554  	}
   555  
   556  	splits := make([]split, len(ranges))
   557  	for i := range ranges {
   558  		splits[i].key = ranges[i].split
   559  	}
   560  	keyScanner := keysutils.MakePrettyScannerForNamedTables(
   561  		map[string]int{"t1": t1ID} /* tableNameToID */, nil /* idxNameToID */)
   562  	rngs, err := processSplits(keyScanner, splits)
   563  	require.NoError(t, err)
   564  
   565  	var zc zoneResolver
   566  	for i, tc := range ranges {
   567  		sameZone := zc.checkSameZone(ctx, &rngs[i])
   568  		newZone := !sameZone
   569  		require.Equal(t, tc.newZone, newZone, "failed at: %d (%s)", i, tc.split)
   570  		if newZone {
   571  			objectID, _ := config.DecodeKeyIntoZoneIDAndSuffix(rngs[i].StartKey)
   572  			zc.setZone(objectID, tc.newZoneKey, tc.newRootZoneCfg)
   573  		}
   574  	}
   575  }
   576  
   577  // TestRangeIteration checks that visitRanges() correctly informs range
   578  // visitors whether ranges fall in the same zone vs a new zone.
   579  func TestRangeIteration(t *testing.T) {
   580  	defer leaktest.AfterTest(t)()
   581  	ctx := context.Background()
   582  
   583  	schema := baseReportTestCase{
   584  		schema: []database{{
   585  			name: "db1",
   586  			zone: &zone{
   587  				replicas: 3,
   588  			},
   589  			tables: []table{
   590  				{
   591  					name: "t1",
   592  					partitions: []partition{
   593  						{
   594  							name:  "p1",
   595  							start: []int{100},
   596  							end:   []int{200},
   597  							zone:  &zone{},
   598  						},
   599  						{
   600  							name:  "p2",
   601  							start: []int{200},
   602  							end:   []int{300},
   603  							zone:  &zone{},
   604  						},
   605  					},
   606  				},
   607  				{
   608  					name: "t2",
   609  				},
   610  			},
   611  		},
   612  		},
   613  		splits: []split{
   614  			{key: "/Table/t1/pk/1"},
   615  			{key: "/Table/t1/pk/2"},
   616  			{key: "/Table/t1/pk/100"},
   617  			{key: "/Table/t1/pk/101"},
   618  			{key: "/Table/t1/pk/200"},
   619  			{key: "/Table/t1/pk/305"},
   620  			{key: "/Table/t2/pk/1"},
   621  		},
   622  		defaultZone: zone{},
   623  	}
   624  
   625  	compiled, err := compileTestCase(schema)
   626  	require.NoError(t, err)
   627  	v := recordingRangeVisitor{}
   628  	require.NoError(t, visitRanges(ctx, &compiled.iter, compiled.cfg, &v))
   629  
   630  	type entry struct {
   631  		newZone bool
   632  		key     string
   633  	}
   634  	exp := []entry{
   635  		{newZone: true, key: "/Table/51/1/1"},
   636  		{newZone: false, key: "/Table/51/1/2"},
   637  		{newZone: true, key: "/Table/51/1/100"},
   638  		{newZone: false, key: "/Table/51/1/101"},
   639  		{newZone: true, key: "/Table/51/1/200"},
   640  		{newZone: true, key: "/Table/51/1/305"},
   641  		{newZone: true, key: "/Table/52/1/1"},
   642  	}
   643  	got := make([]entry, len(v.rngs))
   644  	for i, r := range v.rngs {
   645  		got[i].newZone = r.newZone
   646  		got[i].key = r.rng.StartKey.String()
   647  	}
   648  	require.Equal(t, exp, got)
   649  }
   650  
   651  type recordingRangeVisitor struct {
   652  	rngs []visitorEntry
   653  }
   654  
   655  var _ rangeVisitor = &recordingRangeVisitor{}
   656  
   657  func (r *recordingRangeVisitor) visitNewZone(
   658  	_ context.Context, rng *roachpb.RangeDescriptor,
   659  ) error {
   660  	r.rngs = append(r.rngs, visitorEntry{newZone: true, rng: *rng})
   661  	return nil
   662  }
   663  
   664  func (r *recordingRangeVisitor) visitSameZone(_ context.Context, rng *roachpb.RangeDescriptor) {
   665  	r.rngs = append(r.rngs, visitorEntry{newZone: false, rng: *rng})
   666  }
   667  
   668  func (r *recordingRangeVisitor) failed() bool {
   669  	return false
   670  }
   671  
   672  func (r *recordingRangeVisitor) reset(ctx context.Context) {
   673  	r.rngs = nil
   674  }
   675  
   676  type visitorEntry struct {
   677  	newZone bool
   678  	rng     roachpb.RangeDescriptor
   679  }