Merge pull request #67366 from irfansharif/backport21.1-67315

release-21.1: migration,keys: fix incorrect range descriptor iteration

pkg/keys/keys.go

@@ -657,6 +657,11 @@ func UserKey(key roachpb.RKey) roachpb.RKey {
 	return buf
 }
 
+// InMeta1 returns true iff a key is in the meta1 range (which includes RKeyMin).
+func InMeta1(k roachpb.RKey) bool {
+	return k.Equal(roachpb.RKeyMin) || bytes.HasPrefix(k, MustAddr(Meta1Prefix))
+}
+
 // validateRangeMetaKey validates that the given key is a valid Range Metadata
 // key. This checks only the constraints common to forward and backwards scans:
 // correct prefix and not exceeding KeyMax.

pkg/migration/migrationcluster/BUILD.bazel

@@ -36,6 +36,7 @@ go_test(
     ],
     embed = [":migrationcluster"],
     deps = [
+        "//pkg/keys",
         "//pkg/kv/kvserver",
         "//pkg/migration/nodelivenesstest",
         "//pkg/roachpb",

pkg/migration/migrationcluster/client_test.go

@@ -12,8 +12,10 @@ package migrationcluster_test
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
+	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver"
 	"github.com/cockroachdb/cockroach/pkg/migration/migrationcluster"
 	"github.com/cockroachdb/cockroach/pkg/migration/nodelivenesstest"
@@ -23,43 +25,60 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 )
 
