Merge cockroachdb#73402

73402: spanconfig/store: implement cloning and an overlapping iterator r=irfansharif a=irfansharif We'll use both in spanconfig.Reconciler (cockroachdb#71994). Release note: None Co-authored-by: irfan sharif <[email protected]>
irfansharif · Dec 3, 2021 · eea11ee · eea11ee
2 parents d6e5378 + 18f2269
commit eea11ee
Show file tree

Hide file tree

Showing 3 changed files with 153 additions and 61 deletions.
diff --git 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
@@ -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
@@ -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]))
 	}
 }