asim: add set span config

Previously, a span config could only be set for a range with and an ID
in the allocation simulator. This commit adds the `SetSpanConfig`
command which allows setting a span config for an arbitrary key span.
The key span is split accordingly to apply the span config.

Part of: cockroachdb#90141

Release note: None

pkg/kv/kvserver/asim/queue/replicate_queue_test.go

@@ -84,7 +84,7 @@ func TestReplicateQueue(t *testing.T) {
 		s := state.NewStateWithReplCounts(replicaCounts, 2 /* replsPerRange */, 1000 /* keyspace */, testSettings)
 		spanConfig := roachpb.SpanConfig{NumVoters: replicationFactor, NumReplicas: replicationFactor}
 		for _, r := range s.Ranges() {
-			s.SetSpanConfig(r.RangeID(), spanConfig)
+			s.SetSpanConfigForRange(r.RangeID(), spanConfig)
 		}
 		return s
 	}

pkg/kv/kvserver/asim/queue/split_queue_test.go

@@ -48,7 +48,7 @@ func TestSplitQueue(t *testing.T) {
 			NumReplicas: replicationFactor,
 		}
 		for _, r := range s.Ranges() {
-			s.SetSpanConfig(r.RangeID(), spanConfig)
+			s.SetSpanConfigForRange(r.RangeID(), spanConfig)
 		}
 		s.TransferLease(state.RangeID(2 /* The interesting range */), leaseholder)
 		return s

pkg/kv/kvserver/asim/state/config_loader.go

@@ -322,7 +322,7 @@ func LoadRangeInfo(s State, rangeInfos ...RangeInfo) {
 			))
 		}
 
