diff --git a/pkg/kv/kvserver/asim/queue/replicate_queue_test.go b/pkg/kv/kvserver/asim/queue/replicate_queue_test.go index 1af4f403f8d6..cb59d061f5aa 100644 --- a/pkg/kv/kvserver/asim/queue/replicate_queue_test.go +++ b/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 } diff --git a/pkg/kv/kvserver/asim/queue/split_queue_test.go b/pkg/kv/kvserver/asim/queue/split_queue_test.go index 58116410ffc0..7c2a6ac7a816 100644 --- a/pkg/kv/kvserver/asim/queue/split_queue_test.go +++ b/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 diff --git a/pkg/kv/kvserver/asim/state/config_loader.go b/pkg/kv/kvserver/asim/state/config_loader.go index 65a34a1b64cf..f82744a60a3f 100644 --- a/pkg/kv/kvserver/asim/state/config_loader.go +++ b/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, diff --git a/pkg/kv/kvserver/asim/state/impl.go b/pkg/kv/kvserver/asim/state/impl.go index 4d71c6ebdb52..622a2ee85c16 100644 --- a/pkg/kv/kvserver/asim/state/impl.go +++ b/pkg/kv/kvserver/asim/state/impl.go @@ -544,8 +544,8 @@ 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 @@ -553,6 +553,82 @@ func (s *state) SetSpanConfig(rangeID RangeID, spanConfig roachpb.SpanConfig) bo 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 diff --git a/pkg/kv/kvserver/asim/state/state.go b/pkg/kv/kvserver/asim/state/state.go index 202f52f86c1e..45333eb4bdb8 100644 --- a/pkg/kv/kvserver/asim/state/state.go +++ b/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 diff --git a/pkg/kv/kvserver/asim/state/state_test.go b/pkg/kv/kvserver/asim/state/state_test.go index 2d7f13f58bb5..2b0f122db5b7 100644 --- a/pkg/kv/kvserver/asim/state/state_test.go +++ b/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(