kvserver: clear RHS when learning about a split through a snapshot

Previously, it was possible for us to leak keyspans on receiving a
narrower snapshot. This could happen if the end key of the keyspace
covered by the original descriptor + any subsumed descriptors was less
than the end key of the snapshot. This hazard was shown in #148707.

Practically speaking, the simplest case that's fixed here corresponds
to a replica learning about a split through the snapshot. The more
general case involves the replica learning about a series of merges
and at least one split through the snapshot.

Closes #73462

Release note: None

Author: arulajmani

pkg/kv/kvserver/client_split_test.go

@@ -4566,8 +4566,6 @@ func TestStoreRangeSplitRaftSnapshotAfterRHSRebalanced(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
 
-	skip.WithIssue(t, 73462)
-
 	ctx := context.Background()
 	// Start a 5 node cluster.
 	tc := testcluster.StartTestCluster(t, 5, base.TestClusterArgs{

pkg/kv/kvserver/replica_raftstorage.go

@@ -7,6 +7,7 @@ package kvserver
 
 import (
 	"context"
+	"slices"
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
@@ -24,6 +25,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/storage"
 	"github.com/cockroachdb/cockroach/pkg/storage/fs"
 	"github.com/cockroachdb/cockroach/pkg/util"
+	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/humanizeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -587,6 +589,13 @@ func (r *Replica) applySnapshotRaftMuLocked(
 		subsume = append(subsume, destroyReplicaInfo{id: sr.ID(), desc: srDesc})
 	}
 
+	// NB: subsumedDescs in snapWriteBuilder must be sorted by start key. This
+	// should be the case, by construction, but add a test-only assertion just in
+	// case this ever changes.
+	testingAssert(slices.IsSortedFunc(subsume, func(a, b destroyReplicaInfo) int {
+		return a.desc.StartKey.Compare(b.desc.StartKey)
+	}), "subsumedDescs must be sorted by start key")
+
 	sb := snapWriteBuilder{
 		id: r.ID(),
 
@@ -597,6 +606,7 @@ func (r *Replica) applySnapshotRaftMuLocked(
 		truncState: truncState,
 		hardState:  hs,
 		desc:       desc,
+		origDesc:   r.shMu.state.Desc,
 		subsume:    subsume,
 
 		cleared: inSnap.clearedSpans,
@@ -826,3 +836,9 @@ func (r *Replica) clearSubsumedReplicaInMemoryData(
 	}
 	return phs, nil
 }
+
+func testingAssert(cond bool, msg string) {
+	if buildutil.CrdbTestBuild && !cond {
+		panic(msg)
+	}
+}

pkg/kv/kvserver/snapshot_apply_prepare.go

@@ -39,8 +39,11 @@ type snapWriteBuilder struct {
 
 	truncState kvserverpb.RaftTruncatedState
 	hardState  raftpb.HardState
-	desc       *roachpb.RangeDescriptor
-	subsume    []destroyReplicaInfo
+	desc       *roachpb.RangeDescriptor // corresponds to the range descriptor in the snapshot
+	origDesc   *roachpb.RangeDescriptor // pre-snapshot range descriptor
+	// NB: subsume, if set, must be in sorted (by destroyReplicaInfo.desc start
+	// key) order.
+	subsume []destroyReplicaInfo
 
 	// cleared contains the spans that this snapshot application clears before
 	// writing new state on top.
@@ -54,7 +57,12 @@ func (s *snapWriteBuilder) prepareSnapApply(ctx context.Context) error {
 		return err
 	}
 	_ = applySnapshotTODO // 1.2 + 2.1 + 2.2 + 2.3 (diff) + 3.2
-	return s.clearSubsumedReplicaDiskData(ctx)
+	if err := s.clearSubsumedReplicaDiskData(ctx); err != nil {
+		return err
+	}
+
+	_ = applySnapshotTODO // 2.3 (split)
+	return s.clearResidualDataOnNarrowSnapshot(ctx)
 }
 
 // rewriteRaftState clears and rewrites the unreplicated rangeID-local key space
@@ -104,25 +112,50 @@ func (s *snapWriteBuilder) rewriteRaftState(ctx context.Context, w storage.Write
 // the Reader reflects the latest I/O each of the subsumed replicas has done
 // (i.e. Reader was instantiated after all raftMu were acquired).
 //
-// NB: does nothing if subsumedDescs is empty.
+// NB: does nothing if s.subsumedDescs is empty.
 func (s *snapWriteBuilder) clearSubsumedReplicaDiskData(ctx context.Context) error {
+	if len(s.subsume) == 0 {
+		return nil // no subsumed replicas to speak of; early return
+	}
+	// NB: The snapshot must never subsume a replica that extends the range of the
+	// replica to the left. This is because splits and merges (the only operation
+	// that change the key bounds) always leave the start key intact. Extending to
+	// the left implies that either we merged "to the left" (we don't), or that
+	// we're applying a snapshot for another range (we don't do that either).
+	// Something is severely wrong for this to happen, so perform a sanity check.
+	if s.subsume[0].desc.StartKey.Compare(s.desc.StartKey) < 0 { // subsumedDescs are sorted by StartKey
+		log.Dev.Fatalf(ctx,
+			"subsuming replica to our left; subsumed desc start key: %v; snapshot desc start key %v",
+			s.subsume[0].desc.StartKey, s.desc.StartKey,
+		)
+	}
+
+	// In the common case, the subsumed replicas' end key does not extend beyond
+	// the snapshot end key (sn <= b):
+	//
+	//	   subsumed: [a---s1---...---sn)
+	//	   snapshot: [a---------------b)
+	//	or snapshot: [a-------------------b)
+	//
+	// In this case, we do not need to clear the range-local, lock table, and user
+	// keys owned by subsumed replicas, since clearing [a, b) will cover it. We
+	// only need to clear the per-replica RangeID-local key spans here.
+	//
+	// This leaves the only other case, where the subsumed replicas extend past
+	// the snapshot's end key (s[n-1] < b < sn):
+	//
+	//	subsumed: [a---s1---...---s[n-1]---sn)
+	//	snapshot: [a--------------------b)
+	//
+	// This is only possible if we're not only learning about merges through the
+	// snapshot, but also a split -- that's the only way the bounds of the
+	// snapshot can be narrower than the bounds of all the subsumed replicas. In
+	// this case, we do need to clear range-local, lock table, and user keys in
+	// the span [b, sn). We do this in
+	// clearResidualDataOnNarrowSnapshot, not here.
+
 	// TODO(sep-raft-log): need different readers for raft and state engine.
 	reader := storage.Reader(s.todoEng)
-
-	// NB: we don't clear RangeID local key spans here. That happens
-	// via the call to DestroyReplica.
-	getKeySpans := func(d *roachpb.RangeDescriptor) []roachpb.Span {
-		return rditer.Select(d.RangeID, rditer.SelectOpts{
-			Ranged: rditer.SelectRangedOptions{
-				RSpan:      d.RSpan(),
-				SystemKeys: true,
-				UserKeys:   true,
-				LockTable:  true,
-			},
-		})
-	}
-	keySpans := getKeySpans(s.desc)
-	totalKeySpans := append([]roachpb.Span(nil), keySpans...)
 	for _, sub := range s.subsume {
 		// We have to create an SST for the subsumed replica's range-id local keys.
 		if err := s.writeSST(ctx, func(ctx context.Context, w storage.Writer) error {
@@ -138,83 +171,89 @@ func (s *snapWriteBuilder) clearSubsumedReplicaDiskData(ctx context.Context) err
 				ReplicatedByRangeID:   opts.ClearReplicatedByRangeID,
 				UnreplicatedByRangeID: opts.ClearUnreplicatedByRangeID,
 			})...)
+			// NB: Actually clear RangeID local key spans.
 			return kvstorage.DestroyReplica(ctx, sub.id, reader, w, mergedTombstoneReplicaID, opts)
 		}); err != nil {
 			return err
 		}
-
-		srKeySpans := getKeySpans(sub.desc)
-		// Compute the total key space covered by the current replica and all
-		// subsumed replicas.
-		for i := range srKeySpans {
-			if srKeySpans[i].Key.Compare(totalKeySpans[i].Key) < 0 {
-				totalKeySpans[i].Key = srKeySpans[i].Key
-			}
-			if srKeySpans[i].EndKey.Compare(totalKeySpans[i].EndKey) > 0 {
-				totalKeySpans[i].EndKey = srKeySpans[i].EndKey
-			}
-		}
 	}
 
-	// We might have to create SSTs for the range local keys, lock table keys,
-	// and user keys depending on if the subsumed replicas are not fully
-	// contained by the replica in our snapshot. The following is an example to
-	// this case happening.
-	//
-	// a       b       c       d
-	// |---1---|-------2-------|  S1
-	// |---1-------------------|  S2
-	// |---1-----------|---3---|  S3
-	//
-	// Since the merge is the first operation to happen, a follower could be down
-	// before it completes. It is reasonable for a snapshot for r1 from S3 to
-	// subsume both r1 and r2 in S1.
-	for i := range keySpans {
-		// The snapshot must never subsume a replica that extends the range of the
-		// replica to the left. This is because splits and merges (the only
-		// operation that change the key bounds) always leave the start key intact.
-		// Extending to the left implies that either we merged "to the left" (we
-		// don't), or that we're applying a snapshot for another range (we don't do
-		// that either). Something is severely wrong for this to happen.
-		if totalKeySpans[i].Key.Compare(keySpans[i].Key) < 0 {
-			log.Dev.Fatalf(ctx, "subsuming replica to our left; key span: %v; total key span %v",
-				keySpans[i], totalKeySpans[i])
-		}
+	return nil
+}
 
-		// In the comments below, s1, ..., sn are the end keys for the subsumed
-		// replicas (for the current keySpan).
-		// Note that if there aren't any subsumed replicas (the common case), the
-		// next comparison is always zero and this loop is a no-op.
-
-		if totalKeySpans[i].EndKey.Compare(keySpans[i].EndKey) <= 0 {
-			// The subsumed replicas are fully contained in the snapshot:
-			//
-			// [a---s1---...---sn)
-			// [a---------------------b)
-			//
-			// We don't need to clear additional keyspace here, since clearing `[a,b)`
-			// will also clear all subsumed replicas.
-			continue
-		}
+// clearResidualDataOnNarrowSnapshot clears the overlapping replicas' data not
+// covered by the snapshot. Specifically, the data between the snapshot's end
+// key and that of the keyspace covered by s.origDesc + s.subsume.descs, if the
+// former is lower (i.e. the snapshot is "narrower"). Note that the start keys
+// of the snapshot and s.origDesc[^1] match, so the residual data may only exist
+// between the end keys. clearResidualDataOnNarrowSnapshot is a no-op if there
+// is no narrowing business to speak of.
+//
+// Visually, the picture looks as follows:
+//
+// The simplest case is when the snapshot isn't subsuming any replicas.
+// original descriptor: [a-----------------------------c)
+// snapshot descriptor: [a---------------------b)
+// cleared: [b, c)
+//
+// In the more general case[^2]:
+//
+// store descriptors:   [a----------------s1---...---sn)
+// snapshot descriptor: [a---------------------b)
+// cleared: [b, sn)
+//
+// Practically speaking, the simple case above corresponds to a replica learning
+// about a split through a snapshot. The more general case corresponds to a
+// replica learning about a series of merges and at least one split through the
+// snapshot.
+//
+// [1] Assuming s.origDesc is initialized. If it isn't, and we're applying an
+// initial snapshot, the snapshot may still be narrower than the end key of the
+// right-most subsumed replica.
+//
+// [2] In the diagram , S1...Sn correspond to subsumed replicas with end keys
+// s1...sn respectively. These are all replicas on the store that overlap with
+// the snapshot descriptor, covering the range [a,b), and the right-most
+// descriptor is that of replica Sn.
+func (s *snapWriteBuilder) clearResidualDataOnNarrowSnapshot(ctx context.Context) error {
+	if !s.origDesc.IsInitialized() && len(s.subsume) == 0 {
+		// Early return in the case where we're ingesting an initial snapshot and
+		// there are no subsumed replicas that the snapshot could be making
+		// narrower.
+		return nil
+	}
 
-		// The subsumed replicas extend past the snapshot:
-		//
-		// [a----------------s1---...---sn)
-		// [a---------------------b)
-		//
-		// We need to additionally clear [b,sn).
+	rightMostDesc := s.origDesc
+	if len(s.subsume) != 0 {
+		// NB: s.subsume are non-overlapping and sorted by start key. Pick the last
+		// one to determine whether the snapshot is narrowing the keyspace or not.
+		rightMostDesc = s.subsume[len(s.subsume)-1].desc
+	}
 
+	if rightMostDesc.EndKey.Compare(s.desc.EndKey) <= 0 {
+		return nil // we aren't narrowing anything; no-op
+	}
+
+	// TODO(sep-raft-log): read from the state machine engine here.
+	reader := storage.Reader(s.todoEng)
+	for _, span := range rditer.Select(0, rditer.SelectOpts{
+		Ranged: rditer.SelectRangedOptions{RSpan: roachpb.RSpan{
+			Key: s.desc.EndKey, EndKey: rightMostDesc.EndKey,
+		},
+			SystemKeys: true,
+			LockTable:  true,
+			UserKeys:   true,
+		},
+	}) {
 		if err := s.writeSST(ctx, func(ctx context.Context, w storage.Writer) error {
 			return storage.ClearRangeWithHeuristic(
-				ctx, reader, w,
-				keySpans[i].EndKey, totalKeySpans[i].EndKey,
-				kvstorage.ClearRangeThresholdPointKeys,
+				ctx, reader, w, span.Key, span.EndKey, kvstorage.ClearRangeThresholdPointKeys,
 			)
 		}); err != nil {
 			return err
 		}
-		s.cleared = append(s.cleared,
-			roachpb.Span{Key: keySpans[i].EndKey, EndKey: totalKeySpans[i].EndKey})
+		s.cleared = append(s.cleared, span)
 	}
+
 	return nil
 }

pkg/kv/kvserver/snapshot_apply_prepare_test.go

@@ -92,6 +92,7 @@ func TestPrepareSnapApply(t *testing.T) {
 		truncState: kvserverpb.RaftTruncatedState{Index: 100, Term: 20},
 		hardState:  raftpb.HardState{Term: 20, Commit: 100},
 		desc:       desc(id.RangeID, "a", "k"),
+		origDesc:   desc(id.RangeID, "a", "k"),
 		subsume: []destroyReplicaInfo{
 			{id: roachpb.FullReplicaID{RangeID: descA.RangeID, ReplicaID: replicaID}, desc: descA},
 			{id: roachpb.FullReplicaID{RangeID: descB.RangeID, ReplicaID: replicaID}, desc: descB},