From cbefbed93bdd1d8a1a1e8447e48b722da47cf7a1 Mon Sep 17 00:00:00 2001 From: Vlad Dmitriev Date: Wed, 12 Aug 2026 13:59:30 -0700 Subject: [PATCH 1/6] core: fix and extend region write/scan benchmarks Three defects made the existing benchmarks unable to measure the region heartbeat and range-scan paths: - BenchmarkRandomSetRegionWithGetRegionSizeByRange and its Parallel variant scanned an empty key range, which is answered from regionTree.totalSize in O(1) and never reaches the range scan. Both now scan a partial range. - Their background reader goroutines never stopped, so they outlived their benchmark and perturbed every later one in the same binary. - BenchmarkRandomSetRegion mutates an already published *RegionInfo and re-reports the same pointer, so origin == region and updateStat adds and subtracts the same values. It does not model a heartbeat. Add benchmarks for the three paths that matter: - BenchmarkSetRegionSameRange alternates two immutable same-range versions of every region, which is what a region heartbeat does. - BenchmarkSetRegionRangeChanged alternates two end keys, the only heartbeat shape that touches the btree structure. - BenchmarkBatchScanRegionsHighConcurrency runs many short scans concurrently with heartbeat writers and point readers, which is the realistic client shape; a single wide scan amortises away the per-call cost. Signed-off-by: Vlad Dmitriev --- pkg/core/region_test.go | 187 ++++++++++++++++++++++++++++++++++++---- 1 file changed, 171 insertions(+), 16 deletions(-) diff --git a/pkg/core/region_test.go b/pkg/core/region_test.go index 00b40bf062..0c8747d282 100644 --- a/pkg/core/region_test.go +++ b/pkg/core/region_test.go @@ -768,6 +768,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 @@ -836,18 +839,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 +872,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 +887,160 @@ 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) + } +} + +// 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 From 8fe893e7328e205a37747d7aee5e819a07effa27 Mon Sep 17 00:00:00 2001 From: Vlad Dmitriev Date: Wed, 12 Aug 2026 14:13:14 -0700 Subject: [PATCH 2/6] core: hold regionItem's RegionInfo in an atomic pointer regionItem embedded a *RegionInfo, and setRegionLocked replaces that pointer in place while other goroutines may be reading the item. Today every reader holds RegionsInfo.t, so the replacement is safe; it stops being safe as soon as anything reads a region tree without that lock. Replace the embedded pointer with an atomic.Pointer slot reached through getRegion/setRegion. atomic.Pointer is one word, so regionItem stays 8 bytes and there is no per-region memory cost. Two notes recorded on the type, because both are load-bearing and neither was written down before: - Ordering immutability: an item that is in a tree must never change its key range. This is what a later copy-on-write snapshot will depend on. - Publish-once: the atomic protects the slot, not the RegionInfo behind it, so a stored *RegionInfo must not be mutated afterwards. RegionInfo.ref and RegionInfo.reportBuckets are the exceptions, as both carry their own atomic. Dropping the embedded field also removes the shadowing of RegionInfo's GetID, GetStartKey and GetEndKey by the identically named methods on regionItem, which previously made it impossible to tell at a call site which of the two was being invoked. No behaviour change. Signed-off-by: Vlad Dmitriev --- pkg/core/region.go | 42 ++++++------- pkg/core/region_test.go | 10 +-- pkg/core/region_tree.go | 118 +++++++++++++++++++++++++---------- pkg/core/region_tree_test.go | 62 +++++++++--------- 4 files changed, 141 insertions(+), 91 deletions(-) diff --git a/pkg/core/region.go b/pkg/core/region.go index 40701a5137..f4b9c9e1bb 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,7 +1355,7 @@ 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. @@ -1372,17 +1372,17 @@ func (r *RegionsInfo) setRegionLocked(region *RegionInfo, withOverlaps bool, ol } r.tree.remove(origin) // Update the RegionInfo in the regionItem. - item.RegionInfo = region + item.setRegion(region) } 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 + 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 @@ -1449,7 +1449,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 +1581,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 +1884,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 +1998,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 +2008,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 } @@ -2306,10 +2306,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 0c8747d282..c39662b284 100644 --- a/pkg/core/region_test.go +++ b/pkg/core/region_test.go @@ -435,14 +435,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 +450,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) } diff --git a/pkg/core/region_tree.go b/pkg/core/region_tree.go index e7bebf09fd..8570622b21 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. @@ -319,7 +364,7 @@ func (t *regionTree) find(item *regionItem) *regionItem { return false }) - if result == nil || !result.contain(item.GetStartKey()) { + if result == nil || !result.getRegion().contain(item.GetStartKey()) { return nil } @@ -333,9 +378,9 @@ func (t *regionTree) scanRange(startKey []byte, f func(*RegionInfo) bool) { // 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) + return f(r.getRegion()) } - start := ®ionItem{RegionInfo: region} + start := newRegionItem(region) startItem := t.find(start) if startItem == nil { startItem = start @@ -358,7 +403,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 +447,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 +466,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 +477,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 +505,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 +528,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..1826b05d0a 100644 --- a/pkg/core/region_tree_test.go +++ b/pkg/core/region_tree_test.go @@ -99,21 +99,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 +161,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 +191,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 +221,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"))) @@ -246,7 +246,7 @@ func updateRegions(re *require.Assertions, tree *regionTree, regions []*RegionIn 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 +408,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 +431,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 +519,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 +528,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])) } } From da4ef918b036fa6d493272f45e62195d68a6aa05 Mon Sep 17 00:00:00 2001 From: Vlad Dmitriev Date: Wed, 12 Aug 2026 14:32:41 -0700 Subject: [PATCH 3/6] core: allocate a fresh regionItem when a region's range changes setRegionLocked reused the regionItem from r.regions even when the region's key range changed, rewriting the item's start and end keys in place. That is safe only while every reader holds RegionsInfo.t, because the item's keys are what orders it in the btree: a reader holding a reference to that item across the change would see it at one position while it reports another range. Allocate a fresh item on the range-changed path instead and repoint the map entry. The path already performs tree.remove plus overlaps, Delete and ReplaceOrInsert, so one 8-byte allocation is not measurable against it. The same-range path is untouched and keeps replacing the value in place, which is what the region heartbeat does. The reference count is unchanged. tree.remove still finds the old item, which still holds the origin, so origin loses its reference; tree.update then adds one for the region behind the fresh item. Two things that were previously relied on but not written down are now comments: that tree.remove has to run before the item is replaced, and why the new map assignment cannot be undone by the overlap cleanup that follows. TestSetRegionRangeChangedRepointsMap covers the map, the tree and the reference counts together. Without the repointing it fails with the pre-shrink region. Signed-off-by: Vlad Dmitriev --- pkg/core/region.go | 23 +++++++++++++++++++---- pkg/core/region_test.go | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 4 deletions(-) diff --git a/pkg/core/region.go b/pkg/core/region.go index f4b9c9e1bb..c3db85d6f0 100644 --- a/pkg/core/region.go +++ b/pkg/core/region.go @@ -1359,7 +1359,6 @@ func (r *RegionsInfo) setRegionLocked(region *RegionInfo, withOverlaps bool, ol 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,13 +1369,29 @@ 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.setRegion(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. + // 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 } diff --git a/pkg/core/region_test.go b/pkg/core/region_test.go index c39662b284..878341ece3 100644 --- a/pkg/core/region_test.go +++ b/pkg/core/region_test.go @@ -530,6 +530,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() From 372edea0b29bcd807866da91f87c2db2e1500567 Mon Sep 17 00:00:00 2001 From: Vlad Dmitriev Date: Wed, 12 Aug 2026 14:42:10 -0700 Subject: [PATCH 4/6] core: scan a copy-on-write snapshot in GetRegionSizeByRange GetRegionSizeByRange walked the region tree in ScanRegionLimit-sized chunks, releasing and immediately re-acquiring RegionsInfo.t between them. On a cluster with 1.5M regions one call therefore took the lock about 1500 times, and because Go's RWMutex is write-preferring, each acquisition let at most one waiting writer through. checkStores calls this for every preparing store on a 10s ticker, so a scale-out could hold region-heartbeat writes near that rate for as long as the stores stayed in the preparing state. Chunking is also not exact. Each chunk resumed at the previous region's end key, and scanRange resolves a 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 partly counted. pkg/btree is copy-on-write and its Clone is O(1), so add rootRangeSnapshot: a clone of the root tree that can be walked with no lock at all. Taking it needs the write lock, because Clone rewrites the source tree's copy-on-write context, but that critical section is O(1) regardless of region count. A snapshot freezes the tree's shape, not its values. The RegionInfo behind an item may still be replaced concurrently, though only by one covering the same key range, so a scan sees exactly one region per key. That is weaker than a scan under one read lock, where every region comes from a single instant, and the type documents the difference. Measured on 1M regions, arm64, with a partial-range scanner running against concurrent writers, writer cost per operation: GOMAXPROCS=1 2.1-6.2us -> 1.3-2.8us GOMAXPROCS=4 12.1-12.9us -> 0.81-1.19us GOMAXPROCS=16 13.7-15.9us -> 1.08-1.16us BatchScanRegions is deliberately left on the read lock. It has the opposite shape -- short scans at a high call rate -- so an exclusive clone per call costs much more than the read lock it would replace: eight-region scans went from 3.1us to 20.7us at GOMAXPROCS=4. Only a shared snapshot invalidated on structural change would suit that caller. A comment records the measurement so the next reader does not have to rediscover it. Signed-off-by: Vlad Dmitriev --- pkg/core/region.go | 59 +++++++++++-------- pkg/core/region_tree.go | 108 ++++++++++++++++++++++++++++++----- pkg/core/region_tree_test.go | 72 +++++++++++++++++++++++ 3 files changed, 202 insertions(+), 37 deletions(-) diff --git a/pkg/core/region.go b/pkg/core/region.go index c3db85d6f0..ceb25f7115 100644 --- a/pkg/core/region.go +++ b/pkg/core/region.go @@ -1453,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() @@ -2159,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 { @@ -2250,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 } diff --git a/pkg/core/region_tree.go b/pkg/core/region_tree.go index 8570622b21..8b683ca5cf 100644 --- a/pkg/core/region_tree.go +++ b/pkg/core/region_tree.go @@ -356,10 +356,80 @@ 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 }) @@ -371,25 +441,35 @@ func (t *regionTree) find(item *regionItem) *regionItem { 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}} - // find if there is a region with key range [s, d), s <= startKey < d - fn := func(item *regionItem) bool { - r := item - return f(r.getRegion()) - } start := newRegionItem(region) - startItem := t.find(start) + // find if there is a region with key range [s, d), s <= startKey < d + 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 diff --git a/pkg/core/region_tree_test.go b/pkg/core/region_tree_test.go index 1826b05d0a..fd35c29e8b 100644 --- a/pkg/core/region_tree_test.go +++ b/pkg/core/region_tree_test.go @@ -243,6 +243,78 @@ 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) *RegionInfo { + region := NewTestRegionInfo(id, 1, []byte(start), []byte(end), opts...) + _, err := regions.AtomicCheckAndPutRegion(ContextTODO(), region) + re.NoError(err) + return region + } + + 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)) +} + func TestRegionTreeSplitAndMerge(t *testing.T) { re := require.New(t) tree := newRegionTree() From a927d39935cc9f97fac458c19d4577a3692fe3b9 Mon Sep 17 00:00:00 2001 From: Vlad Dmitriev Date: Wed, 12 Aug 2026 14:52:20 -0700 Subject: [PATCH 5/6] core, btree: cover clone-under-mutation and snapshot semantics btree Clone was covered by one test that compares Ascend order and nothing else, while the test that checks the per-node index bookkeeping never clones. Nothing exercised mutableFor's copying of the indices array, even though GetAt and GetWithIndex depend on it, and through them GetCountByRange and RandomRegions. A fault there would return a wrong rank or a wrong item with no panic and no error. - TestBTreeSizeInfoAfterCloneG clones repeatedly while filling and then emptying a tree, and re-checks the length, rank and index invariants on the original and on every clone. - TestCloneReadWhileWritingOriginalG reads a clone while the original is written, which is how PD uses Clone. Writes are serialised, because BTreeG does not support concurrent writers; only the lock-free reads are concurrent. Both pass at degree 2, 32 and 64. On the core side: - TestRootRangeSnapshotIsolation applies a same-range update, a split and a removal after taking a snapshot, and pins down both halves of the contract: the key ranges are frozen, the values are not. - TestRootRangeSnapshotEquivalence checks a snapshot against the live tree, checks GetRegionSizeByRange against a brute-force sum over ScanRegions for random ranges, and checks that taking a snapshot changes no reference count, which is what keeps it O(1) and keeps RegionInfo.ref meaning "present in a tree". - TestGetRegionSizeByRangeConcurrent runs the lock-free walk against writers doing same-range updates, shrinks and restores, and cross-checks the subtrees afterwards. - TestRegionItemRangeImmutability holds a snapshot across a thousand updates and re-checks the ordering of the live tree and the snapshot after each one. - TestRootRangeSnapshotAdjacencyUnderMerges checks that a snapshot of a contiguous key space stays contiguous while the tree changes underneath it. Merges drive the churn, and the split half of each cycle is made atomic with respect to taking a snapshot: a split reaches PD as two separate updates, so between them the key space really is missing a range and a scanner is right to report a hole however it locks. - TestGetRegionSizeByRange gains cases for boundaries inside a region, a range past the last region's start key, and an empty range. The last two fail without the fresh-item allocation, reporting start keys that go backwards -- the signature of a shared item that moved in one tree while a clone still holds it at its old position. Benchmarks for the snapshot itself, on 1M regions, arm64: BenchmarkSnapshotRootTree ~95-107ns, 4 allocs, flat from 100k to 4M regions BenchmarkSetRegionSameRangeWithSnapshot ~365ns, 0 allocs, same as without ...RangeChangedWithSnapshot 2557ns vs 2496ns, +36B/op Signed-off-by: Vlad Dmitriev --- pkg/btree/btree_generic_test.go | 140 +++++++++++++ pkg/core/region_test.go | 361 +++++++++++++++++++++++++++++++- pkg/core/region_tree_test.go | 76 +++++++ 3 files changed, 575 insertions(+), 2 deletions(-) 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_test.go b/pkg/core/region_test.go index 878341ece3..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) { @@ -846,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) { @@ -1002,6 +1026,85 @@ func BenchmarkSetRegionSameRange(b *testing.B) { } } +// 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 { @@ -1272,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_test.go b/pkg/core/region_tree_test.go index fd35c29e8b..96128e749b 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" @@ -315,6 +316,81 @@ func TestRootRangeSnapshotIsolation(t *testing.T) { 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() From 722f690226970631ed0972e5a1b8b5c08807b484 Mon Sep 17 00:00:00 2001 From: Vlad Dmitriev Date: Wed, 12 Aug 2026 15:04:55 -0700 Subject: [PATCH 6/6] core: drop an unused return value in a snapshot test helper Satisfies the unparam linter. Signed-off-by: Vlad Dmitriev --- pkg/core/region_tree_test.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pkg/core/region_tree_test.go b/pkg/core/region_tree_test.go index 96128e749b..b72a1ab811 100644 --- a/pkg/core/region_tree_test.go +++ b/pkg/core/region_tree_test.go @@ -274,11 +274,10 @@ func treeRanges(regions *RegionsInfo) [][2]string { func TestRootRangeSnapshotIsolation(t *testing.T) { re := require.New(t) regions := NewRegionsInfo() - put := func(id uint64, start, end string, opts ...RegionCreateOption) *RegionInfo { + 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) - return region } put(1, "a", "b", SetApproximateSize(10))