-		if !s.SetSpanConfig(rng.RangeID(), *r.Config) {
+		if !s.SetSpanConfigForRange(rng.RangeID(), *r.Config) {
 			panic(fmt.Sprintf(
 				"Unable to load config: cannot set span config for range %s",
 				rng,

pkg/kv/kvserver/asim/state/impl.go

@@ -544,15 +544,91 @@ func (s *state) removeReplica(rangeID RangeID, storeID StoreID) bool {
 	return true
 }
 
-// SetSpanConfig set the span config for the Range with ID RangeID.
-func (s *state) SetSpanConfig(rangeID RangeID, spanConfig roachpb.SpanConfig) bool {
+// SetSpanConfigForRange set the span config for the Range with ID RangeID.
+func (s *state) SetSpanConfigForRange(rangeID RangeID, spanConfig roachpb.SpanConfig) bool {
 	if rng, ok := s.ranges.rangeMap[rangeID]; ok {
 		rng.config = spanConfig
 		return true
 	}
 	return false
 }
 
+// SetSpanConfig sets the span config for all ranges represented by the span,
+// splitting if necessary.
+func (s *state) SetSpanConfig(span roachpb.Span, config roachpb.SpanConfig) {
+	startKey := ToKey(span.Key)
+	endKey := ToKey(span.EndKey)
+
+	// Decide whether we need to split due to the config intersecting an existing
+	// range boundary. Split if necessary. Then apply the span config to all the
+	// ranges contained within the span. e.g.
+	//   ranges r1: [a, c) r2: [c, z)
+	//   span: [b, f)
+	// resulting ranges:
+	//   [a, b)         - keeps old span config from [a,c)
+	//   [b, c) [c, f)  - gets the new span config passed in
+	//   [f, z)         - keeps old span config from [c,z)
+
+	splitsRequired := []Key{}
+	s.ranges.rangeTree.DescendLessOrEqual(&rng{startKey: startKey}, func(i btree.Item) bool {
+		cur, _ := i.(*rng)
+		rStart := cur.startKey
+		// There are two cases we handle:
+		// (1) rStart == startKey: We don't need to split.
+		// (2) rStart < startKey:  We need to split into lhs [rStart, startKey) and
+		//     rhs [startKey, ...). Where the lhs does not have the span config
+		//     applied and the rhs does.
+		if rStart < startKey {
+			splitsRequired = append(splitsRequired, startKey)
+		}
+		return false
+	})
+
+	s.ranges.rangeTree.DescendLessOrEqual(&rng{startKey: endKey}, func(i btree.Item) bool {
+		cur, _ := i.(*rng)
+		rEnd := cur.endKey
+		rStart := cur.startKey
+		if rStart == endKey {
+			return false
+		}
+		// There are two cases we handle:
+		// (1) rEnd == endKey: We don't need to split.
+		// (2) rEnd >  endKey: We need to split into lhs [..., endKey) and rhs
+		//     [endKey, rEnd). Where the lhs has the span config applied and the rhs
+		//     does not.
+		// Split required if its the last range we will hit.
+		if rEnd > endKey {
+			splitsRequired = append(splitsRequired, endKey)
+		}
+		return false
+	})
+
+	for _, splitKey := range splitsRequired {
+		// We panic here as we don't have any way to roll back the split if one
+		// succeeds and another doesn't
+		if _, _, ok := s.SplitRange(splitKey); !ok {
+			panic(fmt.Sprintf(
+				"programming error: unable to split range (key=%d) for set span "+
+					"config=%s, state=%s", splitKey, config.String(), s))
+		}
+	}
+
+	// Apply the span config to all the ranges affected.
+	s.ranges.rangeTree.AscendGreaterOrEqual(&rng{startKey: startKey}, func(i btree.Item) bool {
+		cur, _ := i.(*rng)
+		if cur.startKey == endKey {
+			return false
+		}
+		if cur.startKey > endKey {
+			panic("programming error: unexpected range found with start key > end key")
+		}
+		if !s.SetSpanConfigForRange(cur.rangeID, config) {
+			panic("programming error: unable to set span config for range")
+		}
+		return true
+	})
+}
+
 // SplitRange splits the Range which contains Key in [StartKey, EndKey).
 // The Range is partitioned into [StartKey, Key), [Key, EndKey) and
 // returned. The right hand side of this split, is the new Range. If any

pkg/kv/kvserver/asim/state/state.go

@@ -115,8 +115,11 @@ type State interface {
 	// RangeSpan returns the [StartKey, EndKey) for the range with ID RangeID
 	// if it exists, otherwise it returns false.
 	RangeSpan(RangeID) (Key, Key, bool)
-	// SetSpanConfig set the span config for the Range with ID RangeID.
-	SetSpanConfig(RangeID, roachpb.SpanConfig) bool
+	// SetSpanConfigForRange set the span config for the Range with ID RangeID.
+	SetSpanConfigForRange(RangeID, roachpb.SpanConfig) bool
+	// SetSpanConfig sets the span config for all ranges represented by the span,
+	// splitting if necessary.
+	SetSpanConfig(roachpb.Span, roachpb.SpanConfig)
 	// ValidTransfer returns whether transferring the lease for the Range with ID
 	// RangeID, to the Store with ID StoreID is valid.
 	ValidTransfer(RangeID, StoreID) bool

pkg/kv/kvserver/asim/state/state_test.go

@@ -452,6 +452,143 @@ func TestSplitRangeDeterministic(t *testing.T) {
 	}
 }
 
+func TestSetSpanConfig(t *testing.T) {
+	settings := config.DefaultSimulationSettings()
+
+	setupState := func() State {
+		s := newState(settings)
+		node := s.AddNode()
+		_, ok := s.AddStore(node.NodeID())
+		require.True(t, ok)
+
+		// Setup with the following ranges:
+		// [1,50) [50, 100) [100, 1000)
+		_, _, ok = s.SplitRange(1)
+		require.True(t, ok)
+		_, _, ok = s.SplitRange(50)
+		require.True(t, ok)
+		_, _, ok = s.SplitRange(100)
+		require.True(t, ok)
+		return s
+	}
+
+	keySet := func(keys ...Key) map[Key]struct{} {
+		ret := make(map[Key]struct{}, len(keys))
+		for _, key := range keys {
+			ret[key] = struct{}{}
+		}
+		return ret
+	}
+
+	testCases := []struct {
+		desc                        string
+		end, start                  Key
+		expectedAppliedStartKeys    map[Key]struct{}
+		expectedNonAppliedStartKeys map[Key]struct{}
+	}{
+		{
+			// Matching start and end key over a single range.
+			desc:                        "[1,50)",
+			start:                       1,
+			end:                         50,
+			expectedAppliedStartKeys:    keySet(1),
+			expectedNonAppliedStartKeys: keySet(50, 100),
+		},
+		{
+			// Matching start and end key over multiple ranges.
+			desc:                        "[1,100)",
+			start:                       1,
+			end:                         100,
+			expectedAppliedStartKeys:    keySet(1, 50),
+			expectedNonAppliedStartKeys: keySet(100),
+		},
+		{
+			// Matching start key, overlapping end key over single range.
+			desc:                        "[1,40)",
+			start:                       1,
+			end:                         40,
+			expectedAppliedStartKeys:    keySet(1),
+			expectedNonAppliedStartKeys: keySet(40, 50, 100),
+		},
+		{
+			// Matching start key, overlapping end key over multiple ranges.
+			desc:                        "[1,70)",
+			start:                       1,
+			end:                         70,
+			expectedAppliedStartKeys:    keySet(1, 50),
+			expectedNonAppliedStartKeys: keySet(70, 100),
+		},
+		{
+			// Overlapping start key, matching end key over single range.
+			desc:                        "[20,50)",
+			start:                       20,
+			end:                         50,
+			expectedAppliedStartKeys:    keySet(20),
+			expectedNonAppliedStartKeys: keySet(1, 50, 100),
+		},
+		{
+			// Overlapping start key, matching end key over multiple ranges.
+			desc:                        "[20,100)",
+			start:                       20,
+			end:                         100,
+			expectedAppliedStartKeys:    keySet(20, 50),
+			expectedNonAppliedStartKeys: keySet(1, 100),
+		},
+		{
+			// Overlapping start and end key over single range.
+			desc:                        "[20,40)",
+			start:                       20,
+			end:                         40,
+			expectedAppliedStartKeys:    keySet(20),
+			expectedNonAppliedStartKeys: keySet(1, 40, 50, 100),
+		},
+		{
+			// Overlapping start and end key over multiple ranges.
+			desc:                        "[20,70)",
+			start:                       20,
+			end:                         70,
+			expectedAppliedStartKeys:    keySet(20, 50),
+			expectedNonAppliedStartKeys: keySet(1, 70, 100),
+		},
+	}
+
+	const sentinel = int64(-1)
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			s := setupState()
+			config := roachpb.SpanConfig{
+				RangeMinBytes: sentinel,
+			}
+			span := roachpb.Span{
+				Key:    tc.start.ToRKey().AsRawKey(),
+				EndKey: tc.end.ToRKey().AsRawKey(),
+			}
+			s.SetSpanConfig(span, config)
+			for _, rng := range s.Ranges() {
+				start, _, ok := s.RangeSpan(rng.RangeID())
+				require.True(t, ok)
+				config := rng.SpanConfig()
+
+				// We ignore the range starting with the min key as we have pre-split
+				// the keyspace to start at key=1.
+				if start == MinKey {
+					continue
+				}
+
+				if _, ok := tc.expectedAppliedStartKeys[start]; ok {
+					require.Equal(t, sentinel, config.RangeMinBytes,
+						"sentinel not set, when should be for start key %d", start)
+				} else if _, ok := tc.expectedNonAppliedStartKeys[start]; ok {
+					require.NotEqual(t, sentinel, config.RangeMinBytes,
+						"sentinel set, when it should not be for start key %d", start)
+				} else {
+					t.Fatalf("Start key not found in either expected apply keys or unapplied keys %s", rng)
+				}
+			}
+		})
+	}
+}
+
 func TestSetNodeLiveness(t *testing.T) {
 	t.Run("liveness func", func(t *testing.T) {
 		s := LoadClusterInfo(