diff --git a/pkg/btree/btree_generic_test.go b/pkg/btree/btree_generic_test.go index 4d36ef2a7c..4958d35ea9 100644 --- a/pkg/btree/btree_generic_test.go +++ b/pkg/btree/btree_generic_test.go @@ -818,6 +818,146 @@ func TestCloneConcurrentOperationsG(t *testing.T) { } } +// assertSizeInfo checks the per-node index bookkeeping that GetAt and +// GetWithIndex rely on. A tree can be correctly ordered, so that Ascend returns +// everything in the right sequence, while its indices are wrong; then GetAt and +// GetWithIndex silently return the wrong item or the wrong rank. +func assertSizeInfo(t *testing.T, desc string, tr *BTreeG[Int]) { + assertEq(t, desc+" root length", tr.getRootLength(), tr.Len()) + if tr.Len() == 0 { + return + } + min, _ := tr.Min() + assertEq(t, desc+" min", tr.GetAt(0), min) + max, _ := tr.Max() + assertEq(t, desc+" max", tr.GetAt(tr.Len()-1), max) + // GetAt must agree with Ascend, and every item's own rank must round-trip. + i := 0 + tr.Ascend(func(item Int) bool { + assertEq(t, desc+" get k-th", tr.GetAt(i), item) + got, rank := tr.GetWithIndex(item) + assertEq(t, desc+" get", got, item) + assertEq(t, desc+" rank", rank, i) + i++ + return true + }) + assertEq(t, desc+" ascend count", i, tr.Len()) +} + +// TestBTreeSizeInfoAfterCloneG checks that the index bookkeeping survives +// copy-on-write cloning, on the clone and on the original. +// +// TestBTreeSizeInfo covers the indices but never clones, and +// TestCloneConcurrentOperationsG clones but only compares Ascend order, so +// nothing exercised mutableFor's copying of the indices array. Callers that then +// read a clone through GetAt or GetWithIndex would get wrong answers with no +// panic and no error. +func TestBTreeSizeInfoAfterCloneG(t *testing.T) { + const treeSize = 2000 + tr := NewG[Int](*btreeDegree) + clones := []*BTreeG[Int]{} + for i, item := range perm(treeSize) { + tr.ReplaceOrInsert(item) + if i%(treeSize/10) == 0 { + // Cloning makes every existing node foreign to tr's write context, so + // the following inserts have to copy each node they touch. + clones = append(clones, tr.Clone()) + assertSizeInfo(t, "after clone original", tr) + assertSizeInfo(t, "after clone clone", clones[len(clones)-1]) + } + } + assertSizeInfo(t, "filled original", tr) + + // Deleting from the original must not disturb any clone, and both sides must + // keep consistent indices while nodes are being split, merged and freed. + lengths := make([]int, len(clones)) + for i, clone := range clones { + lengths[i] = clone.Len() + } + for _, item := range perm(treeSize) { + if item%3 == 0 { + tr.Delete(item) + } + } + assertSizeInfo(t, "after delete original", tr) + for i, clone := range clones { + assertEq(t, "clone length unchanged", clone.Len(), lengths[i]) + assertSizeInfo(t, "after delete clone", clone) + } +} + +// TestCloneReadWhileWritingOriginalG reads a clone while the original is being +// mutated, which is how PD uses Clone: one snapshot serving a scan while the +// heartbeat path keeps writing. Run it with -race. +func TestCloneReadWhileWritingOriginalG(t *testing.T) { + const treeSize = 5000 + tr := NewG[Int](*btreeDegree) + for _, item := range perm(treeSize) { + tr.ReplaceOrInsert(item) + } + snap := tr.Clone() + want := rang(treeSize) + + var writers, readers sync.WaitGroup + done := make(chan struct{}) + // BTreeG does not support concurrent writes, so serialise them the way PD + // does, under one lock. Only the original tree is written; a clone is + // read-only by contract, and no lock is taken to read it. + var writeMu syncutil.Mutex + + // Writers: keep inserting and deleting items outside the snapshot's range, and + // replacing items inside it, so nodes are copied, split and merged throughout. + for w := range 4 { + writers.Add(1) + go func() { + defer writers.Done() + for i := 0; ; i++ { + select { + case <-done: + return + default: + } + item := Int(treeSize + w*treeSize + i%treeSize) + writeMu.Lock() + tr.ReplaceOrInsert(item) + tr.ReplaceOrInsert(Int(i % treeSize)) + tr.Delete(item) + writeMu.Unlock() + } + }() + } + + // Readers: the snapshot must not change in any observable way. + for range 4 { + readers.Add(1) + go func() { + defer readers.Done() + for range 50 { + if got := snap.Len(); got != treeSize { + t.Errorf("snapshot length changed: got %d want %d", got, treeSize) + return + } + if got := all(snap); !reflect.DeepEqual(want, got) { + t.Errorf("snapshot contents changed: got %d items", len(got)) + return + } + for k := range treeSize { + if got := snap.GetAt(k); got != Int(k) { + t.Errorf("snapshot GetAt(%d) = %v", k, got) + return + } + } + } + }() + } + readers.Wait() + close(done) + writers.Wait() + + assertSizeInfo(t, "snapshot after concurrent writes", snap) + assertEq(t, "snapshot contents after concurrent writes", all(snap), want) +} + func BenchmarkDeleteAndRestore(b *testing.B) { items := perm(16392) b.ResetTimer() diff --git a/pkg/core/region.go b/pkg/core/region.go index 40701a5137..ceb25f7115 100644 --- a/pkg/core/region.go +++ b/pkg/core/region.go @@ -1086,7 +1086,7 @@ func (r *RegionsInfo) GetRegion(regionID uint64) *RegionInfo { func (r *RegionsInfo) getRegionLocked(regionID uint64) *RegionInfo { if item := r.regions[regionID]; item != nil { - return item.RegionInfo + return item.getRegion() } return nil } @@ -1097,7 +1097,7 @@ func (r *RegionsInfo) CheckAndPutRegion(region *RegionInfo) []*RegionInfo { origin := r.getRegionLocked(region.GetID()) var ols []*RegionInfo if origin == nil || !bytes.Equal(origin.GetStartKey(), region.GetStartKey()) || !bytes.Equal(origin.GetEndKey(), region.GetEndKey()) { - ols = r.tree.overlaps(®ionItem{RegionInfo: region}) + ols = r.tree.overlaps(newRegionItem(region)) } err := check(region, origin, ols) if err != nil { @@ -1133,7 +1133,7 @@ func (r *RegionsInfo) AtomicCheckAndPutRegion(ctx *MetaProcessContext, region *R var ols []*RegionInfo origin := r.getRegionLocked(region.GetID()) if origin == nil || !bytes.Equal(origin.GetStartKey(), region.GetStartKey()) || !bytes.Equal(origin.GetEndKey(), region.GetEndKey()) { - ols = r.tree.overlaps(®ionItem{RegionInfo: region}) + ols = r.tree.overlaps(newRegionItem(region)) } tracer.OnCheckOverlapsFinished() err := check(region, origin, ols) @@ -1159,7 +1159,7 @@ func (r *RegionsInfo) CheckAndPutRootTree(ctx *MetaProcessContext, region *Regio var ols []*RegionInfo origin := r.getRegionLocked(region.GetID()) if origin == nil || !bytes.Equal(origin.GetStartKey(), region.GetStartKey()) || !bytes.Equal(origin.GetEndKey(), region.GetEndKey()) { - ols = r.tree.overlaps(®ionItem{RegionInfo: region}) + ols = r.tree.overlaps(newRegionItem(region)) } tracer.OnCheckOverlapsFinished() err := check(region, origin, ols) @@ -1201,7 +1201,7 @@ func (r *RegionsInfo) UpdateSubTreeOrderInsensitive(region *RegionInfo) { defer r.st.Unlock() originItem, ok := r.subRegions[region.GetID()] if ok { - origin = originItem.RegionInfo + origin = originItem.getRegion() } rangeChanged := true if origin != nil { @@ -1240,14 +1240,14 @@ func (r *RegionsInfo) preUpdateSubTreeLocked( // to keep region tree consistent with subtree, we need to drop this update. if tree, ok := r.subRegions[region.GetID()]; ok { // Fetch the origin region info from the subtree again to ensure it is up-to-date. - origin := tree.RegionInfo + origin := tree.getRegion() r.updateSubTreeStat(origin, region) // overlapTree is the only ref-counted subtree and the shared item is // repointed to region below, so transfer its reference here. Otherwise // a flow-only update leaves the live region stuck at ref 1 while it is // still present in the subtree. r.overlapTree.updateRef(origin, region) - tree.RegionInfo = region + tree.setRegion(region) } return true } @@ -1270,7 +1270,7 @@ func (r *RegionsInfo) updateSubTreeLocked(rangeChanged bool, overlaps []*RegionI } } // Reinsert the region into all subtrees. - item := ®ionItem{region} + item := newRegionItem(region) r.subRegions[region.GetID()] = item r.overlapTree.update(item, false) // Add leaders and followers. @@ -1302,7 +1302,7 @@ func (r *RegionsInfo) updateSubTreeLocked(rangeChanged bool, overlaps []*RegionI } func (r *RegionsInfo) getOverlapRegionFromOverlapTreeLocked(region *RegionInfo) []*RegionInfo { - return r.overlapTree.overlaps(®ionItem{RegionInfo: region}) + return r.overlapTree.overlaps(newRegionItem(region)) } // GetRelevantRegions returns the relevant regions for a given region. @@ -1311,7 +1311,7 @@ func (r *RegionsInfo) GetRelevantRegions(region *RegionInfo) (origin *RegionInfo defer r.t.RUnlock() origin = r.getRegionLocked(region.GetID()) if origin == nil || !bytes.Equal(origin.GetStartKey(), region.GetStartKey()) || !bytes.Equal(origin.GetEndKey(), region.GetEndKey()) { - return origin, r.tree.overlaps(®ionItem{RegionInfo: region}) + return origin, r.tree.overlaps(newRegionItem(region)) } return } @@ -1355,11 +1355,10 @@ func (r *RegionsInfo) setRegionLocked(region *RegionInfo, withOverlaps bool, ol if item = r.regions[region.GetID()]; item != nil { // If this ID already exists, use the existing regionItem and pick out the origin. - origin = item.RegionInfo + origin = item.getRegion() rangeChanged = !origin.rangeEqualsTo(region) if rangeChanged { // Delete itself in regionTree so that overlaps will not contain itself. - // Because the regionItem is reused, there is no need to delete it in the regionMap. idx := -1 for i, o := range ol { if o.GetID() == region.GetID() { @@ -1370,19 +1369,35 @@ func (r *RegionsInfo) setRegionLocked(region *RegionInfo, withOverlaps bool, ol if idx >= 0 { ol = append(ol[:idx], ol[idx+1:]...) } + // Remove the origin first, while the old item still holds it: remove + // locates the tree item by the origin's own key range, and adjusts the + // tree statistics and the reference count through it. r.tree.remove(origin) - // Update the RegionInfo in the regionItem. - item.RegionInfo = region + // The key range changed, so the old item cannot be reused. An item that + // has been in a tree must never change its range, or an outstanding + // copy-on-write snapshot would keep it at its old position while + // reporting the new range. Allocate a fresh item and repoint the map + // entry; see regionItem. + // + // The assignment below cannot be undone by the overlap cleanup further + // down, because overlaps never contains this region itself: either it + // was filtered out of ol just above, or tree.update recomputes the + // overlaps after the tree.remove call above has already taken this + // region out of the tree. + item = newRegionItem(region) + r.regions[region.GetID()] = item } else { // If the range is not changed, only the statistical on the regionTree needs to be updated. r.tree.updateStat(origin, region) - // Update the RegionInfo in the regionItem. - item.RegionInfo = region + // The range is unchanged, so the item may keep its position and only + // its value is replaced. This is the region-heartbeat path, and it is + // the only place allowed to replace a tree-resident item's value. + item.setRegion(region) return origin, nil, rangeChanged } } else { // If this ID does not exist, generate a new regionItem and save it in the regionMap. - item = ®ionItem{RegionInfo: region} + item = newRegionItem(region) r.regions[region.GetID()] = item } var overlaps []*RegionInfo @@ -1438,6 +1453,19 @@ func (r *RegionsInfo) updateSubTreeStat(origin *RegionInfo, region *RegionInfo) updatePeersStat(r.pendingPeers, region.GetPendingPeers()) } +// snapshotRootTree returns a read-only snapshot of the root region tree, so that +// a long scan does not have to hold r.t for its whole duration. See +// rootRangeSnapshot for what the snapshot does and does not guarantee. +// +// This takes the write lock rather than the read lock, because btree Clone +// rewrites the source tree's copy-on-write context. The critical section is O(1): +// a struct copy and two small allocations, independent of the number of regions. +func (r *RegionsInfo) snapshotRootTree() *rootRangeSnapshot { + r.t.Lock() + defer r.t.Unlock() + return r.tree.snapshot() +} + // TreeLen returns the RegionsInfo tree length(now only used in test) func (r *RegionsInfo) TreeLen() int { r.t.RLock() @@ -1449,7 +1477,7 @@ func (r *RegionsInfo) TreeLen() int { func (r *RegionsInfo) GetOverlaps(region *RegionInfo) []*RegionInfo { r.t.RLock() defer r.t.RUnlock() - return r.tree.overlaps(®ionItem{RegionInfo: region}) + return r.tree.overlaps(newRegionItem(region)) } // RemoveRegion removes RegionInfo from regionTree and regionMap @@ -1581,7 +1609,7 @@ func (r *RegionsInfo) GetRegions() []*RegionInfo { defer r.t.RUnlock() regions := make([]*RegionInfo, 0, len(r.regions)) for _, item := range r.regions { - regions = append(regions, item.RegionInfo) + regions = append(regions, item.getRegion()) } return regions } @@ -1884,7 +1912,7 @@ func (r *RegionsInfo) GetMetaRegions() []*metapb.Region { defer r.t.RUnlock() regions := make([]*metapb.Region, 0, len(r.regions)) for _, item := range r.regions { - regions = append(regions, typeutil.DeepClone(item.meta, RegionFactory)) + regions = append(regions, typeutil.DeepClone(item.getRegion().meta, RegionFactory)) } return regions } @@ -1998,7 +2026,7 @@ func (r *RegionsInfo) GetLeader(storeID uint64, region *RegionInfo) *RegionInfo r.st.RLock() defer r.st.RUnlock() if leaders, ok := r.leaders[storeID]; ok { - return leaders.find(®ionItem{RegionInfo: region}).RegionInfo + return leaders.find(newRegionItem(region)).getRegion() } return nil } @@ -2008,7 +2036,7 @@ func (r *RegionsInfo) GetFollower(storeID uint64, region *RegionInfo) *RegionInf r.st.RLock() defer r.st.RUnlock() if followers, ok := r.followers[storeID]; ok { - return followers.find(®ionItem{RegionInfo: region}).RegionInfo + return followers.find(newRegionItem(region)).getRegion() } return nil } @@ -2144,6 +2172,13 @@ func (r *RegionsInfo) BatchScanRegions(keyRanges *keyutil.KeyRanges, opts ...Bat opt(scanOptions) } + // This holds r.t for the whole scan, which a rootRangeSnapshot could avoid. + // Measured on 1M regions, that is a bad trade here: the scan is short and the + // call rate is high, so taking r.t exclusively once per call to clone the tree + // costs far more than the read lock it replaces. At GOMAXPROCS=4 with + // concurrent heartbeat writers, eight-region scans went from 3.1us to 20.7us + // per call. Only a shared, invalidated-on-structural-change snapshot would help + // this shape of caller; see rootRangeSnapshot. r.t.RLock() defer r.t.RUnlock() for _, keyRange := range krs { @@ -2235,31 +2270,24 @@ func (r *RegionsInfo) GetRegionSizeByRange(startKey, endKey []byte) int64 { defer r.t.RUnlock() return r.tree.totalSize } + // Walk a snapshot instead of scanning the live tree in chunks. + // + // The chunked version took r.t once per ScanRegionLimit regions, so on a large + // cluster one call acquired the lock hundreds of times, and each acquisition + // let at most one waiting writer through. It was also not exact: it resumed + // each chunk at the previous region's end key, and scanRange resolves that key + // to the region *containing* it, so a merge across a chunk boundary made the + // merged region contribute its whole size to a range that had already been + // counted in part. + snap := r.snapshotRootTree() var size int64 - for { - r.t.RLock() - var cnt int - r.tree.scanRange(startKey, func(region *RegionInfo) bool { - if len(endKey) > 0 && bytes.Compare(region.GetStartKey(), endKey) >= 0 { - return false - } - if cnt >= ScanRegionLimit { - return false - } - cnt++ - startKey = region.GetEndKey() - size += region.GetApproximateSize() - return true - }) - r.t.RUnlock() - if cnt == 0 { - break - } - if len(startKey) == 0 { - break + snap.scanRange(startKey, func(region *RegionInfo) bool { + if len(endKey) > 0 && bytes.Compare(region.GetStartKey(), endKey) >= 0 { + return false } - } - + size += region.GetApproximateSize() + return true + }) return size } @@ -2306,10 +2334,10 @@ func (r *RegionsInfo) GetAdjacentRegions(region *RegionInfo) (prev, next *Region p, n := r.tree.getAdjacentRegions(region) // check key to avoid key range hole if p != nil && bytes.Equal(p.GetEndKey(), region.GetStartKey()) { - prev = p.RegionInfo + prev = p.getRegion() } if n != nil && bytes.Equal(region.GetEndKey(), n.GetStartKey()) { - next = n.RegionInfo + next = n.getRegion() } return prev, next } diff --git a/pkg/core/region_test.go b/pkg/core/region_test.go index 00b40bf062..7256439569 100644 --- a/pkg/core/region_test.go +++ b/pkg/core/region_test.go @@ -15,6 +15,7 @@ package core import ( + "bytes" "crypto/rand" "encoding/json" "fmt" @@ -22,6 +23,7 @@ import ( mrand "math/rand/v2" "strconv" "strings" + "sync" "sync/atomic" "testing" "time" @@ -35,6 +37,7 @@ import ( "github.com/tikv/pd/pkg/errs" "github.com/tikv/pd/pkg/utils/keyutil" + "github.com/tikv/pd/pkg/utils/syncutil" ) func TestNeedMerge(t *testing.T) { @@ -435,14 +438,14 @@ func TestRegionMap(t *testing.T) { re := require.New(t) rm := make(map[uint64]*regionItem) checkMap(re, rm) - rm[1] = ®ionItem{RegionInfo: regionInfo(1)} + rm[1] = newRegionItem(regionInfo(1)) checkMap(re, rm, 1) - rm[2] = ®ionItem{RegionInfo: regionInfo(2)} - rm[3] = ®ionItem{RegionInfo: regionInfo(3)} + rm[2] = newRegionItem(regionInfo(2)) + rm[3] = newRegionItem(regionInfo(3)) checkMap(re, rm, 1, 2, 3) - rm[3] = ®ionItem{RegionInfo: regionInfo(3)} + rm[3] = newRegionItem(regionInfo(3)) delete(rm, 4) checkMap(re, rm, 1, 2, 3) @@ -450,7 +453,7 @@ func TestRegionMap(t *testing.T) { delete(rm, 1) checkMap(re, rm, 2) - rm[3] = ®ionItem{RegionInfo: regionInfo(3)} + rm[3] = newRegionItem(regionInfo(3)) checkMap(re, rm, 2, 3) } @@ -530,6 +533,38 @@ func TestSetRegionConcurrence(t *testing.T) { re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/core/UpdateSubTree")) } +// TestSetRegionRangeChangedRepointsMap checks that a region whose key range +// changes ends up reachable through both the region map and the region tree. +// +// A range change allocates a fresh regionItem and repoints r.regions at it, so +// the map entry must survive the overlap cleanup that follows. If it did not, +// GetRegion would keep returning the pre-split region for ever, with no error and +// no metric to show it. +func TestSetRegionRangeChangedRepointsMap(t *testing.T) { + re := require.New(t) + regions := NewRegionsInfo() + + origin := NewTestRegionInfo(1, 1, []byte("a"), []byte("c")) + _, err := regions.AtomicCheckAndPutRegion(ContextTODO(), origin) + re.NoError(err) + re.Equal(int32(2), origin.GetRef()) // root tree plus overlap tree + + // Shrink the region: same ID, same start key, smaller end key, higher version. + shrunk := NewTestRegionInfo(1, 1, []byte("a"), []byte("b"), SetRegionVersion(2)) + _, err = regions.AtomicCheckAndPutRegion(ContextTODO(), shrunk) + re.NoError(err) + + re.Same(shrunk, regions.GetRegion(1)) + re.Same(shrunk, regions.tree.search([]byte("a"))) + re.Equal(1, regions.tree.length()) + re.Same(shrunk, regions.regions[1].getRegion()) + // The fresh item took over the reference the old one held. + re.Equal(int32(2), shrunk.GetRef()) + re.Zero(origin.GetRef()) + // [b, c) is no longer covered by any region. + re.Nil(regions.tree.search([]byte("b"))) +} + func TestSetRegion(t *testing.T) { re := require.New(t) regions := NewRegionsInfo() @@ -768,6 +803,9 @@ func BenchmarkRandomRegion(b *testing.B) { } } +// BenchmarkRandomSetRegion re-reports the same *RegionInfo after mutating it in +// place, so origin == region and updateStat adds and subtracts the same values. +// It does not model a region heartbeat; use BenchmarkSetRegionSameRange for that. func BenchmarkRandomSetRegion(b *testing.B) { regions := NewRegionsInfo() var items []*RegionInfo @@ -811,14 +849,35 @@ func TestGetRegionSizeByRange(t *testing.T) { origin, overlaps, rangeChanged := regions.SetRegion(region) regions.UpdateSubTree(region, origin, overlaps, rangeChanged) } + re := require.New(t) totalSize := regions.GetRegionSizeByRange([]byte(""), []byte("")) - require.Equal(t, int64(nums*10), totalSize) + re.Equal(int64(nums*10), totalSize) for i := 1; i < 10; i++ { verifyNum := nums / i endKey := fmt.Sprintf("%20d", verifyNum) totalSize := regions.GetRegionSizeByRange([]byte(""), []byte(endKey)) - require.Equal(t, int64(verifyNum*10), totalSize) + re.Equal(int64(verifyNum*10), totalSize) } + + // A range starting part-way through the keyspace, still on region boundaries. + re.Equal(int64(100*10), regions.GetRegionSizeByRange( + []byte(fmt.Sprintf("%20d", 500)), []byte(fmt.Sprintf("%20d", 600)))) + + // Boundaries that fall *inside* a region rather than on its edges. scanRange + // starts at the region containing the start key, and stops at the first region + // whose start key is at or past the end key, so both partial regions count in + // full: [500, 501) through [600, 601) is 101 regions. + re.Equal(int64(101*10), regions.GetRegionSizeByRange( + []byte(fmt.Sprintf("%20d", 500)+"x"), []byte(fmt.Sprintf("%20d", 600)+"x"))) + + // A range extending past the last region's start key. The last region has an + // empty end key, so it is the final one counted. + re.Equal(int64(2*10), regions.GetRegionSizeByRange( + []byte(fmt.Sprintf("%20d", nums-2)), []byte(""))) + + // An empty range in the middle of one region contributes that region. + re.Equal(int64(10), regions.GetRegionSizeByRange( + []byte(fmt.Sprintf("%20d", 500)+"a"), []byte(fmt.Sprintf("%20d", 500)+"b"))) } func BenchmarkRandomSetRegionWithGetRegionSizeByRange(b *testing.B) { @@ -836,18 +895,21 @@ func BenchmarkRandomSetRegionWithGetRegionSizeByRange(b *testing.B) { regions.UpdateSubTree(region, origin, overlaps, rangeChanged) items = append(items, region) } + // Scan a partial range: an empty range is answered from regionTree.totalSize + // in O(1) and never reaches the range scan we want to measure against. + scanStart, scanEnd := benchScanRange(len(items)) b.ResetTimer() - go func() { - for { - regions.GetRegionSizeByRange([]byte(""), []byte("")) - time.Sleep(time.Millisecond) - } - }() + stopBackgroundScan(b, func() { + regions.GetRegionSizeByRange(scanStart, scanEnd) + time.Sleep(time.Millisecond) + }) for i := range b.N { item := items[i%len(items)] - item.approximateKeys = int64(200000) - origin, overlaps, rangeChanged := regions.SetRegion(item) - regions.UpdateSubTree(item, origin, overlaps, rangeChanged) + // Publish a new version instead of mutating the region already stored in + // the tree: a stored *RegionInfo is immutable, see regionItem. + n := item.Clone(SetApproximateKeys(int64(200000))) + origin, overlaps, rangeChanged := regions.SetRegion(n) + regions.UpdateSubTree(n, origin, overlaps, rangeChanged) } } @@ -866,13 +928,14 @@ func BenchmarkRandomSetRegionWithGetRegionSizeByRangeParallel(b *testing.B) { regions.UpdateSubTree(region, origin, overlaps, rangeChanged) items = append(items, region) } + // Scan a partial range: an empty range is answered from regionTree.totalSize + // in O(1) and never reaches the range scan we want to measure against. + scanStart, scanEnd := benchScanRange(len(items)) b.ResetTimer() - go func() { - for { - regions.GetRegionSizeByRange([]byte(""), []byte("")) - time.Sleep(time.Millisecond) - } - }() + stopBackgroundScan(b, func() { + regions.GetRegionSizeByRange(scanStart, scanEnd) + time.Sleep(time.Millisecond) + }) b.RunParallel( func(pb *testing.PB) { @@ -880,12 +943,239 @@ func BenchmarkRandomSetRegionWithGetRegionSizeByRangeParallel(b *testing.B) { item := items[mrand.IntN(len(items))] n := item.Clone(SetApproximateSize(20)) origin, overlaps, rangeChanged := regions.SetRegion(n) - regions.UpdateSubTree(item, origin, overlaps, rangeChanged) + regions.UpdateSubTree(n, origin, overlaps, rangeChanged) } }, ) } +// benchScanRange returns a key range covering the first half of a benchmark +// region set built with sequential `fmt.Sprintf("%20d", i)` keys. +func benchScanRange(count int) (startKey, endKey []byte) { + return []byte(fmt.Sprintf("%20d", 0)), []byte(fmt.Sprintf("%20d", count/2)) +} + +// stopBackgroundScan runs f in a loop on a background goroutine and stops it when +// the benchmark ends. Benchmarks in one binary share a process, so a background +// loop that outlives its benchmark perturbs every later one. +func stopBackgroundScan(b *testing.B, f func()) { + done := make(chan struct{}) + stopped := make(chan struct{}) + go func() { + defer close(stopped) + for { + select { + case <-done: + return + default: + } + f() + } + }() + b.Cleanup(func() { + close(done) + <-stopped + }) +} + +// buildBenchRegions returns a RegionsInfo holding count adjacent regions with +// sequential keys, together with the regions themselves. +func buildBenchRegions(count int) (*RegionsInfo, []*RegionInfo) { + regions := NewRegionsInfo() + items := make([]*RegionInfo, 0, count) + for i := range count { + peer := &metapb.Peer{StoreId: 1, Id: uint64(i + 1)} + region := NewRegionInfo(&metapb.Region{ + Id: uint64(i + 1), + Peers: []*metapb.Peer{peer}, + StartKey: []byte(fmt.Sprintf("%20d", i)), + EndKey: []byte(fmt.Sprintf("%20d", i+1)), + RegionEpoch: &metapb.RegionEpoch{ConfVer: 1, Version: 1}, + }, peer, SetApproximateSize(10)) + origin, overlaps, rangeChanged := regions.SetRegion(region) + regions.UpdateSubTree(region, origin, overlaps, rangeChanged) + items = append(items, region) + } + return regions, items +} + +// BenchmarkSetRegionSameRange measures the dominant region-heartbeat write path: +// a region is re-reported with an unchanged key range and fresher statistics. +// +// Unlike BenchmarkRandomSetRegion, this alternates between two pre-built +// immutable versions of every region, so origin != region on every iteration and +// no already-published *RegionInfo is mutated. That matters twice over: it is what +// production does, and mutating a stored region in place would break the +// immutability that regionItem relies on. +func BenchmarkSetRegionSameRange(b *testing.B) { + regions, items := buildBenchRegions(1000000) + versions := make([][2]*RegionInfo, len(items)) + for i, item := range items { + versions[i] = [2]*RegionInfo{ + item.Clone(SetApproximateSize(11), SetWrittenBytes(1000)), + item.Clone(SetApproximateSize(12), SetWrittenBytes(2000)), + } + } + b.ReportAllocs() + b.ResetTimer() + for i := range b.N { + idx := i % len(versions) + region := versions[idx][(i/len(versions))%2] + origin, overlaps, rangeChanged := regions.SetRegion(region) + regions.UpdateSubTree(region, origin, overlaps, rangeChanged) + } +} + +// liveBenchSnapshot keeps a snapshot reachable for the duration of a benchmark. +// It is a package-level variable so the compiler cannot decide the snapshot is +// dead and let it be collected. +var liveBenchSnapshot *rootRangeSnapshot + +// BenchmarkSetRegionSameRangeWithSnapshot is BenchmarkSetRegionSameRange with a +// snapshot outstanding. +// +// The point is that it should cost the same. A same-range update never calls +// ReplaceOrInsert or Delete, so it never reaches btree's copy-on-write node copy, +// and an outstanding snapshot cannot make the region-heartbeat path pay anything. +func BenchmarkSetRegionSameRangeWithSnapshot(b *testing.B) { + regions, items := buildBenchRegions(1000000) + versions := make([][2]*RegionInfo, len(items)) + for i, item := range items { + versions[i] = [2]*RegionInfo{ + item.Clone(SetApproximateSize(11), SetWrittenBytes(1000)), + item.Clone(SetApproximateSize(12), SetWrittenBytes(2000)), + } + } + liveBenchSnapshot = regions.snapshotRootTree() + b.Cleanup(func() { liveBenchSnapshot = nil }) + b.ReportAllocs() + b.ResetTimer() + for i := range b.N { + idx := i % len(versions) + region := versions[idx][(i/len(versions))%2] + origin, overlaps, rangeChanged := regions.SetRegion(region) + regions.UpdateSubTree(region, origin, overlaps, rangeChanged) + } +} + +// BenchmarkSetRegionRangeChangedWithSnapshot is BenchmarkSetRegionRangeChanged +// with a snapshot re-taken periodically. +// +// This is the path that does pay for copy-on-write. Every clone makes the whole +// live tree foreign to its own write context, so the next write to each root-to- +// leaf path has to copy the nodes on it. Re-cloning often keeps the tree in that +// state instead of letting it converge back. +func BenchmarkSetRegionRangeChangedWithSnapshot(b *testing.B) { + for _, clonePeriod := range []int{100000, 1000} { + b.Run(fmt.Sprintf("clone_every=%d", clonePeriod), func(b *testing.B) { + regions, items := buildBenchRegions(1000000) + shapes := benchRegionShapes(items) + liveBenchSnapshot = regions.snapshotRootTree() + b.Cleanup(func() { liveBenchSnapshot = nil }) + b.ReportAllocs() + b.ResetTimer() + for i := range b.N { + if i%clonePeriod == 0 { + liveBenchSnapshot = regions.snapshotRootTree() + } + idx := i % len(shapes) + region := shapes[idx][(i/len(shapes)+1)%2] + origin, overlaps, rangeChanged := regions.SetRegion(region) + regions.UpdateSubTree(region, origin, overlaps, rangeChanged) + } + }) + } +} + +// BenchmarkSnapshotRootTree checks that taking a snapshot is O(1): the cost must +// not grow with the number of regions. That is what justifies taking the write +// lock to do it. +func BenchmarkSnapshotRootTree(b *testing.B) { + for _, size := range []int{100000, 1000000, 4000000} { + b.Run(fmt.Sprintf("regions=%d", size), func(b *testing.B) { + regions, _ := buildBenchRegions(size) + b.ReportAllocs() + b.ResetTimer() + for range b.N { + liveBenchSnapshot = regions.snapshotRootTree() + } + b.StopTimer() + liveBenchSnapshot = nil + }) + } +} + +// benchRegionShapes returns, for every region, two versions that differ in end +// key, so alternating between them always changes the region's range. +func benchRegionShapes(items []*RegionInfo) [][2]*RegionInfo { + shapes := make([][2]*RegionInfo, len(items)) + for i, item := range items { + // "x" sorts after the start key and before the original end key. + shrunk := item.Clone(WithEndKey([]byte(fmt.Sprintf("%20d", i) + "x"))) + shapes[i] = [2]*RegionInfo{item, shrunk} + } + return shapes +} + +// BenchmarkSetRegionRangeChanged measures the structural write path: a region is +// re-reported with a different end key, so its tree item has to be removed and +// re-inserted. This is the only heartbeat shape that touches the btree structure. +func BenchmarkSetRegionRangeChanged(b *testing.B) { + regions, items := buildBenchRegions(1000000) + shapes := benchRegionShapes(items) + b.ReportAllocs() + b.ResetTimer() + for i := range b.N { + idx := i % len(shapes) + region := shapes[idx][(i/len(shapes)+1)%2] + origin, overlaps, rangeChanged := regions.SetRegion(region) + regions.UpdateSubTree(region, origin, overlaps, rangeChanged) + } +} + +// BenchmarkBatchScanRegionsHighConcurrency measures many short scans at high +// concurrency, mixed with heartbeat writers and point readers. A short scan is +// the realistic client shape, and it is the case where per-call snapshot overhead +// would show up; a single wide scan would amortise it away and hide the cost. +func BenchmarkBatchScanRegionsHighConcurrency(b *testing.B) { + const count = 1000000 + regions, items := buildBenchRegions(count) + versions := make([]*RegionInfo, len(items)) + for i, item := range items { + versions[i] = item.Clone(SetApproximateSize(11)) + } + + // Background heartbeat writers and point readers, so the scan contends the + // same lock it does in production. + stopBackgroundScan(b, func() { + i := mrand.IntN(len(versions)) + region := versions[i] + origin, overlaps, rangeChanged := regions.SetRegion(region) + regions.UpdateSubTree(region, origin, overlaps, rangeChanged) + }) + stopBackgroundScan(b, func() { + regions.GetRegion(uint64(mrand.IntN(count) + 1)) + }) + + b.ReportAllocs() + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + // Roughly 8 regions per scan. + start := mrand.IntN(count - 8) + krs := keyutil.NewKeyRanges([]keyutil.KeyRange{ + keyutil.NewKeyRange( + fmt.Sprintf("%20d", start), + fmt.Sprintf("%20d", start+8), + ), + }) + if _, err := regions.BatchScanRegions(krs); err != nil { + b.Fatal(err) + } + } + }) +} + const ( peerNum = 3 storeNum = 10 @@ -1085,6 +1375,260 @@ func TestUpdateRegionEquivalence(t *testing.T) { checksEquivalence() } +// TestRootRangeSnapshotEquivalence checks that a snapshot of a quiescent tree +// agrees with the live tree, and that taking one changes nothing observable -- +// in particular that it takes no reference on any region, which is both what +// keeps it O(1) and what keeps RegionInfo.ref meaning "present in a tree". +func TestRootRangeSnapshotEquivalence(t *testing.T) { + re := require.New(t) + regions := NewRegionsInfo() + items := generateTestRegions(1000, 5) + for _, item := range items { + origin, overlaps, rangeChanged := regions.SetRegion(item) + regions.UpdateSubTree(item, origin, overlaps, rangeChanged) + } + + refsBefore := make(map[uint64]int32, len(items)) + for _, item := range items { + refsBefore[item.GetID()] = item.GetRef() + } + lenBefore := regions.tree.length() + + snap := regions.snapshotRootTree() + + // Same regions, same order, same count. + re.Equal(lenBefore, snap.length()) + re.Equal(lenBefore, regions.tree.length()) + var fromSnap []uint64 + snap.scanRange([]byte(""), func(r *RegionInfo) bool { + fromSnap = append(fromSnap, r.GetID()) + return true + }) + var fromTree []uint64 + regions.tree.scanRange([]byte(""), func(r *RegionInfo) bool { + fromTree = append(fromTree, r.GetID()) + return true + }) + re.Equal(fromTree, fromSnap) + + // GetRegionSizeByRange, which now walks a snapshot, agrees with a brute-force + // sum over ScanRegions for arbitrary ranges, including keys that fall inside a + // region rather than on a boundary. + for range 200 { + i, j := mrand.IntN(len(items)*10), mrand.IntN(len(items)*10) + if i > j { + i, j = j, i + } + startKey := []byte(fmt.Sprintf("%20d", i)) + endKey := []byte(fmt.Sprintf("%20d", j)) + var want int64 + for _, r := range regions.ScanRegions(startKey, endKey, -1) { + want += r.GetApproximateSize() + } + re.Equal(want, regions.GetRegionSizeByRange(startKey, endKey), + "range [%s, %s)", startKey, endKey) + } + + // Taking and dropping a snapshot must not touch any reference count. + for _, item := range items { + re.Equal(refsBefore[item.GetID()], item.GetRef(), "region %d", item.GetID()) + } +} + +// TestGetRegionSizeByRangeConcurrent exercises the lock-free snapshot walk in +// GetRegionSizeByRange against concurrent writers of every shape. Run with -race. +// +// This guards the new code rather than reproducing an old bug: before the walk +// moved onto a snapshot, every scan held r.t, so there was no unsynchronised +// reader to race with. +func TestGetRegionSizeByRangeConcurrent(t *testing.T) { + re := require.New(t) + regions := NewRegionsInfo() + const count = 500 + items := generateTestRegions(count, 5) + for _, item := range items { + origin, overlaps, rangeChanged := regions.SetRegion(item) + regions.UpdateSubTree(item, origin, overlaps, rangeChanged) + } + total := int64(count * 10) + + var readers, writers sync.WaitGroup + done := make(chan struct{}) + + for range 4 { + readers.Add(1) + go func() { + defer readers.Done() + for range 300 { + startKey := []byte(fmt.Sprintf("%20d", 0)) + endKey := []byte(fmt.Sprintf("%20d", count*10)) + size := regions.GetRegionSizeByRange(startKey, endKey) + if size < 0 || size > total*2 { + t.Errorf("implausible size %d", size) + return + } + } + }() + } + + // Writers: same-range flow updates, splits, and merges back. + for w := range 3 { + writers.Add(1) + go func() { + defer writers.Done() + for i := 0; ; i++ { + select { + case <-done: + return + default: + } + idx := (i*7 + w) % count + base := items[idx] + switch i % 3 { + case 0: // same range, fresher stats + n := base.Clone(SetApproximateSize(11), SetWrittenBytes(uint64(i))) + origin, overlaps, rangeChanged := regions.SetRegion(n) + regions.UpdateSubTree(n, origin, overlaps, rangeChanged) + case 1: // shrink: range change + n := base.Clone( + WithEndKey([]byte(fmt.Sprintf("%20d", idx*10)+"x")), + WithIncVersion()) + origin, overlaps, rangeChanged := regions.SetRegion(n) + regions.UpdateSubTree(n, origin, overlaps, rangeChanged) + default: // restore: range change back + n := base.Clone(WithIncVersion()) + origin, overlaps, rangeChanged := regions.SetRegion(n) + regions.UpdateSubTree(n, origin, overlaps, rangeChanged) + } + } + }() + } + + readers.Wait() + close(done) + writers.Wait() + + // Restore every region so the subtree bookkeeping can be cross-checked. + for _, item := range items { + n := item.Clone(WithIncVersion()) + origin, overlaps, rangeChanged := regions.SetRegion(n) + regions.UpdateSubTree(n, origin, overlaps, rangeChanged) + } + checkRegions(re, regions) +} + +// TestRootRangeSnapshotAdjacencyUnderMerges checks the property that makes a +// snapshot a valid replacement for a scan under one read lock: the key ranges it +// reports are exactly those that existed when it was taken, so an adjacency check +// over a snapshot cannot see a hole that was not already there. +// +// Merges are used rather than splits on purpose. A split reaches PD as two +// independent updates -- the shrunk parent, then the new child -- so between them +// the key space really does have a hole, and a scanner would be right to report +// one even while holding the lock for its whole scan. +func TestRootRangeSnapshotAdjacencyUnderMerges(t *testing.T) { + re := require.New(t) + regions := NewRegionsInfo() + const count = 200 + items := generateTestRegions(count, 5) + for _, item := range items { + origin, overlaps, rangeChanged := regions.SetRegion(item) + regions.UpdateSubTree(item, origin, overlaps, rangeChanged) + } + + // scanAdjacent walks the snapshot and fails if two successive regions do not + // meet, or if the covered range stops short of endKey. + scanAdjacent := func(snap *rootRangeSnapshot, startKey, endKey []byte) error { + var lastEnd []byte + var err error + first := true + snap.scanRange(startKey, func(r *RegionInfo) bool { + if len(endKey) > 0 && bytes.Compare(r.GetStartKey(), endKey) >= 0 { + return false + } + if !first && !bytes.Equal(lastEnd, r.GetStartKey()) { + err = errs.ErrRegionNotAdjacent.FastGen( + "hole between %s and %s", lastEnd, r.GetStartKey()) + return false + } + first = false + lastEnd = r.GetEndKey() + return true + }) + if err == nil && bytes.Compare(lastEnd, endKey) < 0 { + err = errs.ErrRegionNotAdjacent.FastGen("short of %s, reached %s", endKey, lastEnd) + } + return err + } + + var readers, writers sync.WaitGroup + done := make(chan struct{}) + // A merge is one tree update and never leaves a hole. Undoing it is a split, + // which is inherently two updates -- shrink the parent, then insert the child + // -- and between them the key space really is missing a range. That is true of + // PD itself, where the two halves arrive as separate heartbeats, so a scanner + // is right to report a hole there whatever locking it uses. + // + // pairMu makes each such pair atomic with respect to taking a snapshot, which + // is the only way to ask the question this test is about: given a snapshot of a + // contiguous key space, does it stay contiguous while the tree changes + // underneath it? + var pairMu syncutil.Mutex + writers.Add(1) + go func() { + defer writers.Done() + for i := 1; ; i++ { + select { + case <-done: + return + default: + } + idx := 1 + i%(count-1) + // Merge idx into idx-1. One update, no hole. + merged := items[idx-1].Clone( + WithEndKey(items[idx].GetEndKey()), WithIncVersion()) + origin, overlaps, rangeChanged := regions.SetRegion(merged) + regions.UpdateSubTree(merged, origin, overlaps, rangeChanged) + + // Split them apart again. Two updates, so hold pairMu across both. + pairMu.Lock() + restored := items[idx-1].Clone(WithIncVersion()) + origin, overlaps, rangeChanged = regions.SetRegion(restored) + regions.UpdateSubTree(restored, origin, overlaps, rangeChanged) + back := items[idx].Clone(WithIncVersion()) + origin, overlaps, rangeChanged = regions.SetRegion(back) + regions.UpdateSubTree(back, origin, overlaps, rangeChanged) + pairMu.Unlock() + } + }() + + for range 4 { + readers.Add(1) + go func() { + defer readers.Done() + for range 300 { + // Only taking the snapshot is serialised against a split pair. The + // scan itself runs with nothing held, concurrently with the writer. + pairMu.Lock() + snap := regions.snapshotRootTree() + pairMu.Unlock() + + startKey := []byte(fmt.Sprintf("%20d", 0)) + endKey := []byte(fmt.Sprintf("%20d", count*10)) + if err := scanAdjacent(snap, startKey, endKey); err != nil { + t.Errorf("snapshot reported a hole: %v", err) + return + } + } + }() + } + + readers.Wait() + close(done) + writers.Wait() + re.NotNil(regions) +} + func generateTestRegions(count int, storeNum int) []*RegionInfo { var items []*RegionInfo for i := range count { diff --git a/pkg/core/region_tree.go b/pkg/core/region_tree.go index e7bebf09fd..8b683ca5cf 100644 --- a/pkg/core/region_tree.go +++ b/pkg/core/region_tree.go @@ -17,6 +17,7 @@ package core import ( "bytes" "math/rand/v2" + "sync/atomic" "go.uber.org/zap" @@ -29,29 +30,73 @@ import ( "github.com/tikv/pd/pkg/utils/logutil" ) +// regionItem is the value a region tree stores. It holds a *RegionInfo and +// supplies the btree ordering for it. +// +// # Ordering immutability +// +// Once a regionItem has been inserted into a tree, its key range must never +// change. The RegionInfo it holds may only be replaced by one whose start and end +// keys are identical, see RegionInfo.rangeEqualsTo. A range change must allocate a +// fresh regionItem instead; setRegionLocked does that. +// +// This is what lets an O(1) copy-on-write btree Clone serve as a usable read-only +// snapshot: a clone shares regionItem pointers with the live tree, so changing a +// shared item's range would leave it at the wrong position in the clone and +// silently break the clone's ordering. See rootRangeSnapshot. +// +// # Publish-once +// +// The RegionInfo is held atomically so that snapshot readers can walk a cloned +// tree without holding the tree's lock. The atomic protects the slot only. It does +// not protect the RegionInfo, its metapb.Region, or its key slices, so a +// *RegionInfo that has been stored here must never be mutated afterwards. The +// exceptions are the fields that carry their own atomics: RegionInfo.ref and +// RegionInfo.reportBuckets. type regionItem struct { - *RegionInfo + region atomic.Pointer[RegionInfo] +} + +// newRegionItem returns a regionItem holding the given region. +func newRegionItem(region *RegionInfo) *regionItem { + item := ®ionItem{} + item.region.Store(region) + return item +} + +// getRegion returns the RegionInfo held by the item. +func (r *regionItem) getRegion() *RegionInfo { + return r.region.Load() +} + +// setRegion replaces the RegionInfo held by the item. +// +// The caller must hold the write lock of every tree the item belongs to, and the +// new region must cover the same key range as the current one. See the ordering +// immutability note on regionItem. +func (r *regionItem) setRegion(region *RegionInfo) { + r.region.Store(region) } // GetStartKey returns the start key of the region. func (r *regionItem) GetStartKey() []byte { - return r.meta.StartKey + return r.getRegion().meta.StartKey } // GetID returns the ID of the region. func (r *regionItem) GetID() uint64 { - return r.meta.GetId() + return r.getRegion().meta.GetId() } // GetEndKey returns the end key of the region. func (r *regionItem) GetEndKey() []byte { - return r.meta.EndKey + return r.getRegion().meta.EndKey } // Less returns true if the region start key is less than the other. func (r *regionItem) Less(other *regionItem) bool { - left := r.meta.StartKey - right := other.meta.StartKey + left := r.getRegion().meta.StartKey + right := other.getRegion().meta.StartKey return bytes.Compare(left, right) < 0 } @@ -94,8 +139,8 @@ func newRegionTreeWithCountRef() *regionTree { // GetCountByRange returns the number of regions in the range [startKey, endKey). func (t *regionTree) GetCountByRange(startKey, endKey []byte) int { - start := ®ionItem{&RegionInfo{meta: &metapb.Region{StartKey: startKey}}} - end := ®ionItem{&RegionInfo{meta: &metapb.Region{StartKey: endKey}}} + start := newRegionItem(&RegionInfo{meta: &metapb.Region{StartKey: startKey}}) + end := newRegionItem(&RegionInfo{meta: &metapb.Region{StartKey: endKey}}) // it returns 0 if startKey is nil. item, startIndex := t.tree.GetWithIndex(start) // if item is nil, it means that the startKey is not found in the tree, we need to check the previous item, avoid @@ -151,7 +196,7 @@ func (t *regionTree) overlaps(item *regionItem) []*RegionInfo { if len(endKey) > 0 && bytes.Compare(endKey, i.GetStartKey()) <= 0 { return false } - overlaps = append(overlaps, i.RegionInfo) + overlaps = append(overlaps, i.getRegion()) return true }) return overlaps @@ -172,7 +217,7 @@ func (t *regionTree) updateRef(origin, region *RegionInfo) { // It finds and deletes all the overlapped regions first, and then // insert the region. func (t *regionTree) update(item *regionItem, withOverlaps bool, overlaps ...*RegionInfo) []*RegionInfo { - region := item.RegionInfo + region := item.getRegion() t.totalSize += region.approximateSize regionWriteBytesRate, regionWriteKeysRate := region.GetWriteRate() t.totalWriteBytesRate += regionWriteBytesRate @@ -186,11 +231,11 @@ func (t *regionTree) update(item *regionItem, withOverlaps bool, overlaps ...*Re } for _, old := range overlaps { - t.tree.Delete(®ionItem{RegionInfo: old}) + t.tree.Delete(newRegionItem(old)) } t.tree.ReplaceOrInsert(item) if t.countRef { - item.IncRef() + item.getRegion().IncRef() } result := make([]*RegionInfo, len(overlaps)) for i, overlap := range overlaps { @@ -245,18 +290,18 @@ func (t *regionTree) remove(region *RegionInfo) { if t.length() == 0 { return } - item := ®ionItem{RegionInfo: region} + item := newRegionItem(region) result := t.find(item) if result == nil || result.GetID() != region.GetID() { return } - t.totalSize -= result.GetApproximateSize() - regionWriteBytesRate, regionWriteKeysRate := result.GetWriteRate() + t.totalSize -= result.getRegion().GetApproximateSize() + regionWriteBytesRate, regionWriteKeysRate := result.getRegion().GetWriteRate() t.totalWriteBytesRate -= regionWriteBytesRate t.totalWriteKeysRate -= regionWriteKeysRate if t.countRef { - result.DecRef() + result.getRegion().DecRef() } if !region.LoadedFromStorage() { t.notFromStorageRegionsCnt-- @@ -267,28 +312,28 @@ func (t *regionTree) remove(region *RegionInfo) { // search returns a region that contains the key. func (t *regionTree) search(regionKey []byte) *RegionInfo { region := &RegionInfo{meta: &metapb.Region{StartKey: regionKey}} - result := t.find(®ionItem{RegionInfo: region}) + result := t.find(newRegionItem(region)) if result == nil { return nil } - return result.RegionInfo + return result.getRegion() } // searchPrev returns the previous region of the region where the regionKey is located. func (t *regionTree) searchPrev(regionKey []byte) *RegionInfo { curRegion := &RegionInfo{meta: &metapb.Region{StartKey: regionKey}} - curRegionItem := t.find(®ionItem{RegionInfo: curRegion}) + curRegionItem := t.find(newRegionItem(curRegion)) if curRegionItem == nil { return nil } - prevRegionItem, _ := t.getAdjacentRegions(curRegionItem.RegionInfo) + prevRegionItem, _ := t.getAdjacentRegions(curRegionItem.getRegion()) if prevRegionItem == nil { return nil } if !bytes.Equal(prevRegionItem.GetEndKey(), curRegionItem.GetStartKey()) { return nil } - return prevRegionItem.RegionInfo + return prevRegionItem.getRegion() } // searchByKeys searches the regions by keys and return a slice of `*RegionInfo` whose order is the same as the input keys. @@ -311,40 +356,120 @@ func (t *regionTree) searchByPrevKeys(prevKeys [][]byte) []*RegionInfo { return regions } -// find returns the range item contains the start key. -func (t *regionTree) find(item *regionItem) *regionItem { +// rootRangeSnapshot is a read-only view of a region tree, produced by an O(1) +// copy-on-write btree Clone. A caller can scan an arbitrarily large key range +// from a snapshot without holding the tree's lock for the duration of the scan. +// +// # What a snapshot freezes +// +// The tree's shape: the set of items, each item's key range, and therefore the +// key-space coverage, the iteration order and the length. Regions created, split, +// merged or removed after the snapshot was taken are invisible to it, and regions +// removed after it was taken stay visible to it. +// +// This is what a chunked scan cannot offer. A scan that releases the lock and +// re-enters the tree at the last end key can see the key space change underneath +// it: a split at a chunk boundary yields a region counted twice, and a merge that +// swallows the boundary region leaves a range no chunk covers. A snapshot has no +// second scan to be inconsistent with. +// +// # What a snapshot does not freeze +// +// The values. The RegionInfo behind an item may be replaced concurrently by the +// heartbeat path, but only by one covering the same key range, see regionItem. So +// a scan observes exactly one region per key, and for each of them either the +// value that was live when the snapshot was taken or a later value for the same +// range, with fresher epoch, leader, peers or statistics. +// +// A scan of a live tree under one read lock is stronger than this: there, every +// region comes from a single instant. Callers that need agreement between two +// different regions' values, rather than agreement about which key ranges exist, +// must not use a snapshot. +// +// # References and lifetime +// +// A snapshot deliberately takes no reference on the regions it holds, see +// RegionInfo.IncRef. Doing so would make it O(n), and the reference count carries +// the functional "already present in the subtree" signal that RaftCluster relies +// on. A region reached through a snapshot may therefore already be gone from the +// live tree. +// +// A snapshot pins the btree nodes that existed when it was taken. Those nodes can +// no longer be recycled through the shared free list, and the regions they point +// at stay reachable. Keep a snapshot function-scoped and short-lived; never store +// one in a long-lived struct. +type rootRangeSnapshot struct { + tree *btree.BTreeG[*regionItem] +} + +// snapshot returns a read-only view of the tree. +// +// The caller must hold the tree's *write* lock. btree Clone rewrites the source +// tree's copy-on-write context, so it must not run concurrently with a write or +// with another Clone; see btree.BTreeG.Clone. +func (t *regionTree) snapshot() *rootRangeSnapshot { + return &rootRangeSnapshot{tree: t.tree.Clone()} +} + +// scanRange scans from the first region containing or behind the start key until +// f returns false. It takes no lock. +func (s *rootRangeSnapshot) scanRange(startKey []byte, f func(*RegionInfo) bool) { + scanTreeRange(s.tree, startKey, f) +} + +// length returns the number of regions the snapshot holds. +func (s *rootRangeSnapshot) length() int { + return s.tree.Len() +} + +// findItem returns the item in tree whose key range contains the start key of +// the given item. +// +// It only reads tree, so it is also safe to call on a snapshot. See +// rootRangeSnapshot. +func findItem(tree *btree.BTreeG[*regionItem], item *regionItem) *regionItem { var result *regionItem - t.tree.DescendLessOrEqual(item, func(i *regionItem) bool { + tree.DescendLessOrEqual(item, func(i *regionItem) bool { result = i return false }) - if result == nil || !result.contain(item.GetStartKey()) { + if result == nil || !result.getRegion().contain(item.GetStartKey()) { return nil } return result } -// scanRage scans from the first region containing or behind the start key -// until f return false -func (t *regionTree) scanRange(startKey []byte, f func(*RegionInfo) bool) { +// scanTreeRange scans tree from the first region containing or behind the start +// key until f returns false. +// +// It only reads tree, so it is also safe to call on a snapshot. See +// rootRangeSnapshot. +func scanTreeRange(tree *btree.BTreeG[*regionItem], startKey []byte, f func(*RegionInfo) bool) { region := &RegionInfo{meta: &metapb.Region{StartKey: startKey}} + start := newRegionItem(region) // find if there is a region with key range [s, d), s <= startKey < d - fn := func(item *regionItem) bool { - r := item - return f(r.RegionInfo) - } - start := ®ionItem{RegionInfo: region} - startItem := t.find(start) + startItem := findItem(tree, start) if startItem == nil { startItem = start } - t.tree.AscendGreaterOrEqual(startItem, func(item *regionItem) bool { - return fn(item) + tree.AscendGreaterOrEqual(startItem, func(item *regionItem) bool { + return f(item.getRegion()) }) } +// find returns the range item contains the start key. +func (t *regionTree) find(item *regionItem) *regionItem { + return findItem(t.tree, item) +} + +// scanRage scans from the first region containing or behind the start key +// until f return false +func (t *regionTree) scanRange(startKey []byte, f func(*RegionInfo) bool) { + scanTreeRange(t.tree, startKey, f) +} + func (t *regionTree) scanRanges() []*RegionInfo { if t.length() == 0 { return nil @@ -358,7 +483,7 @@ func (t *regionTree) scanRanges() []*RegionInfo { } func (t *regionTree) getAdjacentRegions(region *RegionInfo) (prev, next *regionItem) { - item := ®ionItem{RegionInfo: &RegionInfo{meta: &metapb.Region{StartKey: region.GetStartKey()}}} + item := newRegionItem(&RegionInfo{meta: &metapb.Region{StartKey: region.GetStartKey()}}) return t.getAdjacentItem(item) } @@ -402,10 +527,15 @@ func (t *regionTree) RandomRegions(n int, ranges []keyutil.KeyRange) []*RegionIn startIndex, endIndex = 0, treeLen randIndex int startItem *regionItem - pivotItem = ®ionItem{&RegionInfo{meta: &metapb.Region{}}} - region *RegionInfo - regions = make([]*RegionInfo, 0, n) - curLen = len(regions) + // pivotRegion is the region behind pivotItem. pivotItem is a scratch item + // that is only ever used to look up a key, and is never inserted into the + // tree, so its start key may be rewritten in place through pivotRegion. + // That does not violate the ordering immutability of tree-resident items. + pivotRegion = &RegionInfo{meta: &metapb.Region{}} + pivotItem = newRegionItem(pivotRegion) + region *RegionInfo + regions = make([]*RegionInfo, 0, n) + curLen = len(regions) // setStartEndIndices is a helper function to set `startIndex` and `endIndex` // according to the `startKey` and `endKey` and check if the range is invalid // to skip the iteration. @@ -416,10 +546,10 @@ func (t *regionTree) RandomRegions(n int, ranges []keyutil.KeyRange) []*RegionIn startIndex, endIndex = 0, treeLen return false } - pivotItem.meta.StartKey = startKey + pivotRegion.meta.StartKey = startKey startItem, startIndex = t.tree.GetWithIndex(pivotItem) if endKeyLen > 0 { - pivotItem.meta.StartKey = endKey + pivotRegion.meta.StartKey = endKey _, endIndex = t.tree.GetWithIndex(pivotItem) } else { endIndex = treeLen @@ -427,7 +557,7 @@ func (t *regionTree) RandomRegions(n int, ranges []keyutil.KeyRange) []*RegionIn // Consider that the item in the tree may not be continuous, // we need to check if the previous item contains the key. if startIndex != 0 && startItem == nil { - region = t.tree.GetAt(startIndex - 1).RegionInfo + region = t.tree.GetAt(startIndex - 1).getRegion() if region.contain(startKey) { startIndex-- } @@ -455,7 +585,7 @@ func (t *regionTree) RandomRegions(n int, ranges []keyutil.KeyRange) []*RegionIn } for curLen < n { randIndex = rand.IntN(endIndex-startIndex) + startIndex - region = t.tree.GetAt(randIndex).RegionInfo + region = t.tree.GetAt(randIndex).getRegion() if region.isInvolved(startKey, endKey) { regions = append(regions, region) curLen++ @@ -478,7 +608,7 @@ func (t *regionTree) RandomRegions(n int, ranges []keyutil.KeyRange) []*RegionIn } randIndex = rand.IntN(endIndex-startIndex) + startIndex - region = t.tree.GetAt(randIndex).RegionInfo + region = t.tree.GetAt(randIndex).getRegion() if region.isInvolved(startKey, endKey) { regions = append(regions, region) curLen++ diff --git a/pkg/core/region_tree_test.go b/pkg/core/region_tree_test.go index cf3ddee2f1..b72a1ab811 100644 --- a/pkg/core/region_tree_test.go +++ b/pkg/core/region_tree_test.go @@ -15,6 +15,7 @@ package core import ( + "bytes" "fmt" "math/rand/v2" "testing" @@ -99,21 +100,21 @@ func TestRegionInfo(t *testing.T) { func TestRegionItem(t *testing.T) { re := require.New(t) - item := newRegionItem([]byte("b"), []byte{}) + item := newRegionItemFromRange([]byte("b"), []byte{}) - re.False(item.Less(newRegionItem([]byte("a"), []byte{}))) - re.False(item.Less(newRegionItem([]byte("b"), []byte{}))) - re.True(item.Less(newRegionItem([]byte("c"), []byte{}))) + re.False(item.Less(newRegionItemFromRange([]byte("a"), []byte{}))) + re.False(item.Less(newRegionItemFromRange([]byte("b"), []byte{}))) + re.True(item.Less(newRegionItemFromRange([]byte("c"), []byte{}))) - re.False(item.contain([]byte("a"))) - re.True(item.contain([]byte("b"))) - re.True(item.contain([]byte("c"))) + re.False(item.getRegion().contain([]byte("a"))) + re.True(item.getRegion().contain([]byte("b"))) + re.True(item.getRegion().contain([]byte("c"))) - item = newRegionItem([]byte("b"), []byte("d")) - re.False(item.contain([]byte("a"))) - re.True(item.contain([]byte("b"))) - re.True(item.contain([]byte("c"))) - re.False(item.contain([]byte("d"))) + item = newRegionItemFromRange([]byte("b"), []byte("d")) + re.False(item.getRegion().contain([]byte("a"))) + re.True(item.getRegion().contain([]byte("b"))) + re.True(item.getRegion().contain([]byte("c"))) + re.False(item.getRegion().contain([]byte("d"))) } func newRegionWithStat(start, end string, size, keys int64) *RegionInfo { @@ -161,9 +162,9 @@ func TestRegionTree(t *testing.T) { updateNewItem(tree, regionA) updateNewItem(tree, regionC) - re.Nil(tree.overlaps(newRegionItem([]byte("b"), []byte("c")))) - re.Equal(regionC, tree.overlaps(newRegionItem([]byte("c"), []byte("d")))[0]) - re.Equal(regionC, tree.overlaps(newRegionItem([]byte("a"), []byte("cc")))[1]) + re.Nil(tree.overlaps(newRegionItemFromRange([]byte("b"), []byte("c")))) + re.Equal(regionC, tree.overlaps(newRegionItemFromRange([]byte("c"), []byte("d")))[0]) + re.Equal(regionC, tree.overlaps(newRegionItemFromRange([]byte("a"), []byte("cc")))[1]) re.Nil(tree.search([]byte{})) re.Equal(regionA, tree.search([]byte("a"))) re.Nil(tree.search([]byte("b"))) @@ -191,28 +192,28 @@ func TestRegionTree(t *testing.T) { // check get adjacent regions prev, next := tree.getAdjacentRegions(regionA) re.Nil(prev) - re.Equal(regionB, next.RegionInfo) + re.Equal(regionB, next.getRegion()) prev, next = tree.getAdjacentRegions(regionB) - re.Equal(regionA, prev.RegionInfo) - re.Equal(regionD, next.RegionInfo) + re.Equal(regionA, prev.getRegion()) + re.Equal(regionD, next.getRegion()) prev, next = tree.getAdjacentRegions(regionC) - re.Equal(regionB, prev.RegionInfo) - re.Equal(regionD, next.RegionInfo) + re.Equal(regionB, prev.getRegion()) + re.Equal(regionD, next.getRegion()) prev, next = tree.getAdjacentRegions(regionD) - re.Equal(regionB, prev.RegionInfo) + re.Equal(regionB, prev.getRegion()) re.Nil(next) // region with the same range and different region id will not be delete. - region0 := newRegionItem([]byte{}, []byte("a")).RegionInfo + region0 := newRegionItemFromRange([]byte{}, []byte("a")).getRegion() updateNewItem(tree, region0) re.Equal(region0, tree.search([]byte{})) - anotherRegion0 := newRegionItem([]byte{}, []byte("a")).RegionInfo + anotherRegion0 := newRegionItemFromRange([]byte{}, []byte("a")).getRegion() anotherRegion0.meta.Id = 123 tree.remove(anotherRegion0) re.Equal(region0, tree.search([]byte{})) // overlaps with 0, A, B, C. - region0D := newRegionItem([]byte(""), []byte("d")).RegionInfo + region0D := newRegionItemFromRange([]byte(""), []byte("d")).getRegion() updateNewItem(tree, region0D) re.Equal(region0D, tree.search([]byte{})) re.Equal(region0D, tree.search([]byte("a"))) @@ -221,7 +222,7 @@ func TestRegionTree(t *testing.T) { re.Equal(regionD, tree.search([]byte("d"))) // overlaps with D. - regionE := newRegionItem([]byte("e"), []byte{}).RegionInfo + regionE := newRegionItemFromRange([]byte("e"), []byte{}).getRegion() updateNewItem(tree, regionE) re.Equal(region0D, tree.search([]byte{})) re.Equal(region0D, tree.search([]byte("a"))) @@ -243,10 +244,156 @@ func updateRegions(re *require.Assertions, tree *regionTree, regions []*RegionIn } } +// snapshotRanges returns the key ranges a snapshot reports, in iteration order. +func snapshotRanges(snap *rootRangeSnapshot) [][2]string { + var got [][2]string + snap.scanRange([]byte(""), func(r *RegionInfo) bool { + got = append(got, [2]string{string(r.GetStartKey()), string(r.GetEndKey())}) + return true + }) + return got +} + +// treeRanges returns the key ranges a live tree reports, in iteration order. +func treeRanges(regions *RegionsInfo) [][2]string { + var got [][2]string + regions.tree.scanRange([]byte(""), func(r *RegionInfo) bool { + got = append(got, [2]string{string(r.GetStartKey()), string(r.GetEndKey())}) + return true + }) + return got +} + +// TestRootRangeSnapshotIsolation checks the contract documented on +// rootRangeSnapshot: the shape is frozen, the values are not. +// +// It applies all three mutation classes after taking the snapshot -- a same-range +// update, a split, and a removal -- and asserts that the snapshot still reports +// the original key ranges in the original order while the live tree reports the +// new ones. +func TestRootRangeSnapshotIsolation(t *testing.T) { + re := require.New(t) + regions := NewRegionsInfo() + put := func(id uint64, start, end string, opts ...RegionCreateOption) { + region := NewTestRegionInfo(id, 1, []byte(start), []byte(end), opts...) + _, err := regions.AtomicCheckAndPutRegion(ContextTODO(), region) + re.NoError(err) + } + + put(1, "a", "b", SetApproximateSize(10)) + put(2, "b", "c", SetApproximateSize(10)) + put(3, "c", "d", SetApproximateSize(10)) + re.Equal(int64(30), regions.GetRegionSizeByRange([]byte("a"), []byte("d"))) + + snap := regions.snapshotRootTree() + original := [][2]string{{"a", "b"}, {"b", "c"}, {"c", "d"}} + re.Equal(original, snapshotRanges(snap)) + + // 1. Same range, new value. + put(2, "b", "c", SetApproximateSize(99), SetRegionVersion(2)) + // 2. Split [c, d) into [c, cc) and [cc, d). + put(3, "c", "cc", SetRegionVersion(2)) + put(4, "cc", "d", SetRegionVersion(2)) + // 3. Remove [a, b). + regions.RemoveRegionIfExist(1) + + // The shape is frozen: same items, same ranges, same order, same length. + re.Equal(3, snap.length()) + re.Equal(original, snapshotRanges(snap)) + + // The value is not frozen. The same-range update has completed and the item is + // shared with the live tree, so the snapshot now reports the new size. This is + // the documented latitude, not a bug. + var sizeOfB int64 + snap.scanRange([]byte("b"), func(r *RegionInfo) bool { + sizeOfB = r.GetApproximateSize() + return false + }) + re.Equal(int64(99), sizeOfB) + + // The live tree has moved on. + re.Equal([][2]string{{"b", "c"}, {"c", "cc"}, {"cc", "d"}}, treeRanges(regions)) +} + +// TestRegionItemRangeImmutability is the direct guard for the invariant +// documented on regionItem: an item that is in the root tree never changes its +// key range. +// +// It holds one snapshot for the whole test and, after every update, checks the +// ordering of both the live tree and that snapshot. An in-place range change +// anywhere on the root path shows up here as a snapshot whose start keys stop +// increasing, which is the state that makes a clone return wrong answers. +func TestRegionItemRangeImmutability(t *testing.T) { + re := require.New(t) + regions := NewRegionsInfo() + const count = 50 + for i := range count { + region := NewTestRegionInfo(uint64(i+1), 1, + []byte(fmt.Sprintf("%04d", i*10)), []byte(fmt.Sprintf("%04d", (i+1)*10))) + _, err := regions.AtomicCheckAndPutRegion(ContextTODO(), region) + re.NoError(err) + } + + snap := regions.snapshotRootTree() + snapLen := snap.length() + + checkOrdered := func(desc string, scan func(func(*RegionInfo) bool)) int { + var prev []byte + n := 0 + scan(func(r *RegionInfo) bool { + if n > 0 { + re.Negative(bytes.Compare(prev, r.GetStartKey()), + "%s: start keys not increasing at %q after %q", desc, r.GetStartKey(), prev) + } + prev = r.GetStartKey() + n++ + return true + }) + return n + } + + version := uint64(1) + for i := range 1000 { + id := uint64(i%count + 1) + origin := regions.GetRegion(id) + re.NotNil(origin) + version++ + var region *RegionInfo + if i%2 == 0 { + // Same range, new value. + region = origin.Clone(SetApproximateSize(int64(i)), SetRegionVersion(version)) + } else { + // Range change: toggle the end key between its original value and a + // shorter one. Appending to the start key keeps the new end key strictly + // inside the region, so this never overlaps a neighbour and never + // deletes one. + fullEnd := []byte(fmt.Sprintf("%04d", int(id)*10)) + shrunkEnd := append(append([]byte{}, origin.GetStartKey()...), '5') + endKey := shrunkEnd + if !bytes.Equal(origin.GetEndKey(), fullEnd) { + endKey = fullEnd + } + region = origin.Clone(WithEndKey(endKey), SetRegionVersion(version)) + } + _, err := regions.AtomicCheckAndPutRegion(ContextTODO(), region) + re.NoError(err) + + checkOrdered("live tree", func(f func(*RegionInfo) bool) { + regions.tree.scanRange([]byte(""), f) + }) + n := checkOrdered("snapshot", func(f func(*RegionInfo) bool) { + snap.scanRange([]byte(""), f) + }) + // The snapshot's shape is frozen, so its length cannot drift either. + re.Equal(snapLen, n) + re.Equal(snapLen, snap.length()) + } +} + func TestRegionTreeSplitAndMerge(t *testing.T) { re := require.New(t) tree := newRegionTree() - regions := []*RegionInfo{newRegionItem([]byte{}, []byte{}).RegionInfo} + regions := []*RegionInfo{newRegionItemFromRange([]byte{}, []byte{}).getRegion()} // Byte will underflow/overflow if n > 7. n := 7 @@ -408,7 +555,7 @@ func TestStoreRegionCount(t *testing.T) { } func updateNewItem(tree *regionTree, region *RegionInfo) { - item := ®ionItem{RegionInfo: region} + item := newRegionItem(region) tree.update(item, false) } @@ -431,8 +578,8 @@ func checkRandomRegion(re *require.Assertions, tree *regionTree, regions []*Regi re.Len(keys, len(regions)) } -func newRegionItem(start, end []byte) *regionItem { - return ®ionItem{RegionInfo: NewTestRegionInfo(1, 1, start, end)} +func newRegionItemFromRange(start, end []byte) *regionItem { + return newRegionItem(NewTestRegionInfo(1, 1, start, end)) } type mockRegionTreeData struct { @@ -519,7 +666,7 @@ func BenchmarkRegionTreeSequentialLookUpRegion(b *testing.B) { b.ResetTimer() for i := range b.N { index := i % MaxCount - data.tree.find(®ionItem{RegionInfo: data.items[index]}) + data.tree.find(newRegionItem(data.items[index])) } } @@ -528,7 +675,7 @@ func BenchmarkRegionTreeRandomLookUpRegion(b *testing.B) { b.ResetTimer() for i := range b.N { index := i % MaxCount - data.tree.find(®ionItem{RegionInfo: data.items[index]}) + data.tree.find(newRegionItem(data.items[index])) } }