1616package art
1717
1818import (
19+ "bytes"
1920 "fmt"
2021 "math"
22+ "sync/atomic"
2123
2224 tikverr "github.com/tikv/client-go/v2/error"
2325 "github.com/tikv/client-go/v2/internal/unionstore/arena"
@@ -44,6 +46,12 @@ type ART struct {
4446 bufferSizeLimit uint64
4547 len int
4648 size int
49+
50+ // The lastTraversedNode stores addr in uint64 of the last traversed node, include search and recursiveInsert.
51+ // Compare to atomic.Pointer, atomic.Uint64 can avoid heap allocation, so it's more efficient.
52+ lastTraversedNode atomic.Uint64
53+ hitCount atomic.Uint64
54+ missCount atomic.Uint64
4755}
4856
4957func New () * ART {
@@ -55,6 +63,7 @@ func New() *ART {
5563 t .allocator .nodeAllocator .freeNode4 = make ([]arena.MemdbArenaAddr , 0 , 1 << 4 )
5664 t .allocator .nodeAllocator .freeNode16 = make ([]arena.MemdbArenaAddr , 0 , 1 << 3 )
5765 t .allocator .nodeAllocator .freeNode48 = make ([]arena.MemdbArenaAddr , 0 , 1 << 2 )
66+ t .lastTraversedNode .Store (arena .NullU64Addr )
5867 return & t
5968}
6069
@@ -102,10 +111,26 @@ func (t *ART) Set(key artKey, value []byte, ops ...kv.FlagsOp) error {
102111 return nil
103112}
104113
105- // search looks up the leaf with the given key.
114+ // search wraps searchImpl with cache.
115+ func (t * ART ) search (key artKey ) (arena.MemdbArenaAddr , * artLeaf ) {
116+ // check cache
117+ addr , leaf , found := t .checkKeyInCache (key )
118+ if found {
119+ t .hitCount .Add (1 )
120+ return addr , leaf
121+ }
122+ t .missCount .Add (1 )
123+ addr , leaf = t .searchImpl (key )
124+ if ! addr .IsNull () {
125+ t .updateLastTraversed (addr )
126+ }
127+ return addr , leaf
128+ }
129+
130+ // searchImpl looks up the leaf with the given key.
106131// It returns the memory arena address and leaf itself it there is a match leaf,
107132// returns arena.NullAddr and nil if the key is not found.
108- func (t * ART ) search (key artKey ) (arena.MemdbArenaAddr , * artLeaf ) {
133+ func (t * ART ) searchImpl (key artKey ) (arena.MemdbArenaAddr , * artLeaf ) {
109134 current := t .root
110135 if current == nullArtNode {
111136 return arena .NullAddr , nil
@@ -154,9 +179,25 @@ func (t *ART) search(key artKey) (arena.MemdbArenaAddr, *artLeaf) {
154179 }
155180}
156181
157- // recursiveInsert returns the node address of the key.
158- // It will insert the key if not exists, returns the newly inserted or existing leaf.
182+ // recursiveInsert wraps recursiveInsertImpl with cache.
159183func (t * ART ) recursiveInsert (key artKey ) (arena.MemdbArenaAddr , * artLeaf ) {
184+ addr , leaf , found := t .checkKeyInCache (key )
185+ if found {
186+ t .hitCount .Add (1 )
187+ return addr , leaf
188+ }
189+ t .missCount .Add (1 )
190+ addr , leaf = t .recursiveInsertImpl (key )
191+ if ! addr .IsNull () {
192+ t .updateLastTraversed (addr )
193+ }
194+ return addr , leaf
195+ }
196+
197+ // recursiveInsertImpl returns the node address of the key.
198+ // It will insert the key if not exists, returns the newly inserted or existing leaf.
199+ func (t * ART ) recursiveInsertImpl (key artKey ) (arena.MemdbArenaAddr , * artLeaf ) {
200+
160201 // lazy init root node and allocator.
161202 // this saves memory for read only txns.
162203 if t .root .addr .IsNull () {
@@ -501,6 +542,7 @@ func (t *ART) Reset() {
501542 t .len = 0
502543 t .allocator .nodeAllocator .Reset ()
503544 t .allocator .vlogAllocator .Reset ()
545+ t .lastTraversedNode .Store (arena .NullU64Addr )
504546}
505547
506548// DiscardValues releases the memory used by all values.
@@ -583,10 +625,31 @@ func (t *ART) RemoveFromBuffer(key []byte) {
583625 panic ("unimplemented" )
584626}
585627
628+ // updateLastTraversed updates the last traversed node atomically
629+ // the addr must be a valid leaf address
630+ func (t * ART ) updateLastTraversed (addr arena.MemdbArenaAddr ) {
631+ t .lastTraversedNode .Store (addr .AsU64 ())
632+ }
633+
634+ // checkKeyInCache retrieves the last traversed node if the key matches
635+ func (t * ART ) checkKeyInCache (key []byte ) (arena.MemdbArenaAddr , * artLeaf , bool ) {
636+ addrU64 := t .lastTraversedNode .Load ()
637+ if addrU64 == arena .NullU64Addr {
638+ return arena .NullAddr , nil , false
639+ }
640+
641+ addr := arena .U64ToAddr (addrU64 )
642+ leaf := t .allocator .getLeaf (addr )
643+ if ! bytes .Equal (leaf .GetKey (), key ) {
644+ return arena .NullAddr , nil , false
645+ }
646+ return addr , leaf , true
647+ }
648+
586649func (t * ART ) GetCacheHitCount () uint64 {
587- return 0
650+ return t . hitCount . Load ()
588651}
589652
590653func (t * ART ) GetCacheMissCount () uint64 {
591- return 0
654+ return t . missCount . Load ()
592655}
0 commit comments