diff --git a/pkg/spanconfig/spanconfigstore/BUILD.bazel b/pkg/spanconfig/spanconfigstore/BUILD.bazel index a17a28eade92..6fcc3a29e773 100644 --- a/pkg/spanconfig/spanconfigstore/BUILD.bazel +++ b/pkg/spanconfig/spanconfigstore/BUILD.bazel @@ -14,6 +14,7 @@ go_library( "//pkg/settings", "//pkg/spanconfig", "//pkg/util/interval", + "//pkg/util/iterutil", "//pkg/util/log", "//pkg/util/syncutil", "@com_github_cockroachdb_errors//:errors", diff --git a/pkg/spanconfig/spanconfigstore/store.go b/pkg/spanconfig/spanconfigstore/store.go index c58bd6b282fc..19dde8e9fe42 100644 --- a/pkg/spanconfig/spanconfigstore/store.go +++ b/pkg/spanconfig/spanconfigstore/store.go @@ -19,6 +19,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/settings" "github.com/cockroachdb/cockroach/pkg/spanconfig" "github.com/cockroachdb/cockroach/pkg/util/interval" + "github.com/cockroachdb/cockroach/pkg/util/iterutil" "github.com/cockroachdb/cockroach/pkg/util/log" "github.com/cockroachdb/cockroach/pkg/util/syncutil" "github.com/cockroachdb/errors" @@ -71,7 +72,7 @@ func (s *Store) NeedsSplit(ctx context.Context, start, end roachpb.RKey) bool { } // ComputeSplitKey is part of the spanconfig.StoreReader interface. -func (s *Store) ComputeSplitKey(ctx context.Context, start, end roachpb.RKey) roachpb.RKey { +func (s *Store) ComputeSplitKey(_ context.Context, start, end roachpb.RKey) roachpb.RKey { sp := roachpb.Span{Key: start.AsRawKey(), EndKey: end.AsRawKey()} // We don't want to split within the system config span while we're still @@ -141,6 +142,42 @@ func (s *Store) Apply( return deleted, added } +// Copy returns a copy of the Store. +func (s *Store) Copy(ctx context.Context) *Store { + clone := New(s.fallback) + _ = s.ForEachOverlapping(ctx, keys.EverythingSpan, func(entry roachpb.SpanConfigEntry) error { + clone.Apply(ctx, false /* dryrun */, spanconfig.Update{ + Span: entry.Span, + Config: entry.Config, + }) + return nil + }) + return clone +} + +// ForEachOverlapping iterates through the set of entries that overlap with the +// given span, in sorted order. It does not return an error if the callback +// doesn't. +func (s *Store) ForEachOverlapping( + _ context.Context, sp roachpb.Span, f func(roachpb.SpanConfigEntry) error, +) error { + s.mu.RLock() + defer s.mu.RUnlock() + + // Iterate over all overlapping ranges and invoke the callback with the + // corresponding span config entries. + for _, overlapping := range s.mu.tree.Get(sp.AsRange()) { + entry := overlapping.(*storeEntry).SpanConfigEntry + if err := f(entry); err != nil { + if iterutil.Done(err) { + err = nil + } + return err + } + } + return nil +} + func (s *Store) applyInternal( dryrun bool, updates ...spanconfig.Update, ) (deleted []roachpb.Span, added []roachpb.SpanConfigEntry, err error) { diff --git a/pkg/spanconfig/spanconfigstore/store_test.go b/pkg/spanconfig/spanconfigstore/store_test.go index 4454c7cc6c91..fca07d3e134e 100644 --- a/pkg/spanconfig/spanconfigstore/store_test.go +++ b/pkg/spanconfig/spanconfigstore/store_test.go @@ -28,23 +28,6 @@ import ( "github.com/stretchr/testify/require" ) -// TestingGetAllOverlapping is a testing only helper to retrieve the set of -// overlapping entries in sorted order. -func (s *Store) TestingGetAllOverlapping( - _ context.Context, sp roachpb.Span, -) []roachpb.SpanConfigEntry { - s.mu.RLock() - defer s.mu.RUnlock() - - // Iterate over all overlapping ranges and return corresponding span config - // entries. - var res []roachpb.SpanConfigEntry - for _, overlapping := range s.mu.tree.Get(sp.AsRange()) { - res = append(res, overlapping.(*storeEntry).SpanConfigEntry) - } - return res -} - // TestingApplyInternal exports an internal method for testing purposes. func (s *Store) TestingApplyInternal( _ context.Context, dryrun bool, updates ...spanconfig.Update, @@ -139,11 +122,14 @@ func TestDataDriven(t *testing.T) { case "overlapping": d.ScanArgs(t, "span", &spanStr) span := spanconfigtestutils.ParseSpan(t, spanStr) - entries := store.TestingGetAllOverlapping(ctx, span) + var results []string - for _, entry := range entries { - results = append(results, spanconfigtestutils.PrintSpanConfigEntry(entry)) - } + _ = store.ForEachOverlapping(ctx, span, + func(entry roachpb.SpanConfigEntry) error { + results = append(results, spanconfigtestutils.PrintSpanConfigEntry(entry)) + return nil + }, + ) return strings.Join(results, "\n") default: @@ -252,53 +238,121 @@ func TestRandomized(t *testing.T) { } } - overlappingConfigs := store.TestingGetAllOverlapping(ctx, testSpan) if !expFound { - require.Len(t, overlappingConfigs, 0) + _ = store.ForEachOverlapping(ctx, testSpan, + func(entry roachpb.SpanConfigEntry) error { + t.Fatalf("found unexpected entry: %s", + spanconfigtestutils.PrintSpanConfigEntry(entry)) + return nil + }, + ) } else { - // Check to see that the set of overlapping span configs is exactly what - // we expect. - require.Len(t, overlappingConfigs, 1) - gotSpan, gotConfig := overlappingConfigs[0].Span, overlappingConfigs[0].Config - - require.Truef(t, gotSpan.Contains(testSpan), - "improper result: expected retrieved span (%s) to contain test span (%s)", - spanconfigtestutils.PrintSpan(gotSpan), spanconfigtestutils.PrintSpan(testSpan)) - - require.Truef(t, expConfig.Equal(gotConfig), - "mismatched configs: expected %s, got %s", - spanconfigtestutils.PrintSpanConfig(expConfig), spanconfigtestutils.PrintSpanConfig(gotConfig)) - - // Ensure that the config accessed through the StoreReader interface is - // the same as above. - storeReaderConfig, err := store.GetSpanConfigForKey(ctx, roachpb.RKey(testSpan.Key)) - require.NoError(t, err) - require.True(t, gotConfig.Equal(storeReaderConfig)) + found := false + _ = store.ForEachOverlapping(ctx, testSpan, + func(entry roachpb.SpanConfigEntry) error { + if found { + t.Fatalf("expected single overlapping entry, found second: %s", + spanconfigtestutils.PrintSpanConfigEntry(entry)) + } + found = true + + // Check that the entry is exactly what we'd expect. + gotSpan, gotConfig := entry.Span, entry.Config + require.Truef(t, gotSpan.Contains(testSpan), + "improper result: expected retrieved span (%s) to contain test span (%s)", + spanconfigtestutils.PrintSpan(gotSpan), spanconfigtestutils.PrintSpan(testSpan)) + + require.Truef(t, expConfig.Equal(gotConfig), + "mismatched configs: expected %s, got %s", + spanconfigtestutils.PrintSpanConfig(expConfig), spanconfigtestutils.PrintSpanConfig(gotConfig)) + + // Ensure that the config accessed through the StoreReader interface is + // the same as above. + storeReaderConfig, err := store.GetSpanConfigForKey(ctx, roachpb.RKey(testSpan.Key)) + require.NoError(t, err) + require.True(t, gotConfig.Equal(storeReaderConfig)) + + return nil + }, + ) } - var last roachpb.SpanConfigEntry everythingSpan := spanconfigtestutils.ParseSpan(t, fmt.Sprintf("[%s,%s)", string(alphabet[0]), string(alphabet[len(alphabet)-1]))) - for i, cur := range store.TestingGetAllOverlapping(ctx, everythingSpan) { - if i == 0 { - last = cur - continue - } - // All spans are expected to be valid. - require.True(t, cur.Span.Valid(), - "expected to only find valid spans, found %s", - spanconfigtestutils.PrintSpan(cur.Span), - ) + var last roachpb.SpanConfigEntry + _ = store.ForEachOverlapping(ctx, everythingSpan, + func(cur roachpb.SpanConfigEntry) error { + // All spans are expected to be valid. + require.True(t, cur.Span.Valid(), + "expected to only find valid spans, found %s", + spanconfigtestutils.PrintSpan(cur.Span), + ) + + if last.Empty() { + last = cur + return nil + } - // Span configs are returned in strictly sorted order. - require.True(t, last.Span.Key.Compare(cur.Span.Key) < 0, - "expected to find spans in strictly sorted order, found %s then %s", - spanconfigtestutils.PrintSpan(last.Span), spanconfigtestutils.PrintSpan(cur.Span)) + // Span configs are returned in strictly sorted order. + require.True(t, last.Span.Key.Compare(cur.Span.Key) < 0, + "expected to find spans in strictly sorted order, found %s then %s", + spanconfigtestutils.PrintSpan(last.Span), spanconfigtestutils.PrintSpan(cur.Span)) + + // Span configs must also be non-overlapping. + require.Falsef(t, last.Span.Overlaps(cur.Span), + "expected non-overlapping spans, found %s and %s", + spanconfigtestutils.PrintSpan(last.Span), spanconfigtestutils.PrintSpan(cur.Span)) + + return nil + }, + ) +} + +// TestStoreClone verifies that a cloned store contains the same contents as the +// original. +func TestStoreClone(t *testing.T) { + defer leaktest.AfterTest(t)() + + ctx := context.Background() + + everything := spanconfigtestutils.ParseSpan(t, "[a, z)") + updates := []spanconfig.Update{ + { + Span: spanconfigtestutils.ParseSpan(t, "[a, b)"), + Config: spanconfigtestutils.ParseConfig(t, "A"), + }, + { + Span: spanconfigtestutils.ParseSpan(t, "[c, d)"), + Config: spanconfigtestutils.ParseConfig(t, "C"), + }, + { + Span: spanconfigtestutils.ParseSpan(t, "[e, f)"), + Config: spanconfigtestutils.ParseConfig(t, "E"), + }, + } - // Span configs must also be non-overlapping. - require.Falsef(t, last.Span.Overlaps(cur.Span), - "expected non-overlapping spans, found %s and %s", - spanconfigtestutils.PrintSpan(last.Span), spanconfigtestutils.PrintSpan(cur.Span)) + original := New(roachpb.TestingDefaultSpanConfig()) + original.Apply(ctx, false, updates...) + clone := original.Copy(ctx) + + var originalEntries, clonedEntries []roachpb.SpanConfigEntry + _ = original.ForEachOverlapping(ctx, everything, + func(entry roachpb.SpanConfigEntry) error { + originalEntries = append(originalEntries, entry) + return nil + }, + ) + + _ = clone.ForEachOverlapping(ctx, everything, + func(entry roachpb.SpanConfigEntry) error { + clonedEntries = append(clonedEntries, entry) + return nil + }, + ) + + require.Equal(t, len(originalEntries), len(clonedEntries)) + for i := 0; i < len(originalEntries); i++ { + require.True(t, originalEntries[i].Equal(clonedEntries[i])) } }