Repository: tikv/client-go
Pull request #1477 (2024-10-16 07:01:12)

membuffer: improve ART's code and test coverage
ref pingcap/tidb#55287
This PR tries to make the code simplify and increase the test coverage (only some unreachable paths are not covered).

ti-chi-bot[bot]: [APPROVALNOTIFIER] This PR is **NOT APPROVED**

This pull-request has been approved by:
**Once this PR has been reviewed and has the lgtm label**, please ask for approval from [you06](https://github.com/you06), **ensuring that each of them provides their approval before proceeding**. For more information see [the Code Review Process](https://book.prow.tidb.net/#/workflows/pr).

The full list of commands accepted by this bot can be found [here](https://prow.tidb.net/command-help?repo=tikv%2Fclient-go).

<details open>
Needs approval from an approver in each of these files:

- **[OWNERS](https://github.com/tikv/client-go/blob/master/OWNERS)**

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":["you06"]} -->

Patch:
diff --git a/error/error.go b/error/error.go
index adc9776..2cf04a6 100644
--- a/error/error.go
+++ b/error/error.go
@@ -212,6 +212,14 @@ func (e *ErrTxnTooLarge) Error() string {
 	return fmt.Sprintf("txn too large, size: %v.", e.Size)
 }
 
+type ErrKeyTooLarge struct {
+	KeySize int
+}
+
+func (e *ErrKeyTooLarge) Error() string {
+	return fmt.Sprintf("key size too large, size: %v.", e.KeySize)
+}
+
 // ErrEntryTooLarge is the error when a key value entry is too large.
 type ErrEntryTooLarge struct {
 	Limit uint64
diff --git a/internal/locate/region_cache.go b/internal/locate/region_cache.go
index 01dfa38..c70e137 100644
--- a/internal/locate/region_cache.go
+++ b/internal/locate/region_cache.go
@@ -2223,7 +2223,18 @@ func (c *RegionCache) scanRegions(bo *retry.Backoffer, startKey, endKey []byte,
 			backoffErr = errors.Errorf("PD returned regions have gaps, limit: %d", limit)
 			continue
 		}
-		return c.handleRegionInfos(bo, regionsInfo, true)
+		validRegions, err := c.handleRegionInfos(bo, regionsInfo, true)
+		if err != nil {
+			return nil, err
+		}
+		// If the region information is loaded from the local disk and the current leader has not
+		// yet reported a heartbeat to PD, the region information scanned at this time will not include the leader.
+		// Retry if there is no valid regions with leaders.
+		if len(validRegions) == 0 {
+			backoffErr = errors.Errorf("All returned regions have no leaders, limit: %d", limit)
+			continue
+		}
+		return validRegions, nil
 	}
 }
 
@@ -2290,7 +2301,18 @@ func (c *RegionCache) batchScanRegions(bo *retry.Backoffer, keyRanges []pd.KeyRa
 			)
 			continue
 		}
-		return c.handleRegionInfos(bo, regionsInfo, opt.needRegionHasLeaderPeer)
+		validRegions, err := c.handleRegionInfos(bo, regionsInfo, opt.needRegionHasLeaderPeer)
+		if err != nil {
+			return nil, err
+		}
+		// If the region information is loaded from the local disk and the current leader has not
+		// yet reported a heartbeat to PD, the region information scanned at this time will not include the leader.
+		// Retry if there is no valid regions with leaders.
+		if len(validRegions) == 0 {
+			backoffErr = errors.Errorf("All returned regions have no leaders, limit: %d", limit)
+			continue
+		}
+		return validRegions, nil
 	}
 }
 