-func TestCluster_IterateRangeDescriptors(t *testing.T) {
+func TestClusterIterateRangeDescriptors(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 
 	ctx := context.Background()
 	const numNodes = 1
 
-	params, _ := tests.CreateTestServerParams()
-	server, _, kvDB := serverutils.StartServer(t, params)
-	defer server.Stopper().Stop(context.Background())
+	for _, splits := range [][]roachpb.Key{
+		{},                                    // no splits
+		{keys.Meta2Prefix},                    // split between meta1 and meta2
+		{keys.SystemPrefix},                   // split after the meta range
+		{keys.Meta2Prefix, keys.SystemPrefix}, // split before and after meta2
+		{keys.RangeMetaKey(roachpb.RKey("middle")).AsRawKey()},                   // split within meta2
+		{keys.Meta2Prefix, keys.RangeMetaKey(roachpb.RKey("middle")).AsRawKey()}, // split at start of and within meta2
+	} {
+		t.Run(fmt.Sprintf("with-splits-at=%s", splits), func(t *testing.T) {
+			params, _ := tests.CreateTestServerParams()
+			server, _, kvDB := serverutils.StartServer(t, params)
+			defer server.Stopper().Stop(context.Background())
 
-	var numRanges int
-	if err := server.GetStores().(*kvserver.Stores).VisitStores(func(s *kvserver.Store) error {
-		numRanges = s.ReplicaCount()
-		return nil
-	}); err != nil {
-		t.Fatal(err)
-	}
+			for _, split := range splits {
+				if _, _, err := server.SplitRange(split); err != nil {
+					t.Fatal(err)
+				}
+			}
+
+			var numRanges int
+			if err := server.GetStores().(*kvserver.Stores).VisitStores(func(s *kvserver.Store) error {
+				numRanges = s.ReplicaCount()
+				return nil
+			}); err != nil {
+				t.Fatal(err)
+			}
 
-	c := nodelivenesstest.New(numNodes)
-	h := migrationcluster.New(migrationcluster.ClusterConfig{
-		NodeLiveness: c,
-		Dialer:       migrationcluster.NoopDialer{},
-		DB:           kvDB,
-	})
+			c := nodelivenesstest.New(numNodes)
+			h := migrationcluster.New(migrationcluster.ClusterConfig{
+				NodeLiveness: c,
+				Dialer:       migrationcluster.NoopDialer{},
+				DB:           kvDB,
+			})
 
-	for _, blockSize := range []int{1, 5, 10, 50} {
-		var numDescs int
-		init := func() { numDescs = 0 }
-		if err := h.IterateRangeDescriptors(ctx, blockSize, init, func(descriptors ...roachpb.RangeDescriptor) error {
-			numDescs += len(descriptors)
-			return nil
-		}); err != nil {
-			t.Fatal(err)
-		}
+			for _, blockSize := range []int{1, 5, 10, 50} {
+				var numDescs int
+				init := func() { numDescs = 0 }
+				if err := h.IterateRangeDescriptors(ctx, blockSize, init, func(descriptors ...roachpb.RangeDescriptor) error {
+					numDescs += len(descriptors)
+					return nil
+				}); err != nil {
+					t.Fatal(err)
+				}
 
-		if numDescs != numRanges {
-			t.Fatalf("expected to find %d ranges, found %d", numRanges+1, numDescs)
-		}
+				if numDescs != numRanges {
+					t.Fatalf("expected to find %d ranges, found %d", numRanges, numDescs)
+				}
+			}
+		})
 	}
 }

pkg/migration/migrationcluster/cluster.go

@@ -134,37 +134,52 @@ func (c *Cluster) ForEveryNode(
 func (c *Cluster) IterateRangeDescriptors(
 	ctx context.Context, blockSize int, init func(), fn func(...roachpb.RangeDescriptor) error,
 ) error {
-	if err := c.DB().Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
+	return c.DB().Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
 		// Inform the caller that we're starting a fresh attempt to page in
 		// range descriptors.
 		init()
 
-		// Iterate through meta2 to pull out all the range descriptors.
-		return txn.Iterate(ctx, keys.Meta2Prefix, keys.MetaMax, blockSize,
+		// Iterate through meta{1,2} to pull out all the range descriptors.
+		var lastRangeIDInMeta1 roachpb.RangeID
+		return txn.Iterate(ctx, keys.MetaMin, keys.MetaMax, blockSize,
 			func(rows []kv.KeyValue) error {
-				descriptors := make([]roachpb.RangeDescriptor, len(rows))
-				for i, row := range rows {
-					if err := row.ValueProto(&descriptors[i]); err != nil {
-						return errors.Wrapf(err,
-							"unable to unmarshal range descriptor from %s",
-							row.Key,
-						)
+				descriptors := make([]roachpb.RangeDescriptor, 0, len(rows))
+				var desc roachpb.RangeDescriptor
+				for _, row := range rows {
+					err := row.ValueProto(&desc)
+					if err != nil {
+						return errors.Wrapf(err, "unable to unmarshal range descriptor from %s", row.Key)
+					}
+
+					// In small enough clusters it's possible for the same range
+					// descriptor to be stored in both meta1 and meta2. This
+					// happens when some range spans both the meta and the user
+					// keyspace. Consider when r1 is [/Min,
+					// /System/NodeLiveness); we'll store the range descriptor
+					// in both /Meta2/<r1.EndKey> and in /Meta1/KeyMax[1].
+					//
+					// As part of iterator we'll de-duplicate this descriptor
+					// away by checking whether we've seen it before in meta1.
+					// Since we're scanning over the meta range in sorted
+					// order, it's enough to check against the last range
+					// descriptor we've seen in meta1.
+					//
+					// [1]: See kvserver.rangeAddressing.
+					if desc.RangeID == lastRangeIDInMeta1 {
+						continue
+					}
+
+					descriptors = append(descriptors, desc)
+					if keys.InMeta1(keys.RangeMetaKey(desc.StartKey)) {
+						lastRangeIDInMeta1 = desc.RangeID
 					}
 				}
 
 				// Invoke fn with the current chunk (of size ~blockSize) of
 				// range descriptors.
-				if err := fn(descriptors...); err != nil {
-					return err
-				}
-
-				return nil
+				return fn(descriptors...)
 			})
-	}); err != nil {
-		return err
-	}
-
-	return nil
+	})
 }
 
 // DB provides exposes the underlying *kv.DB instance.

pkg/server/testserver.go

@@ -1149,10 +1149,6 @@ func (ts *TestServer) SplitRangeWithExpiration(
 	splitKey roachpb.Key, expirationTime hlc.Timestamp,
 ) (roachpb.RangeDescriptor, roachpb.RangeDescriptor, error) {
 	ctx := context.Background()
-	splitRKey, err := keys.Addr(splitKey)
-	if err != nil {
-		return roachpb.RangeDescriptor{}, roachpb.RangeDescriptor{}, err
-	}
 	splitReq := roachpb.AdminSplitRequest{
 		RequestHeader: roachpb.RequestHeader{
 			Key: splitKey,
@@ -1181,40 +1177,27 @@ func (ts *TestServer) SplitRangeWithExpiration(
 	// non-retryable failures and then wrapped when the full transaction fails.
 	var wrappedMsg string
 	if err := ts.DB().Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
-		scanMeta := func(key roachpb.RKey, reverse bool) (desc roachpb.RangeDescriptor, err error) {
-			var kvs []kv.KeyValue
-			if reverse {
-				// Find the last range that ends at or before key.
-				kvs, err = txn.ReverseScan(
-					ctx, keys.Meta2Prefix, keys.RangeMetaKey(key.Next()), 1, /* one result */
-				)
-			} else {
-				// Find the first range that ends after key.
-				kvs, err = txn.Scan(
-					ctx, keys.RangeMetaKey(key.Next()), keys.Meta2Prefix.PrefixEnd(), 1, /* one result */
-				)
-			}
-			if err != nil {
-				return desc, err
-			}
-			if len(kvs) != 1 {
-				return desc, fmt.Errorf("expected 1 result, got %d", len(kvs))
-			}
-			err = kvs[0].ValueProto(&desc)
-			return desc, err
-		}
-
-		rightRangeDesc, err = scanMeta(splitRKey, false /* reverse */)
+		leftRangeDesc, rightRangeDesc = roachpb.RangeDescriptor{}, roachpb.RangeDescriptor{}
+
+		// Discovering the RHS is easy, but the LHS is more difficult. The easiest way to
+		// get both in one operation is to do a reverse range lookup on splitKey.Next();
+		// we need the .Next() because in reverse mode, the end key of a range is inclusive,
+		// i.e. looking up key `c` will match range [a,c), not [c, d).
+		// The result will be the right descriptor, and the first prefetched result will
+		// be the left neighbor, i.e. the resulting left hand side of the split.
+		rs, more, err := kv.RangeLookup(ctx, txn, splitKey.Next(), roachpb.CONSISTENT, 1, true /* reverse */)
 		if err != nil {
-			wrappedMsg = "could not look up right-hand side descriptor"
 			return err
 		}
-
-		leftRangeDesc, err = scanMeta(splitRKey, true /* reverse */)
-		if err != nil {
-			wrappedMsg = "could not look up left-hand side descriptor"
-			return err
+		if len(rs) == 0 {
+			// This is a bug.
+			return errors.AssertionFailedf("no descriptor found for key %s", splitKey)
+		}
+		if len(more) == 0 {
+			return errors.Errorf("looking up post-split descriptor returned first range: %+v", rs[0])
 		}
+		leftRangeDesc = more[0]
+		rightRangeDesc = rs[0]
 
 		if !leftRangeDesc.EndKey.Equal(rightRangeDesc.StartKey) {
 			return errors.Errorf(