From 3b0df858ecbf19ff1f579f66beedfd028bbad015 Mon Sep 17 00:00:00 2001 From: irfan sharif Date: Wed, 27 Apr 2022 01:31:30 -0400 Subject: [PATCH] spanconfig/store: intern span configs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Avoid the expensive proto.Equal() when computing split keys; - Reduce memory overhead of the data structure $ dev bench pkg/spanconfig/spanconfigstore -f=BenchmarkStoreComputeSplitKey -v BenchmarkStoreComputeSplitKey BenchmarkStoreComputeSplitKey/num-entries=10000 BenchmarkStoreComputeSplitKey/num-entries=10000-10 90323 ns/op BenchmarkStoreComputeSplitKey/num-entries=100000 BenchmarkStoreComputeSplitKey/num-entries=100000-10 915936 ns/op BenchmarkStoreComputeSplitKey/num-entries=1000000 BenchmarkStoreComputeSplitKey/num-entries=1000000-10 9575781 ns/op $ benchstat old.txt new.txt # from previous commit name old time/op new time/op delta StoreComputeSplitKey/num-entries=10000-10 431µs ± 0% 90µs ± 0% ~ (p=1.000 n=1+1) StoreComputeSplitKey/num-entries=100000-10 4.31ms ± 0% 0.92ms ± 0% ~ (p=1.000 n=1+1) StoreComputeSplitKey/num-entries=1000000-10 43.8ms ± 0% 9.6ms ± 0% ~ (p=1.000 n=1+1) Release note: None --- pkg/spanconfig/spanconfigstore/BUILD.bazel | 2 + pkg/spanconfig/spanconfigstore/interner.go | 88 +++++++++++++ pkg/spanconfig/spanconfigstore/span_store.go | 124 ++++++++++++------ .../spanconfigstore/span_store_test.go | 72 +++++----- pkg/spanconfig/spanconfigstore/store.go | 27 ++-- pkg/spanconfig/spanconfigstore/store_test.go | 24 ++-- pkg/spanconfig/testing_knobs.go | 4 + 7 files changed, 241 insertions(+), 100 deletions(-) create mode 100644 pkg/spanconfig/spanconfigstore/interner.go diff --git a/pkg/spanconfig/spanconfigstore/BUILD.bazel b/pkg/spanconfig/spanconfigstore/BUILD.bazel index 6cefe05d4a46..3801bdcec577 100644 --- a/pkg/spanconfig/spanconfigstore/BUILD.bazel +++ b/pkg/spanconfig/spanconfigstore/BUILD.bazel @@ -4,6 +4,7 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test") go_library( name = "spanconfigstore", srcs = [ + "interner.go", "span_store.go", "store.go", "system_store.go", @@ -19,6 +20,7 @@ go_library( "//pkg/spanconfig", "//pkg/util/iterutil", "//pkg/util/log", + "//pkg/util/protoutil", "//pkg/util/syncutil", "//pkg/util/timeutil", # keep "@com_github_cockroachdb_errors//:errors", diff --git a/pkg/spanconfig/spanconfigstore/interner.go b/pkg/spanconfig/spanconfigstore/interner.go new file mode 100644 index 000000000000..55a1d5646cba --- /dev/null +++ b/pkg/spanconfig/spanconfigstore/interner.go @@ -0,0 +1,88 @@ +// Copyright 2022 The Cockroach Authors. +// +// Use of this software is governed by the Business Source License +// included in the file licenses/BSL.txt. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0, included in the file +// licenses/APL.txt. + +package spanconfigstore + +import ( + "context" + + "github.com/cockroachdb/cockroach/pkg/roachpb" + "github.com/cockroachdb/cockroach/pkg/util/log" + "github.com/cockroachdb/cockroach/pkg/util/protoutil" +) + +type internerID uint64 + +// interner interns span config protos. It's a ref-counted data structure that +// maintains a single copy of every unique config that that's been added to +// it. Configs that are maintained are identified/retrivable using an +// internerID. Configs can also be removed, and if there are no more references +// to it, we no longer hold onto it in memory. +type interner struct { + internIDAlloc internerID + + configToID map[string]internerID + idToConfig map[internerID]roachpb.SpanConfig + idToRefCount map[internerID]uint64 +} + +func newInterner() *interner { + return &interner{ + configToID: make(map[string]internerID), + idToConfig: make(map[internerID]roachpb.SpanConfig), + idToRefCount: make(map[internerID]uint64), + } +} + +func (i *interner) add(ctx context.Context, conf roachpb.SpanConfig) internerID { + marshalled, err := protoutil.Marshal(&conf) + if err != nil { + log.Fatalf(ctx, "%v", err) + } + + if id, found := i.configToID[string(marshalled)]; found { + i.idToRefCount[id]++ + return id + } + + i.internIDAlloc++ + + id := i.internIDAlloc + i.configToID[string(marshalled)] = id + i.idToConfig[i.internIDAlloc] = conf + i.idToRefCount[id] = 1 + return id +} + +func (i *interner) get(id internerID) (roachpb.SpanConfig, bool) { + conf, found := i.idToConfig[id] + return conf, found +} + +func (i *interner) remove(ctx context.Context, id internerID) { + conf, found := i.idToConfig[id] + if !found { + return // nothing to do + } + + i.idToRefCount[id]-- + if i.idToRefCount[id] != 0 { + return // nothing to do + } + + marshalled, err := protoutil.Marshal(&conf) + if err != nil { + log.Infof(ctx, "%v", err) + } + + delete(i.idToConfig, id) + delete(i.idToRefCount, id) + delete(i.configToID, string(marshalled)) +} diff --git a/pkg/spanconfig/spanconfigstore/span_store.go b/pkg/spanconfig/spanconfigstore/span_store.go index 9512910bd806..b6767f03791c 100644 --- a/pkg/spanconfig/spanconfigstore/span_store.go +++ b/pkg/spanconfig/spanconfigstore/span_store.go @@ -49,8 +49,9 @@ var hostCoalesceAdjacentSetting = settings.RegisterBoolSetting( // interval tree to store non-overlapping span configurations. It isn't safe for // concurrent use. type spanConfigStore struct { - btree *btree - idAlloc uint64 + btree *btree + treeIDAlloc uint64 + interner *interner settings *cluster.Settings knobs *spanconfig.TestingKnobs @@ -67,6 +68,7 @@ func newSpanConfigStore( settings: settings, knobs: knobs, btree: &btree{}, + interner: newInterner(), } return s } @@ -74,12 +76,12 @@ func newSpanConfigStore( // copy returns a copy of the spanConfigStore. func (s *spanConfigStore) copy(ctx context.Context) *spanConfigStore { clone := newSpanConfigStore(s.settings, s.knobs) - _ = s.forEachOverlapping(keys.EverythingSpan, func(entry spanConfigEntry) error { - record, err := spanconfig.MakeRecord(spanconfig.MakeTargetFromSpan(entry.span), entry.config) + _ = s.forEachOverlapping(keys.EverythingSpan, func(sp roachpb.Span, conf roachpb.SpanConfig) error { + record, err := spanconfig.MakeRecord(spanconfig.MakeTargetFromSpan(sp), conf) if err != nil { log.Fatalf(ctx, "%v", err) } - _, _, err = clone.apply(false /* dryrun */, spanconfig.Update(record)) + _, _, err = clone.apply(ctx, false /* dryrun */, spanconfig.Update(record)) if err != nil { log.Fatalf(ctx, "%v", err) } @@ -91,12 +93,15 @@ func (s *spanConfigStore) copy(ctx context.Context) *spanConfigStore { // 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 *spanConfigStore) forEachOverlapping(sp roachpb.Span, f func(spanConfigEntry) error) error { +func (s *spanConfigStore) forEachOverlapping( + sp roachpb.Span, f func(roachpb.Span, roachpb.SpanConfig) error, +) error { // Iterate over all overlapping ranges and invoke the callback with the // corresponding span config entries. iter, query := s.btree.MakeIter(), makeQueryEntry(sp) for iter.FirstOverlap(query); iter.Valid(); iter.NextOverlap(query) { - if err := f(iter.Cur().spanConfigEntry); err != nil { + entry := iter.Cur().spanConfigPairInterned + if err := f(entry.span, entry.conf(s.interner)); err != nil { if iterutil.Done(err) { err = nil } @@ -130,11 +135,11 @@ func (s *spanConfigStore) computeSplitKey( // in the iteration below we'll find all entries overlapping with the given // span but skipping over any start keys <= the given start key. In our // example this could be ['a', 'x') or ['b', 'y'). - var match spanConfigEntry + var match spanConfigPairInterned { iter, query := s.btree.MakeIter(), makeQueryEntry(sp) for iter.FirstOverlap(query); iter.Valid(); iter.NextOverlap(query) { - entry := iter.Cur().spanConfigEntry + entry := iter.Cur().spanConfigPairInterned if entry.span.Key.Compare(sp.Key) <= 0 { continue // more } @@ -177,12 +182,12 @@ func (s *spanConfigStore) computeSplitKey( // different tenant prefix from our original match, our original match key is // the split point we're interested in. If such an entry does not exist, then // too our original match key is the right split point. - var firstMatch spanConfigEntry + var firstMatch spanConfigPairInterned preSplitKeySp := roachpb.Span{Key: sp.Key, EndKey: match.span.Key} { iter, query := s.btree.MakeIter(), makeQueryEntry(preSplitKeySp) for iter.FirstOverlap(query); iter.Valid(); { - firstMatch = iter.Cur().spanConfigEntry + firstMatch = iter.Cur().spanConfigPairInterned break // we're done } if firstMatch.isEmpty() { @@ -192,7 +197,7 @@ func (s *spanConfigStore) computeSplitKey( if err != nil { log.Fatalf(ctx, "%v", err) } - if !firstMatch.config.Equal(match.config) || !firstMatchTenID.Equal(matchTenID) { + if firstMatch.internerID != match.internerID || firstMatchTenID.ToUint64() != matchTenID.ToUint64() { return roachpb.RKey(match.span.Key) } } @@ -200,17 +205,17 @@ func (s *spanConfigStore) computeSplitKey( // At least the first two entries with the given span have the same configs // and part of the same tenant range. Keep seeking ahead until we find a // different config or a different tenant. - var lastMatch spanConfigEntry + var lastMatch spanConfigPairInterned postSplitKeySp := roachpb.Span{Key: match.span.EndKey, EndKey: sp.EndKey} { iter, query := s.btree.MakeIter(), makeQueryEntry(postSplitKeySp) for iter.FirstOverlap(query); iter.Valid(); iter.NextOverlap(query) { - nextEntry := iter.Cur().spanConfigEntry + nextEntry := iter.Cur().spanConfigPairInterned _, entryTenID, err := keys.DecodeTenantPrefix(nextEntry.span.Key) if err != nil { log.Fatalf(ctx, "%v", err) } - if !nextEntry.config.Equal(match.config) || !entryTenID.Equal(matchTenID) { + if nextEntry.internerID != match.internerID || entryTenID.ToUint64() != matchTenID.ToUint64() { lastMatch = nextEntry break // we're done } @@ -229,24 +234,24 @@ func (s *spanConfigStore) computeSplitKey( // key. func (s *spanConfigStore) getSpanConfigForKey( ctx context.Context, key roachpb.RKey, -) (conf roachpb.SpanConfig, found bool, _ error) { +) (conf roachpb.SpanConfig, found bool) { sp := roachpb.Span{Key: key.AsRawKey(), EndKey: key.Next().AsRawKey()} iter, query := s.btree.MakeIter(), makeQueryEntry(sp) for iter.FirstOverlap(query); iter.Valid(); { - conf, found = iter.Cur().config, true + conf, found = iter.Cur().conf(s.interner), true break } if !found && log.ExpensiveLogEnabled(ctx, 1) { log.Warningf(ctx, "span config not found for %s", key.String()) } - return conf, found, nil + return conf, found } // apply takes an incremental set of updates and returns the spans/span<->config // entries deleted/added as a result of applying them. It also updates its state // by applying them if dryrun is false. func (s *spanConfigStore) apply( - dryrun bool, updates ...spanconfig.Update, + ctx context.Context, dryrun bool, updates ...spanconfig.Update, ) (deleted []roachpb.Span, added []spanConfigStoreEntry, err error) { if err := validateApplyArgs(updates...); err != nil { return nil, nil, err @@ -259,13 +264,17 @@ func (s *spanConfigStore) apply( }) updates = sorted // re-use the same variable - entriesToDelete, entriesToAdd := s.accumulateOpsFor(updates) + entriesToDelete, entriesToAdd, err := s.accumulateOpsFor(ctx, dryrun, updates) + if err != nil { + return nil, nil, err + } deleted = make([]roachpb.Span, len(entriesToDelete)) for i := range entriesToDelete { entry := &entriesToDelete[i] if !dryrun { s.btree.Delete(entry) + s.interner.remove(ctx, entry.internerID) } deleted[i] = entry.span } @@ -377,31 +386,31 @@ func (s *spanConfigStore) apply( // add {span=carry-over.span, conf=carry-over.conf} if non-empty // func (s *spanConfigStore) accumulateOpsFor( - updates []spanconfig.Update, -) (toDelete, toAdd []spanConfigStoreEntry) { - var carryOver spanConfigEntry + ctx context.Context, dryrun bool, updates []spanconfig.Update, +) (toDelete, toAdd []spanConfigStoreEntry, _ error) { + var carryOver spanConfigPair for _, update := range updates { - var carriedOver spanConfigEntry - carriedOver, carryOver = carryOver, spanConfigEntry{} + var carriedOver spanConfigPair + carriedOver, carryOver = carryOver, spanConfigPair{} if update.GetTarget().GetSpan().Overlaps(carriedOver.span) { gapBetweenUpdates := roachpb.Span{ Key: carriedOver.span.Key, EndKey: update.GetTarget().GetSpan().Key} if gapBetweenUpdates.Valid() { - toAdd = append(toAdd, s.makeEntry(gapBetweenUpdates, carriedOver.config)) + toAdd = append(toAdd, s.makeEntry(ctx, dryrun, gapBetweenUpdates, carriedOver.config)) } carryOverSpanAfterUpdate := roachpb.Span{ Key: update.GetTarget().GetSpan().EndKey, EndKey: carriedOver.span.EndKey} if carryOverSpanAfterUpdate.Valid() { - carryOver = spanConfigEntry{ + carryOver = spanConfigPair{ span: carryOverSpanAfterUpdate, config: carriedOver.config, } } } else if !carriedOver.isEmpty() { - toAdd = append(toAdd, s.makeEntry(carriedOver.span, carriedOver.config)) + toAdd = append(toAdd, s.makeEntry(ctx, dryrun, carriedOver.span, carriedOver.config)) } skipAddingSelf := false @@ -420,8 +429,9 @@ func (s *spanConfigStore) accumulateOpsFor( post = roachpb.Span{Key: inter.EndKey, EndKey: union.EndKey} ) + existingConf := existing.conf(s.interner) if update.Addition() { - if existing.span.Equal(update.GetTarget().GetSpan()) && existing.config.Equal(update.GetConfig()) { + if existing.span.Equal(update.GetTarget().GetSpan()) && existingConf.Equal(update.GetConfig()) { skipAddingSelf = true break // no-op; peep-hole optimization } @@ -443,7 +453,7 @@ func (s *spanConfigStore) accumulateOpsFor( // Re-add the non-intersecting span, if any. if pre.Valid() { - toAdd = append(toAdd, s.makeEntry(pre, existing.config)) + toAdd = append(toAdd, s.makeEntry(ctx, dryrun, pre, existingConf)) } } @@ -455,16 +465,16 @@ func (s *spanConfigStore) accumulateOpsFor( // up: [---------) // Carry over the non-intersecting span. - carryOver = spanConfigEntry{ + carryOver = spanConfigPair{ span: post, - config: existing.config, + config: existingConf, } } } if update.Addition() && !skipAddingSelf { // Add the update itself. - toAdd = append(toAdd, s.makeEntry(update.GetTarget().GetSpan(), update.GetConfig())) + toAdd = append(toAdd, s.makeEntry(ctx, dryrun, update.GetTarget().GetSpan(), update.GetConfig())) // TODO(irfansharif): If we're adding an entry, we could inspect the // entries before and after and check whether either of them have @@ -482,9 +492,9 @@ func (s *spanConfigStore) accumulateOpsFor( } if !carryOver.isEmpty() { - toAdd = append(toAdd, s.makeEntry(carryOver.span, carryOver.config)) + toAdd = append(toAdd, s.makeEntry(ctx, dryrun, carryOver.span, carryOver.config)) } - return toDelete, toAdd + return toDelete, toAdd, nil } // validateApplyArgs validates the supplied updates can be applied to the @@ -524,28 +534,56 @@ func validateApplyArgs(updates ...spanconfig.Update) error { return nil } -// spanConfigEntry captures a span <->config pair. -type spanConfigEntry struct { +// spanConfigPair represents a span <->config pair. +type spanConfigPair struct { span roachpb.Span config roachpb.SpanConfig } -func (s *spanConfigEntry) isEmpty() bool { +func (s *spanConfigPair) isEmpty() bool { return s.span.Equal(roachpb.Span{}) && s.config.IsEmpty() } +// spanConfigPairInterned represents a span <->config pair, but unlike +// spanConfigPair, doesn't embed the span config itself. It instead holds onto +// an interner ID which can be used to retrieve the corresponding config. +type spanConfigPairInterned struct { + span roachpb.Span + internerID +} + +func (s *spanConfigPairInterned) isEmpty() bool { + return s.span.Equal(roachpb.Span{}) && s.internerID == internerID(0) +} + +func (s *spanConfigPairInterned) conf(interner *interner) roachpb.SpanConfig { + conf, _ := interner.get(s.internerID) + return conf +} + //go:generate ../../util/interval/generic/gen.sh *spanConfigStoreEntry spanconfigstore type spanConfigStoreEntry struct { - spanConfigEntry + spanConfigPairInterned id uint64 } -func (s *spanConfigStore) makeEntry(sp roachpb.Span, conf roachpb.SpanConfig) spanConfigStoreEntry { - s.idAlloc++ +func (s *spanConfigStore) makeEntry( + ctx context.Context, dryrun bool, sp roachpb.Span, conf roachpb.SpanConfig, +) spanConfigStoreEntry { + if !dryrun { + s.treeIDAlloc++ + } + var internID internerID + if !dryrun || s.knobs.StoreInternConfigsInDryRuns { + internID = s.interner.add(ctx, conf) + } return spanConfigStoreEntry{ - spanConfigEntry: spanConfigEntry{span: sp, config: conf}, - id: s.idAlloc, + spanConfigPairInterned: spanConfigPairInterned{ + span: sp, + internerID: internID, + }, + id: s.treeIDAlloc, } } diff --git a/pkg/spanconfig/spanconfigstore/span_store_test.go b/pkg/spanconfig/spanconfigstore/span_store_test.go index 8974e5135d2b..f1f1ae3d07fd 100644 --- a/pkg/spanconfig/spanconfigstore/span_store_test.go +++ b/pkg/spanconfig/spanconfigstore/span_store_test.go @@ -119,7 +119,7 @@ func TestRandomized(t *testing.T) { }) for i := 0; i < numOps; i++ { updates := getRandomUpdates() - _, _, err := store.apply(false /* dryrun */, updates...) + _, _, err := store.apply(ctx, false /* dryrun */, updates...) require.NoError(t, err) for _, update := range updates { if testSpan.Overlaps(update.GetTarget().GetSpan()) { @@ -134,8 +134,8 @@ func TestRandomized(t *testing.T) { if !expFound { _ = store.forEachOverlapping(testSpan, - func(entry spanConfigEntry) error { - record, err := spanconfig.MakeRecord(spanconfig.MakeTargetFromSpan(entry.span), entry.config) + func(sp roachpb.Span, conf roachpb.SpanConfig) error { + record, err := spanconfig.MakeRecord(spanconfig.MakeTargetFromSpan(sp), conf) require.NoError(t, err) t.Fatalf("found unexpected entry: %s", spanconfigtestutils.PrintSpanConfigRecord(t, record)) @@ -143,19 +143,22 @@ func TestRandomized(t *testing.T) { }, ) } else { - var foundEntry spanConfigEntry + var foundSpanConfigPair spanConfigPair _ = store.forEachOverlapping(testSpan, - func(entry spanConfigEntry) error { - if !foundEntry.isEmpty() { - record, err := spanconfig.MakeRecord(spanconfig.MakeTargetFromSpan(entry.span), entry.config) + func(sp roachpb.Span, conf roachpb.SpanConfig) error { + if !foundSpanConfigPair.isEmpty() { + record, err := spanconfig.MakeRecord(spanconfig.MakeTargetFromSpan(sp), conf) require.NoError(t, err) t.Fatalf("expected single overlapping entry, found second: %s", spanconfigtestutils.PrintSpanConfigRecord(t, record)) } - foundEntry = entry + foundSpanConfigPair = spanConfigPair{ + span: sp, + config: conf, + } // Check that the entry is exactly what we'd expect. - gotSpan, gotConfig := entry.span, entry.config + gotSpan, gotConfig := sp, conf require.Truef(t, gotSpan.Contains(testSpan), "improper result: expected retrieved span (%s) to contain test span (%s)", spanconfigtestutils.PrintSpan(gotSpan), spanconfigtestutils.PrintSpan(testSpan)) @@ -170,39 +173,41 @@ func TestRandomized(t *testing.T) { // Ensure that the config accessed through the StoreReader interface is // the same as above. - storeReaderConfig, found, err := store.getSpanConfigForKey(ctx, roachpb.RKey(testSpan.Key)) - require.NoError(t, err) + storeReaderConfig, found := store.getSpanConfigForKey(ctx, roachpb.RKey(testSpan.Key)) require.True(t, found) - require.True(t, foundEntry.config.Equal(storeReaderConfig)) + require.True(t, foundSpanConfigPair.config.Equal(storeReaderConfig)) } everythingSpan := spanconfigtestutils.ParseSpan(t, fmt.Sprintf("[%s,%s)", string(alphabet[0]), string(alphabet[len(alphabet)-1]))) - var lastOverlapping spanConfigEntry + var lastOverlapping spanConfigPair require.NoError(t, store.forEachOverlapping(everythingSpan, - func(cur spanConfigEntry) error { - log.Infof(ctx, "set %s:%s", spanconfigtestutils.PrintSpan(cur.span), spanconfigtestutils.PrintSpanConfig(cur.config)) + func(sp roachpb.Span, conf roachpb.SpanConfig) error { + log.Infof(ctx, "set %s:%s", spanconfigtestutils.PrintSpan(sp), spanconfigtestutils.PrintSpanConfig(conf)) // All spans are expected to be valid. - require.True(t, cur.span.Valid(), + require.True(t, sp.Valid(), "expected to only find valid spans, found %s", - spanconfigtestutils.PrintSpan(cur.span), + spanconfigtestutils.PrintSpan(sp), ) if !lastOverlapping.isEmpty() { // Span configs are returned in strictly sorted order. - require.True(t, lastOverlapping.span.Key.Compare(cur.span.Key) < 0, + require.True(t, lastOverlapping.span.Key.Compare(sp.Key) < 0, "expected to find spans in strictly sorted order, found %s then %s", - spanconfigtestutils.PrintSpan(lastOverlapping.span), spanconfigtestutils.PrintSpan(cur.span)) + spanconfigtestutils.PrintSpan(lastOverlapping.span), spanconfigtestutils.PrintSpan(sp)) // Span configs must also be non-overlapping. - require.Falsef(t, lastOverlapping.span.Overlaps(cur.span), + require.Falsef(t, lastOverlapping.span.Overlaps(sp), "expected non-overlapping spans, found %s and %s", - spanconfigtestutils.PrintSpan(lastOverlapping.span), spanconfigtestutils.PrintSpan(cur.span)) + spanconfigtestutils.PrintSpan(lastOverlapping.span), spanconfigtestutils.PrintSpan(sp)) } - lastOverlapping = cur + lastOverlapping = spanConfigPair{ + span: sp, + config: conf, + } return nil }, )) @@ -214,14 +219,20 @@ func TestRandomized(t *testing.T) { ) numOverlappingWithQuerySp := 0 - var firstOverlappingWithQuerySp, lastOverlappingWithQuerySp spanConfigEntry + var firstOverlappingWithQuerySp, lastOverlappingWithQuerySp spanConfigPair require.NoError(t, store.forEachOverlapping(querySpan, - func(cur spanConfigEntry) error { + func(sp roachpb.Span, conf roachpb.SpanConfig) error { if numOverlappingWithQuerySp == 0 { - firstOverlappingWithQuerySp = cur + firstOverlappingWithQuerySp = spanConfigPair{ + span: sp, + config: conf, + } } numOverlappingWithQuerySp++ - lastOverlappingWithQuerySp = cur + lastOverlappingWithQuerySp = spanConfigPair{ + span: sp, + config: conf, + } return nil }, )) @@ -232,8 +243,7 @@ func TestRandomized(t *testing.T) { require.Truef(t, querySpan.ProperlyContainsKey(curSplitKey.AsRawKey()), "invalid split key %s (over span %s)", curSplitKey, querySpan) - confAtCurSplitKey, found, err := store.getSpanConfigForKey(ctx, curSplitKey) - require.NoError(t, err) + confAtCurSplitKey, found := store.getSpanConfigForKey(ctx, curSplitKey) require.True(t, found) if i == 0 { @@ -269,11 +279,11 @@ func TestRandomized(t *testing.T) { require.NoError(t, store.forEachOverlapping(roachpb.Span{ Key: lastSplitKey.AsRawKey(), EndKey: curSplitKey.AsRawKey(), - }, func(entry spanConfigEntry) error { - require.Truef(t, confAtLastSplitKey.Equal(entry.config), + }, func(sp roachpb.Span, conf roachpb.SpanConfig) error { + require.Truef(t, confAtLastSplitKey.Equal(conf), "expected identical configs, found %s:%s and %s:%s", lastSplitKey, spanconfigtestutils.PrintSpanConfig(confAtLastSplitKey), - entry.span.Key, spanconfigtestutils.PrintSpanConfig(entry.config), + sp.Key, spanconfigtestutils.PrintSpanConfig(conf), ) return nil })) diff --git a/pkg/spanconfig/spanconfigstore/store.go b/pkg/spanconfig/spanconfigstore/store.go index 092039cfb034..28034c5fd30c 100644 --- a/pkg/spanconfig/spanconfigstore/store.go +++ b/pkg/spanconfig/spanconfigstore/store.go @@ -106,10 +106,7 @@ func (s *Store) GetSpanConfigForKey( func (s *Store) getSpanConfigForKeyRLocked( ctx context.Context, key roachpb.RKey, ) (roachpb.SpanConfig, error) { - conf, found, err := s.mu.spanConfigStore.getSpanConfigForKey(ctx, key) - if err != nil { - return roachpb.SpanConfig{}, err - } + conf, found := s.mu.spanConfigStore.getSpanConfigForKey(ctx, key) if !found { conf = s.fallback } @@ -120,7 +117,7 @@ func (s *Store) getSpanConfigForKeyRLocked( func (s *Store) Apply( ctx context.Context, dryrun bool, updates ...spanconfig.Update, ) (deleted []spanconfig.Target, added []spanconfig.Record) { - deleted, added, err := s.applyInternal(dryrun, updates...) + deleted, added, err := s.applyInternal(ctx, dryrun, updates...) if err != nil { log.Fatalf(ctx, "%v", err) } @@ -133,12 +130,12 @@ func (s *Store) ForEachOverlappingSpanConfig( ) error { s.mu.RLock() defer s.mu.RUnlock() - return s.mu.spanConfigStore.forEachOverlapping(span, func(entry spanConfigEntry) error { - config, err := s.getSpanConfigForKeyRLocked(ctx, roachpb.RKey(entry.span.Key)) + return s.mu.spanConfigStore.forEachOverlapping(span, func(sp roachpb.Span, conf roachpb.SpanConfig) error { + config, err := s.getSpanConfigForKeyRLocked(ctx, roachpb.RKey(sp.Key)) if err != nil { return err } - return f(entry.span, config) + return f(sp, config) }) } @@ -154,7 +151,7 @@ func (s *Store) Copy(ctx context.Context) *Store { } func (s *Store) applyInternal( - dryrun bool, updates ...spanconfig.Update, + ctx context.Context, dryrun bool, updates ...spanconfig.Update, ) (deleted []spanconfig.Target, added []spanconfig.Record, err error) { s.mu.Lock() defer s.mu.Unlock() @@ -174,7 +171,7 @@ func (s *Store) applyInternal( return nil, nil, errors.AssertionFailedf("unknown target type") } } - deletedSpans, addedEntries, err := s.mu.spanConfigStore.apply(dryrun, spanStoreUpdates...) + deletedSpans, addedEntries, err := s.mu.spanConfigStore.apply(ctx, dryrun, spanStoreUpdates...) if err != nil { return nil, nil, err } @@ -184,8 +181,10 @@ func (s *Store) applyInternal( } for _, entry := range addedEntries { - record, err := spanconfig.MakeRecord(spanconfig.MakeTargetFromSpan(entry.span), - entry.config) + record, err := spanconfig.MakeRecord( + spanconfig.MakeTargetFromSpan(entry.span), + entry.conf(s.mu.spanConfigStore.interner), + ) if err != nil { return nil, nil, err } @@ -216,8 +215,8 @@ func (s *Store) Iterate(f func(spanconfig.Record) error) error { } return s.mu.spanConfigStore.forEachOverlapping( keys.EverythingSpan, - func(s spanConfigEntry) error { - record, err := spanconfig.MakeRecord(spanconfig.MakeTargetFromSpan(s.span), s.config) + func(sp roachpb.Span, conf roachpb.SpanConfig) error { + record, err := spanconfig.MakeRecord(spanconfig.MakeTargetFromSpan(sp), conf) if err != nil { return err } diff --git a/pkg/spanconfig/spanconfigstore/store_test.go b/pkg/spanconfig/spanconfigstore/store_test.go index 125c361972b1..c06fdfd1b6ba 100644 --- a/pkg/spanconfig/spanconfigstore/store_test.go +++ b/pkg/spanconfig/spanconfigstore/store_test.go @@ -29,9 +29,9 @@ import ( // TestingApplyInternal exports an internal method for testing purposes. func (s *Store) TestingApplyInternal( - dryrun bool, updates ...spanconfig.Update, + ctx context.Context, dryrun bool, updates ...spanconfig.Update, ) (deleted []spanconfig.Target, added []spanconfig.Record, err error) { - return s.applyInternal(dryrun, updates...) + return s.applyInternal(ctx, dryrun, updates...) } // TestingSplitKeys returns the computed list of range split points between @@ -114,6 +114,7 @@ func TestDataDriven(t *testing.T) { cluster.MakeTestingClusterSettings(), &spanconfig.TestingKnobs{ StoreIgnoreCoalesceAdjacentExceptions: true, + StoreInternConfigsInDryRuns: true, }, ) datadriven.RunTest(t, path, func(t *testing.T, d *datadriven.TestData) string { @@ -122,7 +123,7 @@ func TestDataDriven(t *testing.T) { case "apply": updates := spanconfigtestutils.ParseStoreApplyArguments(t, d.Input) dryrun := d.HasArg("dryrun") - deleted, added, err := store.TestingApplyInternal(dryrun, updates...) + deleted, added, err := store.TestingApplyInternal(ctx, dryrun, updates...) if err != nil { return fmt.Sprintf("err: %v", err) } @@ -294,15 +295,14 @@ func BenchmarkStoreComputeSplitKey(b *testing.B) { query := spanconfigtestutils.ParseSpan(b, fmt.Sprintf("[%08d, %08d)", 0, numEntries)) - // XXX: Uncomment. - //overlapping := 0 - //require.NoError(b, store.ForEachOverlappingSpanConfig(ctx, makeQueryEntry, - // func(_ roachpb.Span, _ roachpb.SpanConfig) error { - // overlapping++ - // return nil - // }, - //)) - //require.Equal(b, overlapping, numEntries) + overlapping := 0 + require.NoError(b, store.ForEachOverlappingSpanConfig(ctx, query, + func(_ roachpb.Span, _ roachpb.SpanConfig) error { + overlapping++ + return nil + }, + )) + require.Equal(b, overlapping, numEntries) b.ResetTimer() for i := 0; i < b.N; i++ { diff --git a/pkg/spanconfig/testing_knobs.go b/pkg/spanconfig/testing_knobs.go index 039da8e862b0..501629285f51 100644 --- a/pkg/spanconfig/testing_knobs.go +++ b/pkg/spanconfig/testing_knobs.go @@ -117,6 +117,10 @@ type TestingKnobs struct { // spanconfig.{host,tenant}_coalesce_adjacent.enabled. It also allows // coalescing system database ranges for the host tenant. StoreIgnoreCoalesceAdjacentExceptions bool + + // StoreInternConfigsInDryRuns, if set, will intern span configs even when + // applying mutations in dry run mode. + StoreInternConfigsInDryRuns bool } // ModuleTestingKnobs is part of the base.ModuleTestingKnobs interface.