@@ -2397,7 +2419,7 @@ func (c *RegionCache) handleRegionInfos(bo *retry.Backoffer, regionsInfo []*pd.R
 		regions = append(regions, region)
 	}
 	if len(regions) == 0 {
-		return nil, errors.New("receive Regions with no peer")
+		return nil, nil
 	}
 	if len(regions) < len(regionsInfo) {
 		logutil.Logger(context.Background()).Debug(
diff --git a/internal/unionstore/art/art.go b/internal/unionstore/art/art.go
index 4fa76b4..7433745 100644
--- a/internal/unionstore/art/art.go
+++ b/internal/unionstore/art/art.go
@@ -47,7 +47,7 @@ type ART struct {
 	len             int
 	size            int
 
-	// The lastTraversedNode stores addr in uint64 of the last traversed node, include search and recursiveInsert.
+	// The lastTraversedNode stores addr in uint64 of the last traversed node, includes search and recursiveInsert.
 	// Compare to atomic.Pointer, atomic.Uint64 can avoid heap allocation, so it's more efficient.
 	lastTraversedNode atomic.Uint64
 	hitCount          atomic.Uint64
@@ -68,8 +68,13 @@ func New() *ART {
 }
 
 func (t *ART) Get(key []byte) ([]byte, error) {
+	if t.vlogInvalid {
+		// panic for easier debugging.
+		panic("vlog is reset")
+	}
+
 	// 1. search the leaf node.
-	_, leaf := t.search(key)
+	_, leaf := t.traverse(key, false)
 	if leaf == nil || leaf.vAddr.IsNull() {
 		return nil, tikverr.ErrNotExist
 	}
@@ -79,7 +84,7 @@ func (t *ART) Get(key []byte) ([]byte, error) {
 
 // GetFlags returns the latest flags associated with key.
 func (t *ART) GetFlags(key []byte) (kv.KeyFlags, error) {
-	_, leaf := t.search(key)
+	_, leaf := t.traverse(key, false)
 	if leaf == nil {
 		return 0, tikverr.ErrNotExist
 	}
@@ -90,6 +95,17 @@ func (t *ART) GetFlags(key []byte) (kv.KeyFlags, error) {
 }
 
 func (t *ART) Set(key artKey, value []byte, ops ...kv.FlagsOp) error {
+	if t.vlogInvalid {
+		// panic for easier debugging.
+		panic("vlog is reset")
+	}
+
+	if len(key) > MaxKeyLen {
+		return &tikverr.ErrKeyTooLarge{
+			KeySize: len(key),
+		}
+	}
+
 	if value != nil {
 		if size := uint64(len(key) + len(value)); size > t.entrySizeLimit {
 			return &tikverr.ErrEntryTooLarge{
@@ -98,11 +114,12 @@ func (t *ART) Set(key artKey, value []byte, ops ...kv.FlagsOp) error {
 			}
 		}
 	}
+
 	if len(t.stages) == 0 {
 		t.dirty = true
 	}
 	// 1. create or search the existing leaf in the tree.
-	addr, leaf := t.recursiveInsert(key)
+	addr, leaf := t.traverse(key, true)
 	// 2. set the value and flags.
 	t.setValue(addr, leaf, value, ops)
 	if uint64(t.Size()) > t.bufferSizeLimit {
@@ -111,8 +128,8 @@ func (t *ART) Set(key artKey, value []byte, ops ...kv.FlagsOp) error {
 	return nil
 }
 
-// search wraps searchImpl with cache.
-func (t *ART) search(key artKey) (arena.MemdbArenaAddr, *artLeaf) {
+// traverse wraps search and recursiveInsert with cache.
+func (t *ART) traverse(key artKey, insert bool) (arena.MemdbArenaAddr, *artLeaf) {
 	// check cache
 	addr, leaf, found := t.checkKeyInCache(key)
 	if found {
@@ -120,17 +137,21 @@ func (t *ART) search(key artKey) (arena.MemdbArenaAddr, *artLeaf) {
 		return addr, leaf
 	}
 	t.missCount.Add(1)
-	addr, leaf = t.searchImpl(key)
+	if insert {
+		addr, leaf = t.recursiveInsert(key)
+	} else {
+		addr, leaf = t.search(key)
+	}
 	if !addr.IsNull() {
 		t.updateLastTraversed(addr)
 	}
 	return addr, leaf
 }
 
-// searchImpl looks up the leaf with the given key.
+// search looks up the leaf with the given key.
 // It returns the memory arena address and leaf itself it there is a match leaf,
 // returns arena.NullAddr and nil if the key is not found.
-func (t *ART) searchImpl(key artKey) (arena.MemdbArenaAddr, *artLeaf) {
+func (t *ART) search(key artKey) (arena.MemdbArenaAddr, *artLeaf) {
 	current := t.root
 	if current == nullArtNode {
 		return arena.NullAddr, nil
@@ -179,25 +200,9 @@ func (t *ART) searchImpl(key artKey) (arena.MemdbArenaAddr, *artLeaf) {
 	}
 }
 
-// recursiveInsert wraps recursiveInsertImpl with cache.
-func (t *ART) recursiveInsert(key artKey) (arena.MemdbArenaAddr, *artLeaf) {
-	addr, leaf, found := t.checkKeyInCache(key)
-	if found {
-		t.hitCount.Add(1)
-		return addr, leaf
-	}
-	t.missCount.Add(1)
-	addr, leaf = t.recursiveInsertImpl(key)
-	if !addr.IsNull() {
-		t.updateLastTraversed(addr)
-	}
-	return addr, leaf
-}
-
-// recursiveInsertImpl returns the node address of the key.
+// recursiveInsert returns the node address of the key.
 // It will insert the key if not exists, returns the newly inserted or existing leaf.
-func (t *ART) recursiveInsertImpl(key artKey) (arena.MemdbArenaAddr, *artLeaf) {
-
+func (t *ART) recursiveInsert(key artKey) (arena.MemdbArenaAddr, *artLeaf) {
 	// lazy init root node and allocator.
 	// this saves memory for read only txns.
 	if t.root.addr.IsNull() {
@@ -300,11 +305,7 @@ func (t *ART) expandLeaf(key artKey, depth uint32, prev, current artNode) (arena
 	newAn.addChild(&t.allocator, l2Key.charAt(int(depth)), !l2Key.valid(int(depth)), leaf2Addr)
 
 	// swap the old leaf with the new node4.
-	if prev == nullArtNode {
-		t.root = newAn
-	} else {
-		prev.replaceChild(&t.allocator, key.charAt(prevDepth), newAn)
-	}
+	prev.replaceChild(&t.allocator, key.charAt(prevDepth), newAn)
 	return leaf2Addr.addr, leaf2
 }
 
@@ -328,30 +329,25 @@ func (t *ART) expandNode(key artKey, depth, mismatchIdx uint32, prev, current ar
 	newN4.setPrefix(key[depth:], mismatchIdx)
 
 	// update prefix for old node and move it as a child of the new node.
+	var prefix artKey
 	if currNode.prefixLen <= maxPrefixLen {
-		nodeKey := currNode.prefix[mismatchIdx]
-		currNode.prefixLen -= mismatchIdx + 1
-		copy(currNode.prefix[:], currNode.prefix[mismatchIdx+1:])
-		newAn.addChild(&t.allocator, nodeKey, false, current)
+		// node.prefixLen <= maxPrefixLen means all the prefix is in the prefix array.
+		// The char at mismatchIdx will be stored in the index of new node.
+		prefix = currNode.prefix[:]
 	} else {
-		currNode.prefixLen -= mismatchIdx + 1
+		// Unless, we need to find the prefix in the leaf.
+		// Any leaves in the node should have the same prefix, we use minimum node here.
 		leafArtNode := minimum(&t.allocator, current)
-		leaf := leafArtNode.asLeaf(&t.allocator)
-		leafKey := artKey(leaf.GetKey())
-		kMin := depth + mismatchIdx + 1
-		kMax := depth + mismatchIdx + 1 + min(currNode.prefixLen, maxPrefixLen)
-		copy(currNode.prefix[:], leafKey[kMin:kMax])
-		newAn.addChild(&t.allocator, leafKey.charAt(int(depth+mismatchIdx)), !leafKey.valid(int(depth)), current)
+		prefix = leafArtNode.asLeaf(&t.allocator).GetKey()[depth : depth+currNode.prefixLen]
 	}
+	nodeChar := prefix[mismatchIdx]
+	currNode.setPrefix(prefix[mismatchIdx+1:], currNode.prefixLen-mismatchIdx-1)
+	newAn.addChild(&t.allocator, nodeChar, false, current)
 
 	// insert the artLeaf into new node
 	newLeafAddr, newLeaf := t.newLeaf(key)
 	newAn.addChild(&t.allocator, key.charAt(int(depth+mismatchIdx)), !key.valid(int(depth+mismatchIdx)), newLeafAddr)
-	if prev == nullArtNode {
-		t.root = newAn
-	} else {
-		prev.replaceChild(&t.allocator, key.charAt(prevDepth), newAn)
-	}
+	prev.replaceChild(&t.allocator, key.charAt(prevDepth), newAn)
 	return newLeafAddr.addr, newLeaf
 }
 
@@ -577,7 +573,6 @@ func (t *ART) SelectValueHistory(key []byte, predicate func(value []byte) bool)
 		return nil, nil
 	}
 	return t.allocator.vlogAllocator.GetValue(result), nil
-
 }
 
 func (t *ART) SetMemoryFootprintChangeHook(hook func(uint64)) {
diff --git a/internal/unionstore/art/art_arena.go b/internal/unionstore/art/art_arena.go
index 84ca211..f196bf0 100644
--- a/internal/unionstore/art/art_arena.go
+++ b/internal/unionstore/art/art_arena.go
@@ -25,6 +25,8 @@ import (
 // reusing blocks reduces the memory pieces.
 type nodeArena struct {
 	arena.MemdbArena
+	// The ART node will expand to a higher capacity, and the address of the freed node will be stored in the free list for reuse.
+	// By reusing the freed node, memory usage and fragmentation can be reduced.
 	freeNode4  []arena.MemdbArenaAddr
 	freeNode16 []arena.MemdbArenaAddr
 	freeNode48 []arena.MemdbArenaAddr
diff --git a/internal/unionstore/art/art_iterator.go b/internal/unionstore/art/art_iterator.go
index ad0c852..38f6c6a 100644
--- a/internal/unionstore/art/art_iterator.go
+++ b/internal/unionstore/art/art_iterator.go
@@ -16,6 +16,7 @@ package art
 
 import (
 	"bytes"
+	"fmt"
 	"sort"
 
 	"github.com/pkg/errors"
@@ -156,140 +157,161 @@ func (it *Iterator) init(lowerBound, upperBound []byte) {
 		}
 	}
 
+	startKey, endKey := lowerBound, upperBound
 	if it.reverse {
-		it.inner.idxes, it.inner.nodes = it.seek(upperBound)
-		if len(lowerBound) == 0 {
-			it.endAddr = arena.NullAddr
-		} else {
-			helper := new(baseIter)
-			helper.allocator = &it.tree.allocator
-			helper.idxes, helper.nodes = it.seek(lowerBound)
-			if it.inner.compare(helper) > 0 {
-				// lowerBound is inclusive, call next to find the smallest leaf node that >= lowerBound.
-				it.endAddr = helper.next().addr
-				if it.inner.compare(helper) < 0 || len(helper.idxes) == 0 {
-					it.valid = false
-				}
-				return
-			}
-			it.valid = false
-		}
-		return
+		startKey, endKey = upperBound, lowerBound
 	}
 
-	it.inner.idxes, it.inner.nodes = it.seek(lowerBound)
-	if len(upperBound) == 0 {
-		it.endAddr = arena.NullAddr
+	if len(startKey) == 0 {
+		it.inner.seekToFirst(it.tree.root, it.reverse)
 	} else {
-		helper := new(baseIter)
-		helper.allocator = &it.tree.allocator
-		helper.idxes, helper.nodes = it.seek(upperBound)
-		if it.inner.compare(helper) < 0 {
-			// upperBound is exclusive, so we move the helper cursor to the previous node, which is the true endAddr.
-			it.endAddr = helper.prev().addr
-			if it.inner.compare(helper) > 0 || len(helper.idxes) == 0 {
-				it.valid = false
-			}
-			return
-		}
+		it.inner.seek(it.tree.root, startKey)
+	}
+	if len(endKey) == 0 {
+		it.endAddr = arena.NullAddr
+		return
+	}
+
+	helper := new(baseIter)
+	helper.allocator = &it.tree.allocator
+	helper.seek(it.tree.root, endKey)
+	cmp := it.inner.compare(helper)
+	if cmp == 0 {
+		// no keys exist between start key and end key, set the iterator to invalid.
 		it.valid = false
 		return
 	}
+
+	if it.reverse {
+		it.endAddr = helper.next().addr
+	} else {
+		it.endAddr = helper.prev().addr
+	}
+	cmp = it.inner.compare(helper)
+	if cmp == 0 {
+		// the current key is valid.
+		return
+	}
+	// in asc scan, if cmp > 0, it means current key is larger than end key, set the iterator to invalid.
+	// in desc scan, if cmp < 0, it means current key is less than end key, set the iterator to invalid.
+	if cmp < 0 == it.reverse || len(helper.idxes) == 0 {
+		it.valid = false
+	}
 }
 
-// seek the first node and index that >= key, return the indexes and nodes of the lookup path
-// nodes[0] is the root node
-func (it *Iterator) seek(key artKey) ([]int, []artNode) {
-	curr := it.tree.root
-	depth := uint32(0)
-	idxes := make([]int, 0, 8)
-	nodes := make([]artNode, 0, 8)
+// baseIter is the inner iterator for ART tree.
+// You need to call seek or seekToFirst to initialize the iterator.
+// after initialization, you can call next or prev to iterate the ART.
+// next or prev returns nullArtNode if there is no more leaf node.
+type baseIter struct {
+	allocator *artAllocator
+	// the baseIter iterate the ART tree in DFS order, the idxes and nodes are the current visiting stack.
+	idxes []int
+	nodes []artNode
+}
+
+// seekToFirst seeks the boundary of the tree, sequential or reverse.
+func (it *baseIter) seekToFirst(root artNode, reverse bool) {
+	// if the seek key is empty, it means -inf or +inf, return root node directly.
+	it.nodes = []artNode{root}
+	if reverse {
+		it.idxes = []int{node256cap}
+	} else {
+		it.idxes = []int{inplaceIndex}
+	}
+}
+
+// seek the first node and index that >= key, nodes[0] is the root node
+func (it *baseIter) seek(root artNode, key artKey) {
 	if len(key) == 0 {
-		// if the seek key is empty, it means -inf or +inf, return root node directly.
-		nodes = append(nodes, curr)
-		if it.reverse {
-			idxes = append(idxes, node256cap)
-		} else {
-			idxes = append(idxes, inplaceIndex)
-		}
-		return idxes, nodes
+		panic("seek with empty key is not allowed")
 	}
+	curr := root
+	depth := uint32(0)
 	var node *nodeBase
 	for {
 		if curr.isLeaf() {
 			if key.valid(int(depth)) {
-				lf := curr.asLeaf(&it.tree.allocator)
+				lf := curr.asLeaf(it.allocator)
 				if bytes.Compare(key, lf.GetKey()) > 0 {
 					// the seek key is not exist, and it's longer and larger than the current leaf's key.
 					// e.g. key: [1, 1, 1], leaf: [1, 1].
-					idxes[len(idxes)-1]++
+					it.idxes[len(it.idxes)-1]++
 				}
 			}
-			break
+			return
 		}
 
-		node = curr.asNode(&it.tree.allocator)
+		it.nodes = append(it.nodes, curr)
+		node = curr.asNode(it.allocator)
 		if node.prefixLen > 0 {
-			mismatchIdx := node.matchDeep(&it.tree.allocator, &curr, key, depth)
+			mismatchIdx := node.matchDeep(it.allocator, &curr, key, depth)
 			if mismatchIdx < node.prefixLen {
-				// no leaf node is match with the seek key
-				leafNode := minimum(&it.tree.allocator, curr)
-				leafKey := leafNode.asLeaf(&it.tree.allocator).GetKey()
-				if mismatchIdx+depth == uint32(len(key)) || key[depth+mismatchIdx] < leafKey[depth+mismatchIdx] {
-					// key < leafKey, set index to -1 means all the children are larger than the seek key
-					idxes = append(idxes, -1)
+				// Check whether the seek key is smaller than the prefix, as no leaf node matches the seek key.
+				// If the seek key is smaller than the prefix, all the children are located on the right side of the seek key.
+				// Otherwise, the children are located on the left side of the seek key.
+				var prefix []byte
+				if mismatchIdx < maxPrefixLen {
+					prefix = node.prefix[:]
 				} else {
-					// key > leafKey, set index to 256 means all the children are less than the seek key
-					idxes = append(idxes, node256cap)
+					leafNode := minimum(it.allocator, curr)
+					prefix = leafNode.asLeaf(it.allocator).getKeyDepth(depth)
 				}
-				nodes = append(nodes, curr)
-				return idxes, nodes
+				if mismatchIdx+depth == uint32(len(key)) || key[depth+mismatchIdx] < prefix[mismatchIdx] {
+					// mismatchIdx + depth == len(key) indicates that the seek key is a prefix of any leaf in the current node, implying that key < leafKey.
+					// If key < leafKey, set index to -1 means all the children are larger than the seek key.
+					it.idxes = append(it.idxes, -1)
+				} else {
+					// If key > prefix, set index to 256 means all the children are smaller than the seek key.
+					it.idxes = append(it.idxes, node256cap)
+				}
+				return
 			}
 			depth += min(mismatchIdx, node.prefixLen)
 		}
 
-		nodes = append(nodes, curr)
 		char := key.charAt(int(depth))
-		idx, next := curr.findChild(&it.tree.allocator, char, !key.valid(int(depth)))
+		idx, next := curr.findChild(it.allocator, char, !key.valid(int(depth)))
 		if next.addr.IsNull() {
-			nextIdx := 0
-			switch curr.kind {
-			case typeNode4:
-				n4 := curr.asNode4(&it.tree.allocator)
-				for ; nextIdx < int(n4.nodeNum); nextIdx++ {
-					if n4.keys[nextIdx] >= char {
-						break
-					}
-				}
-			case typeNode16:
-				n16 := curr.asNode16(&it.tree.allocator)
-				nextIdx, _ = sort.Find(int(n16.nodeNum), func(i int) int {
-					if n16.keys[i] < char {
-						return 1
-					}
-					return -1
-				})
-			case typeNode48:
-				n48 := curr.asNode48(&it.tree.allocator)
-				nextIdx = n48.nextPresentIdx(int(char))
-			case typeNode256:
-				n256 := curr.asNode256(&it.tree.allocator)
-				nextIdx = n256.nextPresentIdx(int(char))
-			}
-			idxes = append(idxes, nextIdx)
-			return idxes, nodes
+			nextIdx := seekToIdx(it.allocator, curr, char)
+			it.idxes = append(it.idxes, nextIdx)
+			return
 		}
-		idxes = append(idxes, idx)
+		it.idxes = append(it.idxes, idx)
 		curr = next
 		depth++
 	}
-	return idxes, nodes
 }
 
-type baseIter struct {
-	allocator *artAllocator
-	idxes     []int
-	nodes     []artNode
+// seekToIdx finds the index where all nodes before it are less than the given character.
+func seekToIdx(a *artAllocator, curr artNode, char byte) int {
+	var nextIdx int
+	switch curr.kind {
+	case typeNode4:
+		n4 := curr.asNode4(a)
+		for ; nextIdx < int(n4.nodeNum); nextIdx++ {
+			if n4.keys[nextIdx] >= char {
+				break
+			}
+		}
+	case typeNode16:
+		n16 := curr.asNode16(a)
+		nextIdx, _ = sort.Find(int(n16.nodeNum), func(i int) int {
+			if n16.keys[i] < char {
+				return 1
+			}
+			return -1
+		})
+	case typeNode48:
+		n48 := curr.asNode48(a)
+		nextIdx = n48.nextPresentIdx(int(char))
+	case typeNode256:
+		n256 := curr.asNode256(a)
+		nextIdx = n256.nextPresentIdx(int(char))
+	default:
+		panic("invalid node type")
+	}
+	return nextIdx
 }
 
 // compare compares the path of nodes, return 1 if self > other, -1 if self < other, 0 if self == other
@@ -318,197 +340,209 @@ func (it *baseIter) compare(other *baseIter) int {
 // next returns the next leaf node
 // it returns nullArtNode if there is no more leaf node
 func (it *baseIter) next() artNode {
-	depth := len(it.nodes) - 1
-	curr := it.nodes[depth]
-	idx := it.idxes[depth]
-	switch curr.kind {
-	case typeNode4:
-		n4 := it.allocator.getNode4(curr.addr)
-		if idx == inplaceIndex {
-			idx = 0 // mark in-place leaf is visited
-			it.idxes[depth] = idx
-			if !n4.inplaceLeaf.addr.IsNull() {
-				return n4.inplaceLeaf
+	for {
+		depth := len(it.nodes) - 1
+		curr := it.nodes[depth]
+		idx := it.idxes[depth]
+		var child *artNode
+		switch curr.kind {
+		case typeNode4:
+			n4 := it.allocator.getNode4(curr.addr)
+			if idx == inplaceIndex {
+				idx = 0 // mark in-place leaf is visited
+				it.idxes[depth] = idx
+				if !n4.inplaceLeaf.addr.IsNull() {
+					return n4.inplaceLeaf
+				}
+			} else if idx == node4cap {
+				break
 			}
-		} else if idx == node4cap {
-			break
-		}
-		if idx < int(n4.nodeNum) {
-			it.idxes[depth] = idx
-			child := n4.children[idx]
-			if child.kind == typeLeaf {
-				it.idxes[depth]++
-				return child
+			if idx >= 0 && idx < int(n4.nodeNum) {
+				it.idxes[depth] = idx
+				child = &n4.children[idx]
+			} else if idx == int(n4.nodeNum) {
+				// idx == n4.nodeNum means this node is drain, break to pop stack.
+				break
+			} else {
+				panicForInvalidIndex(idx)
 			}
-			it.nodes = append(it.nodes, child)
-			it.idxes = append(it.idxes, inplaceIndex)
-			return it.next()
-		}
-	case typeNode16:
-		n16 := it.allocator.getNode16(curr.addr)
-		if idx == inplaceIndex {
-			idx = 0 // mark in-place leaf is visited
-			it.idxes[depth] = idx
-			if !n16.inplaceLeaf.addr.IsNull() {
-				return n16.inplaceLeaf
+		case typeNode16:
+			n16 := it.allocator.getNode16(curr.addr)
+			if idx == inplaceIndex {
+				idx = 0 // mark in-place leaf is visited
+				it.idxes[depth] = idx
+				if !n16.inplaceLeaf.addr.IsNull() {
+					return n16.inplaceLeaf
+				}
+			} else if idx == node16cap {
+				break
 			}
-		} else if idx == node16cap {
-			break
-		}
-		if idx < int(n16.nodeNum) {
-			it.idxes[depth] = idx
-			child := n16.children[idx]
-			if child.kind == typeLeaf {
-				it.idxes[depth]++
-				return child
+			if idx >= 0 && idx < int(n16.nodeNum) {
+				it.idxes[depth] = idx
+				child = &n16.children[idx]
+			} else if idx == int(n16.nodeNum) {
+				// idx == n16.nodeNum means this node is drain, break to pop stack.
+				break
+			} else {
+				panicForInvalidIndex(idx)
 			}
-			it.nodes = append(it.nodes, child)
-			it.idxes = append(it.idxes, inplaceIndex)
-			return it.next()
-		}
-	case typeNode48:
-		n48 := it.allocator.getNode48(curr.addr)
-		if idx == inplaceIndex {
-			idx = 0 // mark in-place leaf is visited
-			it.idxes[depth] = idx
-			if !n48.inplaceLeaf.addr.IsNull() {
-				return n48.inplaceLeaf
+		case typeNode48:
+			n48 := it.allocator.getNode48(curr.addr)
+			if idx == inplaceIndex {
+				idx = 0 // mark in-place leaf is visited
+				it.idxes[depth] = idx
+				if !n48.inplaceLeaf.addr.IsNull() {
+					return n48.inplaceLeaf
+				}
+			} else if idx == node256cap {
+				break
 			}
-		} else if idx == node256cap {
-			break
-		}
-		idx = n48.nextPresentIdx(idx)
-		if idx < node256cap {
-			it.idxes[depth] = idx
-			child := n48.children[n48.keys[idx]]
-			if child.kind == typeLeaf {
-				it.idxes[depth]++
-				return child
+			idx = n48.nextPresentIdx(idx)
+			if idx >= 0 && idx < node256cap {
+				it.idxes[depth] = idx
+				child = &n48.children[n48.keys[idx]]
+			} else if idx == node256cap {
+				// idx == node256cap means this node is drain, break to pop stack.
+				break
+			} else {
+				panicForInvalidIndex(idx)
 			}
-			it.nodes = append(it.nodes, child)
-			it.idxes = append(it.idxes, inplaceIndex)
-			return it.next()
-		}
-	case typeNode256:
-		n256 := it.allocator.getNode256(curr.addr)
-		if idx == inplaceIndex {
-			idx = 0 // mark in-place leaf is visited
-			it.idxes[depth] = idx
-			if !n256.inplaceLeaf.addr.IsNull() {
-				return n256.inplaceLeaf
+		case typeNode256:
+			n256 := it.allocator.getNode256(curr.addr)
+			if idx == inplaceIndex {
+				idx = 0 // mark in-place leaf is visited
+				it.idxes[depth] = idx
+				if !n256.inplaceLeaf.addr.IsNull() {
+					return n256.inplaceLeaf
+				}
+			} else if idx == node256cap {
+				break
+			}
+			idx = n256.nextPresentIdx(idx)
+			if idx >= 0 && idx < 256 {
+				it.idxes[depth] = idx
+				child = &n256.children[idx]
+			} else if idx == node256cap {
+				// idx == node256cap means this node is drain, break to pop stack.
+				break
+			} else {
+				panicForInvalidIndex(idx)
 			}
-		} else if idx == 256 {
-			break
+		default:
+			panic("invalid node type")
 		}
-		idx = n256.nextPresentIdx(idx)
-		if idx < 256 {
-			it.idxes[depth] = idx
-			child := n256.children[idx]
+		if child != nil {
 			if child.kind == typeLeaf {
 				it.idxes[depth]++
-				return child
+				return *child
 			}
-			it.nodes = append(it.nodes, child)
+			it.nodes = append(it.nodes, *child)
 			it.idxes = append(it.idxes, inplaceIndex)
-			return it.next()
+			continue
 		}
+		it.nodes = it.nodes[:depth]
+		it.idxes = it.idxes[:depth]
+		if depth == 0 {
+			return nullArtNode
+		}
+		it.idxes[depth-1]++
 	}
-	it.nodes = it.nodes[:depth]
-	it.idxes = it.idxes[:depth]
-	if depth == 0 {
-		return nullArtNode
-	}
-	it.idxes[depth-1]++
-	return it.next()
 }
 
 func (it *baseIter) prev() artNode {
-	depth := len(it.nodes) - 1
-	curr := it.nodes[depth]
-	idx := it.idxes[depth]
-	idx--
-	switch curr.kind {
-	case typeNode4:
-		n4 := it.allocator.getNode4(curr.addr)
-		idx = min(idx, int(n4.nodeNum)-1)
-		if idx >= 0 {
-			it.idxes[depth] = idx
-			child := n4.children[idx]
-			if child.kind == typeLeaf {
-				return child
-			}
-			it.nodes = append(it.nodes, child)
-			it.idxes = append(it.idxes, node256cap)
-			return it.prev()
-		} else if idx == inplaceIndex {
-			it.idxes[depth] = idx
-			if !n4.inplaceLeaf.addr.IsNull() {
-				return n4.inplaceLeaf
-			}
-		}
-	case typeNode16:
-		n16 := it.allocator.getNode16(curr.addr)
-		idx = min(idx, int(n16.nodeNum)-1)
-		if idx >= 0 {
-			it.idxes[depth] = idx
-			child := n16.children[idx]
-			if child.kind == typeLeaf {
-				return child
-			}
-			it.nodes = append(it.nodes, child)
-			it.idxes = append(it.idxes, node256cap)
-			return it.prev()
-		} else if idx == inplaceIndex {
-			it.idxes[depth] = idx
-			if !n16.inplaceLeaf.addr.IsNull() {
-				return n16.inplaceLeaf
-			}
-		}
-	case typeNode48:
-		n48 := it.allocator.getNode48(curr.addr)
-		if idx >= 0 {
-			idx = n48.prevPresentIdx(idx)
-		}
-		if idx >= 0 {
-			it.idxes[depth] = idx
-			child := n48.children[n48.keys[idx]]
-			if child.kind == typeLeaf {
-				return child
+	for {
+		depth := len(it.nodes) - 1
+		curr := it.nodes[depth]
+		idx := it.idxes[depth]
+		idx--
+		if idx != notExistIndex {
+			var child *artNode
+			switch curr.kind {
+			case typeNode4:
+				n4 := it.allocator.getNode4(curr.addr)
+				idx = min(idx, int(n4.nodeNum)-1)
+				if idx >= 0 {
+					it.idxes[depth] = idx
+					child = &n4.children[idx]
+				} else if idx == inplaceIndex {
+					it.idxes[depth] = idx
+					if !n4.inplaceLeaf.addr.IsNull() {
+						return n4.inplaceLeaf
+					}
+				} else {
+					panicForInvalidIndex(idx)
+				}
+			case typeNode16:
+				n16 := it.allocator.getNode16(curr.addr)
+				idx = min(idx, int(n16.nodeNum)-1)
+				if idx >= 0 {
+					it.idxes[depth] = idx
+					child = &n16.children[idx]
+				} else if idx == inplaceIndex {
+					it.idxes[depth] = idx
+					if !n16.inplaceLeaf.addr.IsNull() {
+						return n16.inplaceLeaf
+					}
+				} else {
+					panicForInvalidIndex(idx)
+				}
+			case typeNode48:
+				n48 := it.allocator.getNode48(curr.addr)
+				if idx >= 0 && n48.present[idx>>n48s]&(1<<(idx%n48m)) == 0 {
+					// if idx >= 0 and n48.keys[idx] is not present, goto the previous present key.
+					// for idx < 0, we check inplaceLeaf later.
+					idx = n48.prevPresentIdx(idx)
+				}
+				if idx >= 0 {
+					it.idxes[depth] = idx
+					child = &n48.children[n48.keys[idx]]
+				} else if idx == inplaceIndex {
+					it.idxes[depth] = idx
+					if !n48.inplaceLeaf.addr.IsNull() {
+						return n48.inplaceLeaf
+					}
+				} else {
+					panicForInvalidIndex(idx)
+				}
+			case typeNode256:
+				n256 := it.allocator.getNode256(curr.addr)
+				if idx >= 0 && n256.present[idx>>n48s]&(1<<(idx%n48m)) == 0 {
+					// if idx >= 0 and n256.keys[idx] is not present, goto the previous present key.
+					// for idx < 0, we check inplaceLeaf later.
+					idx = n256.prevPresentIdx(idx)
+				}
+				if idx >= 0 {
+					it.idxes[depth] = idx
+					child = &n256.children[idx]
+				} else if idx == inplaceIndex {
+					it.idxes[depth] = idx
+					if !n256.inplaceLeaf.addr.IsNull() {
+						return n256.inplaceLeaf
+					}
+				} else {
+					panicForInvalidIndex(idx)
+				}
+			default:
+				panic("invalid node type")
 			}
-			it.nodes = append(it.nodes, child)
-			it.idxes = append(it.idxes, node256cap)
-			return it.prev()
-		} else if idx == inplaceIndex {
-			it.idxes[depth] = idx
-			if !n48.inplaceLeaf.addr.IsNull() {
-				return n48.inplaceLeaf
+			if child != nil {
+				if child.kind == typeLeaf {
+					return *child
+				}
+				it.nodes = append(it.nodes, *child)
+				it.idxes = append(it.idxes, node256cap)
+				continue
 			}
 		}
-	case typeNode256:
-		n256 := it.allocator.getNode256(curr.addr)
-		if idx >= 0 {
-			idx = n256.prevPresentIdx(idx)
-		}
-		if idx >= 0 {
-			it.idxes[depth] = idx
-			child := n256.children[idx]
-			if child.kind == typeLeaf {
-				return child
-			}
-			it.nodes = append(it.nodes, child)
-			it.idxes = append(it.idxes, node256cap)
-			return it.prev()
-		} else if idx == -1 {
-			it.idxes[depth] = idx
-			if !n256.inplaceLeaf.addr.IsNull() {
-				return n256.inplaceLeaf
-			}
+		it.nodes = it.nodes[:depth]
+		it.idxes = it.idxes[:depth]
+		if depth == 0 {
+			return nullArtNode
 		}
 	}
-	it.nodes = it.nodes[:depth]
-	it.idxes = it.idxes[:depth]
-	if depth == 0 {
-		return nullArtNode
-	}
-	return it.prev()
+}
+
+func panicForInvalidIndex(idx int) {
+	msg := fmt.Sprintf("ART iterator meets an invalid index %d", idx)
+	panic(msg)
 }
diff --git a/internal/unionstore/art/art_node.go b/internal/unionstore/art/art_node.go
index 75c8d11..01d9169 100644
--- a/internal/unionstore/art/art_node.go
+++ b/internal/unionstore/art/art_node.go
@@ -38,11 +38,14 @@ const (
 )
 
 const (
-	maxPrefixLen  = 20
-	node4cap      = 4
-	node16cap     = 16
-	node48cap     = 48
-	node256cap    = 256
+	maxPrefixLen = 20
+	node4cap     = 4
+	node16cap    = 16
+	node48cap    = 48
+	node256cap   = 256
+	// inplaceIndex is a special index to indicate the index of an in-place leaf in a node,
+	// the in-place leaf has the same key with its parent node and doesn't occupy the quota of the node.
+	// the other valid index of a node is [0, nodeNum), all the other leaves in the node have larger key than the in-place leaf.
 	inplaceIndex  = -1
 	notExistIndex = -2
 )
@@ -117,6 +120,8 @@ type node256 struct {
 	children [node256cap]artNode
 }
 
+const MaxKeyLen = math.MaxUint16
+
 type artLeaf struct {
 	vAddr arena.MemdbArenaAddr
 	klen  uint16
@@ -350,7 +355,7 @@ func (an *artNode) asNode256(a *artAllocator) *node256 {
 // longestCommonPrefix returns the length of the longest common prefix of two keys.
 // the LCP is calculated from the given depth, you need to guarantee l1Key[:depth] equals to l2Key[:depth] before calling this function.
 func longestCommonPrefix(l1Key, l2Key artKey, depth uint32) uint32 {
-	idx, limit := depth, min(uint32(len(l1Key)), uint32(len(l2Key)))
+	idx, limit := depth, uint32(min(len(l1Key), len(l2Key)))
 	// TODO: possible optimization
 	// Compare the key by loop can be very slow if the final LCP is large.
 	// Maybe optimize it by comparing the key in chunks if the limit exceeds certain threshold.
@@ -394,8 +399,8 @@ func minimum(a *artAllocator, an artNode) artNode {
 			}
 			idx := n256.nextPresentIdx(0)
 			an = n256.children[idx]
-		case typeInvalid:
-			return nullArtNode
+		default:
+			panic("invalid node kind")
 		}
 	}
 }
@@ -526,52 +531,52 @@ func (an *artNode) addChild(a *artAllocator, c byte, inplace bool, child artNode
 	case typeNode256:
 		return an.addChild256(a, c, child)
 	}
-	return false
+	panic("add child failed")
 }
 
 func (an *artNode) addChild4(a *artAllocator, c byte, child artNode) bool {
-	node := an.asNode4(a)
+	n4 := an.asNode4(a)
 
-	if node.nodeNum >= node4cap {
-		an.grow(a)
-		an.addChild(a, c, false, child)
+	if n4.nodeNum >= node4cap {
+		an.growNode4(n4, a)
+		an.addChild16(a, c, child)
 		return true
 	}
 
 	i := uint8(0)
-	for ; i < node.nodeNum; i++ {
-		if c <= node.keys[i] {
-			if testing.Testing() && c == node.keys[i] {
+	for ; i < n4.nodeNum; i++ {
+		if c <= n4.keys[i] {
+			if testing.Testing() && c == n4.keys[i] {
 				panic("key already exists")
 			}
 			break
 		}
 	}
 
-	if i < node.nodeNum {
-		copy(node.keys[i+1:node.nodeNum+1], node.keys[i:node.nodeNum])
-		copy(node.children[i+1:node.nodeNum+1], node.children[i:node.nodeNum])
+	if i < n4.nodeNum {
+		copy(n4.keys[i+1:n4.nodeNum+1], n4.keys[i:n4.nodeNum])
+		copy(n4.children[i+1:n4.nodeNum+1], n4.children[i:n4.nodeNum])
 	}
-	node.keys[i] = c
-	node.children[i] = child
-	node.nodeNum++
+	n4.keys[i] = c
+	n4.children[i] = child
+	n4.nodeNum++
 	return false
 }
 
 func (an *artNode) addChild16(a *artAllocator, c byte, child artNode) bool {
-	node := an.asNode16(a)
+	n16 := an.asNode16(a)
 
-	if node.nodeNum >= node16cap {
-		an.grow(a)
-		an.addChild(a, c, false, child)
+	if n16.nodeNum >= node16cap {
+		an.growNode16(n16, a)
+		an.addChild48(a, c, child)
 		return true
 	}
 
-	i, found := sort.Find(int(node.nodeNum), func(i int) int {
-		if node.keys[i] < c {
+	i, found := sort.Find(int(n16.nodeNum), func(i int) int {
+		if n16.keys[i] < c {
 			return 1
 		}
-		if node.keys[i] == c {
+		if n16.keys[i] == c {
 			return 0
 		}
 		return -1
@@ -581,47 +586,47 @@ func (an *artNode) addChild16(a *artAllocator, c byte, child artNode) bool {
 		panic("key already exists")
 	}
 
-	if i < int(node.nodeNum) {
-		copy(node.keys[i+1:node.nodeNum+1], node.keys[i:node.nodeNum])
-		copy(node.children[i+1:node.nodeNum+1], node.children[i:node.nodeNum])
+	if i < int(n16.nodeNum) {
+		copy(n16.keys[i+1:n16.nodeNum+1], n16.keys[i:n16.nodeNum])
+		copy(n16.children[i+1:n16.nodeNum+1], n16.children[i:n16.nodeNum])
 	}
 
-	node.keys[i] = c
-	node.children[i] = child
-	node.nodeNum++
+	n16.keys[i] = c
+	n16.children[i] = child
+	n16.nodeNum++
 	return false
 }
 
 func (an *artNode) addChild48(a *artAllocator, c byte, child artNode) bool {
-	node := an.asNode48(a)
+	n48 := an.asNode48(a)
 
-	if node.nodeNum >= node48cap {
-		an.grow(a)
-		an.addChild(a, c, false, child)
+	if n48.nodeNum >= node48cap {
+		an.growNode48(n48, a)
+		an.addChild256(a, c, child)
 		return true
 	}
 
-	if testing.Testing() && node.present[c>>n48s]&(1<<(c%n48m)) != 0 {
+	if testing.Testing() && n48.present[c>>n48s]&(1<<(c%n48m)) != 0 {
 		panic("key already exists")
 	}
 
-	node.keys[c] = node.nodeNum
-	node.present[c>>n48s] |= 1 << (c % n48m)
-	node.children[node.nodeNum] = child
-	node.nodeNum++
+	n48.keys[c] = n48.nodeNum
+	n48.present[c>>n48s] |= 1 << (c % n48m)
+	n48.children[n48.nodeNum] = child
+	n48.nodeNum++
 	return false
 }
 
 func (an *artNode) addChild256(a *artAllocator, c byte, child artNode) bool {
-	node := an.asNode256(a)
+	n256 := an.asNode256(a)
 
-	if testing.Testing() && node.present[c>>n48s]&(1<<(c%n48m)) != 0 {
+	if testing.Testing() && n256.present[c>>n48s]&(1<<(c%n48m)) != 0 {
 		panic("key already exists")
 	}
 
-	node.present[c>>n48s] |= 1 << (c % n48m)
-	node.children[c] = child
-	node.nodeNum++
+	n256.present[c>>n48s] |= 1 << (c % n48m)
+	n256.children[c] = child
+	n256.nodeNum++
 	return false
 }
 
@@ -632,49 +637,47 @@ func (n *nodeBase) copyMeta(src *nodeBase) {
 	copy(n.prefix[:], src.prefix[:])
 }
 
-func (an *artNode) grow(a *artAllocator) {
-	switch an.kind {
-	case typeNode4:
-		n4 := an.asNode4(a)
-		newAddr, n16 := a.allocNode16()
-		n16.copyMeta(&n4.nodeBase)
+func (an *artNode) growNode4(n4 *node4, a *artAllocator) {
+	newAddr, n16 := a.allocNode16()
+	n16.copyMeta(&n4.nodeBase)
 
-		copy(n16.keys[:], n4.keys[:])
-		copy(n16.children[:], n4.children[:])
+	copy(n16.keys[:], n4.keys[:])
+	copy(n16.children[:], n4.children[:])
 
-		// replace addr and free node4
-		a.freeNode4(an.addr)
-		an.kind = typeNode16
-		an.addr = newAddr
-	case typeNode16:
-		n16 := an.asNode16(a)
-		newAddr, n48 := a.allocNode48()
-		n48.copyMeta(&n16.nodeBase)
-
-		for i := uint8(0); i < n16.nodeBase.nodeNum; i++ {
-			ch := n16.keys[i]
-			n48.keys[ch] = i
-			n48.present[ch>>n48s] |= 1 << (ch % n48m)
-			n48.children[i] = n16.children[i]
-		}
+	// replace addr and free node4
+	a.freeNode4(an.addr)
+	an.kind = typeNode16
+	an.addr = newAddr
+}
 
-		// replace addr and free node16
-		a.freeNode16(an.addr)
-		an.kind = typeNode48
-		an.addr = newAddr
-	case typeNode48:
-		n48 := an.asNode48(a)
-		newAddr, n256 := a.allocNode256()
-		n256.copyMeta(&n48.nodeBase)
+func (an *artNode) growNode16(n16 *node16, a *artAllocator) {
+	newAddr, n48 := a.allocNode48()
+	n48.copyMeta(&n16.nodeBase)
 
-		for i := n48.nextPresentIdx(0); i < node256cap; i = n48.nextPresentIdx(i + 1) {
-			n256.children[i] = n48.children[n48.keys[i]]
-		}
-		copy(n256.present[:], n48.present[:])
+	for i := uint8(0); i < n16.nodeBase.nodeNum; i++ {
+		ch := n16.keys[i]
+		n48.keys[ch] = i
+		n48.present[ch>>n48s] |= 1 << (ch % n48m)
+		n48.children[i] = n16.children[i]
+	}
 
-		// replace addr and free node48
-		a.freeNode48(an.addr)
-		an.kind = typeNode256
-		an.addr = newAddr
+	// replace addr and free node16
+	a.freeNode16(an.addr)
+	an.kind = typeNode48
+	an.addr = newAddr
+}
+
+func (an *artNode) growNode48(n48 *node48, a *artAllocator) {
+	newAddr, n256 := a.allocNode256()
+	n256.copyMeta(&n48.nodeBase)
+
+	for i := n48.nextPresentIdx(0); i < node256cap; i = n48.nextPresentIdx(i + 1) {
+		n256.children[i] = n48.children[n48.keys[i]]
 	}
+	copy(n256.present[:], n48.present[:])
+
+	// replace addr and free node48
+	a.freeNode48(an.addr)
+	an.kind = typeNode256
+	an.addr = newAddr
 }
diff --git a/internal/unionstore/art/art_snapshot.go b/internal/unionstore/art/art_snapshot.go
index 1aab70a..6f470dd 100644
--- a/internal/unionstore/art/art_snapshot.go
+++ b/internal/unionstore/art/art_snapshot.go
@@ -74,7 +74,7 @@ type SnapGetter struct {
 }
 
 func (snap *SnapGetter) Get(ctx context.Context, key []byte) ([]byte, error) {
-	addr, lf := snap.tree.search(key)
+	addr, lf := snap.tree.traverse(key, false)
 	if addr.IsNull() {
 		return nil, tikverr.ErrNotExist
 	}
diff --git a/internal/unionstore/rbt/rbt.go b/internal/unionstore/rbt/rbt.go
index 40b3234..06884fb 100644
--- a/internal/unionstore/rbt/rbt.go
+++ b/internal/unionstore/rbt/rbt.go
@@ -324,6 +324,12 @@ func (db *RBT) Set(key []byte, value []byte, ops ...kv.FlagsOp) error {
 		panic("vlog is reset")
 	}
 
+	if len(key) > MaxKeyLen {
+		return &tikverr.ErrKeyTooLarge{
+			KeySize: len(key),
+		}
+	}
+
 	if value != nil {
 		if size := uint64(len(key) + len(value)); size > db.entrySizeLimit {
 			return &tikverr.ErrEntryTooLarge{
@@ -837,6 +843,8 @@ func (a MemdbNodeAddr) getRight(db *RBT) MemdbNodeAddr {
 	return db.getNode(a.right)
 }
 
+const MaxKeyLen = math.MaxUint16
+
 type memdbNode struct {
 	up    arena.MemdbArenaAddr
 	left  arena.MemdbArenaAddr


Test patch:
diff --git a/internal/locate/region_cache_test.go b/internal/locate/region_cache_test.go
index 8e3bf1a..110a321 100644
--- a/internal/locate/region_cache_test.go
+++ b/internal/locate/region_cache_test.go
@@ -66,7 +66,8 @@ import (
 
 type inspectedPDClient struct {
 	pd.Client
-	getRegion func(ctx context.Context, cli pd.Client, key []byte, opts ...pd.GetRegionOption) (*pd.Region, error)
+	getRegion        func(ctx context.Context, cli pd.Client, key []byte, opts ...pd.GetRegionOption) (*pd.Region, error)
+	batchScanRegions func(ctx context.Context, keyRanges []pd.KeyRange, limit int, opts ...pd.GetRegionOption) ([]*pd.Region, error)
 }
 
 func (c *inspectedPDClient) GetRegion(ctx context.Context, key []byte, opts ...pd.GetRegionOption) (*pd.Region, error) {
@@ -76,6 +77,13 @@ func (c *inspectedPDClient) GetRegion(ctx context.Context, key []byte, opts ...p
 	return c.Client.GetRegion(ctx, key, opts...)
 }
 
+func (c *inspectedPDClient) BatchScanRegions(ctx context.Context, keyRanges []pd.KeyRange, limit int, opts ...pd.GetRegionOption) ([]*pd.Region, error) {
+	if c.batchScanRegions != nil {
+		return c.batchScanRegions(ctx, keyRanges, limit, opts...)
+	}
+	return c.Client.BatchScanRegions(ctx, keyRanges, limit, opts...)
+}
+
 func TestBackgroundRunner(t *testing.T) {
 	t.Run("ShutdownWait", func(t *testing.T) {
 		dur := 100 * time.Millisecond
@@ -459,6 +467,40 @@ func (s *testRegionCacheSuite) TestResolveStateTransition() {
 	s.cluster.AddStore(storeMeta.GetId(), storeMeta.GetAddress(), storeMeta.GetLabels()...)
 }
 
+func (s *testRegionCacheSuite) TestReturnRegionWithNoLeader() {
+	region := s.getRegion([]byte("x"))
+	NoLeaderRegion := &pd.Region{
+		Meta:   region.meta,
+		Leader: nil,
+	}
+
+	originalBatchScanRegions := s.cache.pdClient.BatchScanRegions
+
+	batchScanCnt := 0
+	s.cache.pdClient = &inspectedPDClient{
+		Client: s.cache.pdClient,
+		batchScanRegions: func(ctx context.Context, keyRanges []pd.KeyRange, limit int, opts ...pd.GetRegionOption) ([]*pd.Region, error) {
+			if batchScanCnt == 0 {
+				batchScanCnt++
+				return []*pd.Region{NoLeaderRegion}, nil
+			} else {
+				return originalBatchScanRegions(ctx, keyRanges, limit, opts...)
+			}
+		},
+	}
+
+	bo := retry.NewBackofferWithVars(context.Background(), 1000, nil)
+	returnedRegions, err := s.cache.scanRegions(bo, nil, nil, 100)
+	s.Nil(err)
+	s.Equal(len(returnedRegions), 1)
+	s.Equal(returnedRegions[0].meta.GetId(), region.GetID())
+
+	returnedRegions, err = s.cache.batchScanRegions(bo, []pd.KeyRange{{StartKey: nil, EndKey: nil}}, 100, WithNeedRegionHasLeaderPeer())
+	s.Nil(err)
+	s.Equal(len(returnedRegions), 1)
+	s.Equal(returnedRegions[0].meta.GetId(), region.GetID())
+}
+
 func (s *testRegionCacheSuite) TestNeedExpireRegionAfterTTL() {
 	s.onClosed = func() { SetRegionCacheTTLWithJitter(600, 60) }
 	SetRegionCacheTTLWithJitter(2, 0)
diff --git a/internal/unionstore/art/art_iterator_test.go b/internal/unionstore/art/art_iterator_test.go
index 7dcd504..b21bcce 100644
--- a/internal/unionstore/art/art_iterator_test.go
+++ b/internal/unionstore/art/art_iterator_test.go
@@ -86,7 +86,8 @@ func TestIterSeekLeaf(t *testing.T) {
 			key := []byte{byte(i)}
 			it, err := tree.Iter(key, nil)
 			require.Nil(t, err)
-			idxes, nodes := it.seek(key)
+			it.inner.seek(tree.root, key)
+			idxes, nodes := it.inner.idxes, it.inner.nodes
 			require.Greater(t, len(idxes), 0)
 			require.Equal(t, len(idxes), len(nodes))
 			leafNode := nodes[len(nodes)-1].at(&tree.allocator, idxes[len(idxes)-1])
@@ -198,10 +199,10 @@ func TestSeekInExistNode(t *testing.T) {
 		}
 		require.Equal(t, tree.root.kind, kind)
 		for i := 0; i < cnt-1; i++ {
-			it := &Iterator{
-				tree: tree,
-			}
-			idxes, _ := it.seek([]byte{byte(2*i + 1)})
+			helper := new(baseIter)
+			helper.allocator = &tree.allocator
+			helper.seek(tree.root, []byte{byte(2*i + 1)})
+			idxes := helper.idxes
 			expect := 0
 			switch kind {
 			case typeNode4, typeNode16:
@@ -222,3 +223,147 @@ func TestSeekInExistNode(t *testing.T) {
 	check(New(), typeNode48)
 	check(New(), typeNode256)
 }
+
+func TestSeekToIdx(t *testing.T) {
+	tree := New()
+	check := func(kind nodeKind) {
+		var addr arena.MemdbArenaAddr
+		switch kind {
+		case typeNode4:
+			addr, _ = tree.allocator.allocNode4()
+		case typeNode16:
+			addr, _ = tree.allocator.allocNode16()
+		case typeNode48:
+			addr, _ = tree.allocator.allocNode48()
+		case typeNode256:
+			addr, _ = tree.allocator.allocNode256()
+		}
+		node := artNode{kind: kind, addr: addr}
+		lfAddr, _ := tree.allocator.allocLeaf([]byte{10})
+		lfNode := artNode{kind: typeLeaf, addr: lfAddr}
+		node.addChild(&tree.allocator, 10, false, lfNode)
+
+		var (
+			existIdx int
+			maxIdx   int
+		)
+		switch kind {
+		case typeNode4, typeNode16:
+			existIdx = 0
+			maxIdx = 1
+		case typeNode48, typeNode256:
+			existIdx = 10
+			maxIdx = node256cap
+		}
+
+		nextIdx := seekToIdx(&tree.allocator, node, 1)
+		require.Equal(t, existIdx, nextIdx)
+		nextIdx = seekToIdx(&tree.allocator, node, 11)
+		require.Equal(t, maxIdx, nextIdx)
+	}
+
+	check(typeNode4)
+	check(typeNode16)
+	check(typeNode48)
+	check(typeNode256)
+}
+
+func TestIterateHandle(t *testing.T) {
+	tree := New()
+	h := tree.Staging()
+	require.Nil(t, tree.Set([]byte{1}, []byte{2}))
+	it := tree.IterWithFlags(nil, nil)
+	handle := it.Handle()
+
+	require.Equal(t, tree.GetKeyByHandle(handle), []byte{1})
+	val, valid := tree.GetValueByHandle(handle)
+	require.True(t, valid)
+	require.Equal(t, val, []byte{2})
+
+	tree.Cleanup(h)
+	require.Equal(t, tree.GetKeyByHandle(handle), []byte{1})
+	_, valid = tree.GetValueByHandle(handle)
+	require.False(t, valid)
+}
+
+func TestSeekPrefixMismatch(t *testing.T) {
+	tree := New()
+
+	shortPrefix := make([]byte, 10)
+	longPrefix := make([]byte, 30)
+	for i := 0; i < len(shortPrefix); i++ {
+		shortPrefix[i] = 1
+	}
+	for i := 0; i < len(longPrefix); i++ {
+		longPrefix[i] = 2
+	}
+
+	keys := [][]byte{
+		append(shortPrefix, 1),
+		append(shortPrefix, 2),
+		append(longPrefix, 3),
+		append(longPrefix, 4),
+	}
+	for _, key := range keys {
+		require.Nil(t, tree.Set(key, key))
+	}
+
+	it, err := tree.Iter(append(shortPrefix[:len(shortPrefix)-1], 0), append(longPrefix[:len(longPrefix)-1], 3))
+	require.Nil(t, err)
+	for _, key := range keys {
+		require.True(t, it.Valid())
+		require.Equal(t, it.Key(), key)
+		require.Equal(t, it.Value(), key)
+		require.Nil(t, it.Next())
+	}
+	require.False(t, it.Valid())
+}
+
+func TestIterPositionCompare(t *testing.T) {
+	compare := func(idx1, idx2 []int) int {
+		helper1, helper2 := new(baseIter), new(baseIter)
+		helper1.idxes, helper2.idxes = idx1, idx2
+		return helper1.compare(helper2)
+	}
+
+	require.Equal(t, compare([]int{1, 2, 3}, []int{1, 2, 3}), 0)
+	require.Equal(t, compare([]int{1, 2, 2}, []int{1, 2, 3}), -1)
+	require.Equal(t, compare([]int{1, 2, 4}, []int{1, 2, 3}), 1)
+	require.Equal(t, compare([]int{1, 2, 3}, []int{1, 2}), 1)
+	require.Equal(t, compare([]int{1, 2}, []int{1, 2, 3}), -1)
+}
+
+func TestIterSeekNoResult(t *testing.T) {
+	check := func(kind nodeKind) {
+		var child int
+		switch kind {
+		case typeNode4:
+			child = 0
+		case typeNode16:
+			child = node4cap + 1
+		case typeNode48:
+			child = node16cap + 1
+		case typeNode256:
+			child = node48cap + 1
+		}
+		// let the node expand to target kind
+		tree := New()
+		for i := 0; i < child; i++ {
+			require.Nil(t, tree.Set([]byte{1, byte(i)}, []byte{1, byte(i)}))
+		}
+
+		require.Nil(t, tree.Set([]byte{1, 100}, []byte{1, 100}))
+		require.Nil(t, tree.Set([]byte{1, 200}, []byte{1, 200}))
+		it, err := tree.Iter([]byte{1, 100, 1}, []byte{1, 200})
+		require.Nil(t, err)
+		require.False(t, it.Valid())
+		it, err = tree.IterReverse([]byte{1, 200}, []byte{1, 100, 1})
+		require.Nil(t, err)
+		require.False(t, it.Valid())
+	}
+
+	check(typeNode4)
+	check(typeNode16)
+	check(typeNode48)
+	check(typeNode256)
+}
diff --git a/internal/unionstore/art/art_node_test.go b/internal/unionstore/art/art_node_test.go
index da4ad5d..c9ff55f 100644
--- a/internal/unionstore/art/art_node_test.go
+++ b/internal/unionstore/art/art_node_test.go
@@ -265,3 +265,153 @@ func TestLCP(t *testing.T) {
 		require.Equal(t, uint32(5-i), longestCommonPrefix(k1, k2, uint32(i)))
 	}
 }
+
+func TestNodeAddChild(t *testing.T) {
+	var allocator artAllocator
+	allocator.init()
+
+	check := func(n artNode) {
+		require.Equal(t, uint8(0), n.asNode(&allocator).nodeNum)
+		lfAddr, _ := allocator.allocLeaf([]byte{1, 2, 3, 4, 5})
+		lfNode := artNode{kind: typeLeaf, addr: lfAddr}
+		n.addChild(&allocator, 1, false, lfNode)
+		require.Equal(t, uint8(1), n.asNode(&allocator).nodeNum)
+		require.Panics(t, func() {
+			// addChild should panic if the key is existed.
+			n.addChild(&allocator, 1, false, lfNode)
+		})
+		n.addChild(&allocator, 2, false, lfNode)
+		require.Equal(t, uint8(2), n.asNode(&allocator).nodeNum)
+		// inplace leaf won't be counted in nodeNum
+		n.addChild(&allocator, 0, true, lfNode)
+		require.Equal(t, uint8(2), n.asNode(&allocator).nodeNum)
+	}
+
+	addr, _ := allocator.allocNode4()
+	check(artNode{kind: typeNode4, addr: addr})
+	addr, _ = allocator.allocNode16()
+	check(artNode{kind: typeNode16, addr: addr})
+	addr, _ = allocator.allocNode48()
+	check(artNode{kind: typeNode48, addr: addr})
+	addr, _ = allocator.allocNode256()
+	check(artNode{kind: typeNode256, addr: addr})
+}
+
+func TestNodeGrow(t *testing.T) {
+	var allocator artAllocator
+	allocator.init()
+
+	check := func(n artNode) {
+		capacities := map[nodeKind]int{
+			typeNode4:  node4cap,
+			typeNode16: node16cap,
+			typeNode48: node48cap,
+		}
+		growTypes := map[nodeKind]nodeKind{
+			typeNode4:  typeNode16,
+			typeNode16: typeNode48,
+			typeNode48: typeNode256,
+		}
+
+		capacity, ok := capacities[n.kind]
+		require.True(t, ok)
+		beforeKind := n.kind
+		afterKind, ok := growTypes[n.kind]
+		require.True(t, ok)
+
+		for i := 0; i < capacity; i++ {
+			lfAddr, _ := allocator.allocLeaf([]byte{byte(i)})
+			lfNode := artNode{kind: typeLeaf, addr: lfAddr}
+			n.addChild(&allocator, byte(i), false, lfNode)
+			require.Equal(t, beforeKind, n.kind)
+		}
+		lfAddr, _ := allocator.allocLeaf([]byte{byte(capacity)})
+		lfNode := artNode{kind: typeLeaf, addr: lfAddr}
+		n.addChild(&allocator, byte(capacity), false, lfNode)
+		require.Equal(t, afterKind, n.kind)
+	}
+
+	addr, _ := allocator.allocNode4()
+	check(artNode{kind: typeNode4, addr: addr})
+	addr, _ = allocator.allocNode16()
+	check(artNode{kind: typeNode16, addr: addr})
+	addr, _ = allocator.allocNode48()
+	check(artNode{kind: typeNode48, addr: addr})
+}
+
+func TestReplaceChild(t *testing.T) {
+	var allocator artAllocator
+	allocator.init()
+
+	check := func(n artNode) {
+		require.Equal(t, uint8(0), n.asNode(&allocator).nodeNum)
+		lfAddr, _ := allocator.allocLeaf([]byte{1, 2, 3, 4, 5})
+		lfNode := artNode{kind: typeLeaf, addr: lfAddr}
+		n.addChild(&allocator, 1, false, lfNode)
+		require.Equal(t, uint8(1), n.asNode(&allocator).nodeNum)
+		newLfAddr, _ := allocator.allocLeaf([]byte{1, 2, 3, 4, 4})
+		newLfNode := artNode{kind: typeLeaf, addr: newLfAddr}
+		require.Panics(t, func() {
+			// replaceChild should panic if the key is not existed.
+			n.replaceChild(&allocator, 2, newLfNode)
+		})
+		n.replaceChild(&allocator, 1, newLfNode)
+		require.Equal(t, uint8(1), n.asNode(&allocator).nodeNum)
+		_, childLf := n.findChild(&allocator, 1, false)
+		require.NotEqual(t, childLf.addr, lfAddr)
+		require.Equal(t, childLf.addr, newLfAddr)
+	}
+
+	addr, _ := allocator.allocNode4()
+	check(artNode{kind: typeNode4, addr: addr})
+	addr, _ = allocator.allocNode16()
+	check(artNode{kind: typeNode16, addr: addr})
+	addr, _ = allocator.allocNode48()
+	check(artNode{kind: typeNode48, addr: addr})
+	addr, _ = allocator.allocNode256()
+	check(artNode{kind: typeNode256, addr: addr})
+
+}
+
+func TestMinimumNode(t *testing.T) {
+	var allocator artAllocator
+	allocator.init()
+
+	check := func(kind nodeKind) {
+		var addr arena.MemdbArenaAddr
+
+		switch kind {
+		case typeNode4:
+			addr, _ = allocator.allocNode4()
+		case typeNode16:
+			addr, _ = allocator.allocNode16()
+		case typeNode48:
+			addr, _ = allocator.allocNode48()
+		case typeNode256:
+			addr, _ = allocator.allocNode256()
+		}
+
+		node := artNode{kind: kind, addr: addr}
+
+		for _, char := range []byte{255, 127, 63, 0} {
+			lfAddr, _ := allocator.allocLeaf([]byte{char})
+			lfNode := artNode{kind: typeLeaf, addr: lfAddr}
+			node.addChild(&allocator, char, false, lfNode)
+			minNode := minimum(&allocator, node)
+			require.Equal(t, typeLeaf, minNode.kind)
+			require.Equal(t, lfAddr, minNode.addr)
+		}
+
+		lfAddr, _ := allocator.allocLeaf([]byte{0})
+		lfNode := artNode{kind: typeLeaf, addr: lfAddr}
+		node.addChild(&allocator, 0, true, lfNode)
+		minNode := minimum(&allocator, node)
+		require.Equal(t, typeLeaf, minNode.kind)
+		require.Equal(t, lfAddr, minNode.addr)
+	}
+
+	check(typeNode4)
+	check(typeNode16)
+	check(typeNode48)
+	check(typeNode256)
+}
diff --git a/internal/unionstore/art/art_test.go b/internal/unionstore/art/art_test.go
index f5dc0b7..b5f4148 100644
--- a/internal/unionstore/art/art_test.go
+++ b/internal/unionstore/art/art_test.go
@@ -109,21 +109,21 @@ func TestFlag(t *testing.T) {
 	require.Nil(t, err)
 	require.True(t, flags.HasLocked())
 	// iterate can also see the flags
-	//it, err := tree.Iter(nil, nil)
-	//require.Nil(t, err)
-	//require.True(t, it.Valid())
-	//require.Equal(t, it.Key(), []byte{0})
-	//require.Equal(t, it.Value(), []byte{0})
-	//require.True(t, it.Flags().HasPresumeKeyNotExists())
-	//require.False(t, it.Flags().HasLocked())
-	//require.Nil(t, it.Next())
-	//require.True(t, it.Valid())
-	//require.Equal(t, it.Key(), []byte{1})
-	//require.Equal(t, it.Value(), []byte{1})
-	//require.True(t, it.Flags().HasLocked())
-	//require.False(t, it.Flags().HasPresumeKeyNotExists())
-	//require.Nil(t, it.Next())
-	//require.False(t, it.Valid())
+	it, err := tree.Iter(nil, nil)
+	require.Nil(t, err)
+	require.True(t, it.Valid())
+	require.Equal(t, it.Key(), []byte{0})
+	require.Equal(t, it.Value(), []byte{0})
+	require.True(t, it.Flags().HasPresumeKeyNotExists())
+	require.False(t, it.Flags().HasLocked())
+	require.Nil(t, it.Next())
+	require.True(t, it.Valid())
+	require.Equal(t, it.Key(), []byte{1})
+	require.Equal(t, it.Value(), []byte{1})
+	require.True(t, it.Flags().HasLocked())
+	require.False(t, it.Flags().HasPresumeKeyNotExists())
+	require.Nil(t, it.Next())
+	require.False(t, it.Valid())
 }
 
 func TestLongPrefix1(t *testing.T) {
@@ -184,6 +184,14 @@ func TestFlagOnlyKey(t *testing.T) {
 	require.Error(t, err)
 }
 
+func TestSearchPrefixMisatch(t *testing.T) {
+	tree := New()
+	tree.Set([]byte{1, 1, 1, 1, 1, 1}, []byte{1, 1, 1, 1, 1, 1})
+	tree.Set([]byte{1, 1, 1, 1, 1, 2}, []byte{1, 1, 1, 1, 1, 2})
+	_, err := tree.Get([]byte{1, 1, 1, 3, 1, 1})
+	require.NotNil(t, err)
+}
+
 func TestSearchOptimisticMismatch(t *testing.T) {
 	tree := New()
 	prefix := make([]byte, 22)
@@ -230,3 +238,27 @@ func TestExpansion(t *testing.T) {
 	require.Equal(t, n4.keys[:2], []byte{1, 255})
 	require.Equal(t, n4.children[1].asLeaf(&tree.allocator).GetKey(), append(prefix, []byte{1, 255, 2}...))
 }
+
+func TestDiscardValues(t *testing.T) {
+	tree := New()
+	tree.Set([]byte{1}, []byte{2})
+	it := tree.IterWithFlags(nil, nil)
+	handle := it.Handle()
+	val, exist := tree.GetValueByHandle(handle)
+	require.Equal(t, val, []byte{2})
+	require.True(t, exist)
+	key := tree.GetKeyByHandle(handle)
+	require.Equal(t, key, []byte{1})
+
+	tree.DiscardValues()
+	_, exist = tree.GetValueByHandle(handle)
+	require.False(t, exist)
+	key = tree.GetKeyByHandle(handle)
+	require.Equal(t, key, []byte{1})
+	require.Panics(t, func() {
+		tree.Get([]byte{3})
+	})
+	require.Panics(t, func() {
+		tree.Set([]byte{3}, []byte{4})
+	})
+}
diff --git a/internal/unionstore/memdb_test.go b/internal/unionstore/memdb_test.go
index 2786de1..6003059 100644
--- a/internal/unionstore/memdb_test.go
+++ b/internal/unionstore/memdb_test.go
@@ -37,14 +37,19 @@
 package unionstore
 
 import (
+	"bytes"
 	"context"
 	"encoding/binary"
 	"fmt"
+	"math"
+	"strconv"
+	"strings"
 	"testing"
 
 	leveldb "github.com/pingcap/goleveldb/leveldb/memdb"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	tikverr "github.com/tikv/client-go/v2/error"
 	"github.com/tikv/client-go/v2/kv"
 )
 
@@ -392,6 +397,8 @@ func testReset(t *testing.T, db interface {
 	db.Reset()
 	_, err := db.Get(context.Background(), []byte{0, 0, 0, 0})
 	assert.NotNil(err)
+	_, err = db.GetFlags([]byte{0, 0, 0, 0})
+	assert.NotNil(err)
 	it, _ := db.Iter(nil, nil)
 	assert.False(it.Valid())
 }
@@ -495,6 +502,9 @@ func testDirty(t *testing.T, createDb func() MemBuffer) {
 func TestFlags(t *testing.T) {
 	testFlags(t, newRbtDBWithContext(), func(db MemBuffer) Iterator { return db.(*rbtDBWithContext).IterWithFlags(nil, nil) })
 	testFlags(t, newArtDBWithContext(), func(db MemBuffer) Iterator { return db.(*artDBWithContext).IterWithFlags(nil, nil) })
+
+	testFlags(t, newRbtDBWithContext(), func(db MemBuffer) Iterator { return db.(*rbtDBWithContext).IterReverseWithFlags(nil) })
+	testFlags(t, newArtDBWithContext(), func(db MemBuffer) Iterator { return db.(*artDBWithContext).IterReverseWithFlags(nil) })
 }
 
 func testFlags(t *testing.T, db MemBuffer, iterWithFlags func(db MemBuffer) Iterator) {
@@ -538,6 +548,10 @@ func testFlags(t *testing.T, db MemBuffer, iterWithFlags func(db MemBuffer) Iter
 	for ; it.Valid(); it.Next() {
 		k := binary.BigEndian.Uint32(it.Key())
 		assert.True(k%2 == 0)
+		hasValue := it.(interface {
+			HasValue() bool
+		}).HasValue()
+		assert.False(hasValue)
 	}
 
 	for i := uint32(0); i < cnt; i++ {
@@ -1046,6 +1060,49 @@ func testSnapshotGetIter(t *testing.T, db MemBuffer) {
 		assert.Equal(reverseIter.Key(), []byte{byte(1)})
 		assert.Equal(reverseIter.Value(), []byte{byte(50)})
 	}
+
+	db.(interface {
+		Reset()
+	}).Reset()
+	db.UpdateFlags([]byte{255}, kv.SetPresumeKeyNotExists)
+	// set (2, 2) ... (100, 100) in snapshot
+	for i := 1; i < 50; i++ {
+		db.Set([]byte{byte(2 * i)}, []byte{byte(2 * i)})
+	}
+	h := db.Staging()
+	// set (0, 0) (1, 2) (2, 4) ... (100, 200) in staging
+	for i := 0; i < 100; i++ {
+		db.Set([]byte{byte(i)}, []byte{byte(2 * i)})
+	}
+
+	snapGetter := db.SnapshotGetter()
+	v, err := snapGetter.Get(context.Background(), []byte{byte(2)})
+	assert.Nil(err)
+	assert.Equal(v, []byte{byte(2)})
+	_, err = snapGetter.Get(context.Background(), []byte{byte(1)})
+	assert.NotNil(err)
+	_, err = snapGetter.Get(context.Background(), []byte{byte(254)})
+	assert.NotNil(err)
+	_, err = snapGetter.Get(context.Background(), []byte{byte(255)})
+	assert.NotNil(err)
+
+	it := db.SnapshotIter(nil, nil)
+	// snapshot iter only see the snapshot data
+	for i := 1; i < 50; i++ {
+		assert.Equal(it.Key(), []byte{byte(2 * i)})
+		assert.Equal(it.Value(), []byte{byte(2 * i)})
+		assert.True(it.Valid())
+		it.Next()
+	}
+	it = db.SnapshotIterReverse(nil, nil)
+	for i := 49; i >= 1; i-- {
+		assert.Equal(it.Key(), []byte{byte(2 * i)})
+		assert.Equal(it.Value(), []byte{byte(2 * i)})
+		assert.True(it.Valid())
+		it.Next()
+	}
+	assert.False(it.Valid())
+	db.Release(h)
 }
 
 func TestCleanupKeepPersistentFlag(t *testing.T) {
@@ -1172,3 +1229,101 @@ func testMemBufferCache(t *testing.T, buffer MemBuffer) {
 		assert.Equal(v, []byte{2, 2})
 	})
 }
+
+func TestMemDBLeafFragmentation(t *testing.T) {
+	// RBT cannot pass the leaf fragmentation test.
+	testMemDBLeafFragmentation(t, newArtDBWithContext())
+}
+
+func testMemDBLeafFragmentation(t *testing.T, buffer MemBuffer) {
+	assert := assert.New(t)
+	h := buffer.Staging()
+	mem := buffer.Mem()
+	for i := 0; i < 10; i++ {
+		for k := 0; k < 100; k++ {
+			buffer.Set([]byte(strings.Repeat(strconv.Itoa(k), 256)), []byte("value"))
+		}
+		cur := buffer.Mem()
+		if mem == 0 {
+			mem = cur
+		} else {
+			assert.LessOrEqual(cur, mem)
+		}
+		buffer.Cleanup(h)
+		h = buffer.Staging()
+	}
+}
+
+func TestReadOnlyZeroMem(t *testing.T) {
+	// read only MemBuffer should not allocate heap memory.
+	assert.Zero(t, newRbtDBWithContext().Mem())
+	assert.Zero(t, newArtDBWithContext().Mem())
+}
+
+func TestKeyValueOversize(t *testing.T) {
+	check := func(t *testing.T, db MemBuffer) {
+		key := make([]byte, math.MaxUint16)
+		overSizeKey := make([]byte, math.MaxUint16+1)
+
+		assert.Nil(t, db.Set(key, overSizeKey))
+		err := db.Set(overSizeKey, key)
+		assert.NotNil(t, err)
+		assert.Equal(t, err.(*tikverr.ErrKeyTooLarge).KeySize, math.MaxUint16+1)
+	}
+
+	check(t, newRbtDBWithContext())
+	check(t, newArtDBWithContext())
+}
+
+func TestSetMemoryFootprintChangeHook(t *testing.T) {
+	check := func(t *testing.T, db MemBuffer) {
+		memoryConsumed := uint64(0)
+		assert.False(t, db.MemHookSet())
+		db.SetMemoryFootprintChangeHook(func(mem uint64) {
+			memoryConsumed = mem
+		})
+		assert.True(t, db.MemHookSet())
+
+		assert.Zero(t, memoryConsumed)
+		db.Set([]byte{1}, []byte{1})
+		assert.NotZero(t, memoryConsumed)
+	}
+
+	check(t, newRbtDBWithContext())
+	check(t, newArtDBWithContext())
+}
+
+func TestSelectValueHistory(t *testing.T) {
+	check := func(t *testing.T, db interface {
+		MemBuffer
+		SelectValueHistory(key []byte, predicate func(value []byte) bool) ([]byte, error)
+	}) {
+		db.Set([]byte{1}, []byte{1})
+		h := db.Staging()
+		db.Set([]byte{1}, []byte{1, 1})
+
+		val, err := db.SelectValueHistory([]byte{1}, func(value []byte) bool { return bytes.Equal(value, []byte{1}) })
+		assert.Nil(t, err)
+		assert.Equal(t, val, []byte{1})
+		val, err = db.SelectValueHistory([]byte{1}, func(value []byte) bool { return bytes.Equal(value, []byte{1, 1}) })
+		assert.Nil(t, err)
+		assert.Equal(t, val, []byte{1, 1})
+		val, err = db.SelectValueHistory([]byte{1}, func(value []byte) bool { return bytes.Equal(value, []byte{1, 1, 1}) })
+		assert.Nil(t, err)
+		assert.Nil(t, val)
+		_, err = db.SelectValueHistory([]byte{2}, func([]byte) bool { return false })
+		assert.NotNil(t, err)
+
+		db.Cleanup(h)
+
+		val, err = db.SelectValueHistory([]byte{1}, func(value []byte) bool { return bytes.Equal(value, []byte{1}) })
+		assert.Nil(t, err)
+		assert.Equal(t, val, []byte{1})
+		val, err = db.SelectValueHistory([]byte{1}, func(value []byte) bool { return bytes.Equal(value, []byte{1, 1}) })
+		assert.Nil(t, err)
+		assert.Nil(t, val)
+	}
+
+	check(t, newRbtDBWithContext())
+	check(t, newArtDBWithContext())
+}