From 3044d3484e305fe0fe156da872f338f03c5cd763 Mon Sep 17 00:00:00 2001 From: you06 Date: Mon, 26 Aug 2024 16:36:29 +0900 Subject: [PATCH 1/7] implement ART with basic set & get Signed-off-by: you06 --- internal/unionstore/art/art.go | 321 ++++++++++- internal/unionstore/art/art_arena.go | 122 +++++ internal/unionstore/art/art_node.go | 646 ++++++++++++++++++++++- internal/unionstore/art/art_node_test.go | 261 +++++++++ internal/unionstore/art/art_test.go | 191 +++++++ internal/unionstore/memdb_art.go | 2 +- internal/unionstore/memdb_test.go | 7 +- 7 files changed, 1526 insertions(+), 24 deletions(-) create mode 100644 internal/unionstore/art/art_node_test.go create mode 100644 internal/unionstore/art/art_test.go diff --git a/internal/unionstore/art/art.go b/internal/unionstore/art/art.go index 4ed89e5f3..3bfdd6252 100644 --- a/internal/unionstore/art/art.go +++ b/internal/unionstore/art/art.go @@ -18,6 +18,7 @@ package art import ( "math" + tikverr "github.com/tikv/client-go/v2/error" "github.com/tikv/client-go/v2/internal/unionstore/arena" "github.com/tikv/client-go/v2/kv" ) @@ -47,20 +48,310 @@ func New() *ART { } func (t *ART) Get(key []byte) ([]byte, error) { - panic("unimplemented") + _, leaf := t.search(key) + if leaf == nil || leaf.vAddr.IsNull() { + return nil, tikverr.ErrNotExist + } + return t.allocator.vlogAllocator.GetValue(leaf.vAddr), nil } // GetFlags returns the latest flags associated with key. func (t *ART) GetFlags(key []byte) (kv.KeyFlags, error) { - panic("unimplemented") + _, leaf := t.search(key) + if leaf == nil { + return 0, tikverr.ErrNotExist + } + if leaf.vAddr.IsNull() && leaf.isDeleted() { + return 0, tikverr.ErrNotExist + } + return leaf.getKeyFlags(), nil } -func (t *ART) Set(key artKey, value []byte, ops []kv.FlagsOp) error { - panic("unimplemented") +func (t *ART) Set(key artKey, value []byte, ops ...kv.FlagsOp) error { + if value != nil { + if size := uint64(len(key) + len(value)); size > t.entrySizeLimit { + return &tikverr.ErrEntryTooLarge{ + Limit: t.entrySizeLimit, + Size: size, + } + } + } + if len(t.stages) == 0 { + t.dirty = true + } + addr, leaf := t.recursiveInsert(key) + t.setValue(addr, leaf, value, ops) + if uint64(t.Size()) > t.bufferSizeLimit { + return &tikverr.ErrTxnTooLarge{Size: t.Size()} + } + return nil } func (t *ART) search(key artKey) (arena.MemdbArenaAddr, *artLeaf) { - panic("unimplemented") + current := t.root + if current == nullArtNode { + return arena.NullAddr, nil + } + depth := uint32(0) + var node *nodeBase + for { + if current.isLeaf() { + lf := current.leaf(&t.allocator) + if lf.match(0, key) { + return current.addr, lf + } + return arena.NullAddr, nil + } + + // inline: performance critical path + // get the basic node information. + switch current.kind { + case typeNode4: + node = ¤t.node4(&t.allocator).nodeBase + case typeNode16: + node = ¤t.node16(&t.allocator).nodeBase + case typeNode48: + node = ¤t.node48(&t.allocator).nodeBase + case typeNode256: + node = ¤t.node256(&t.allocator).nodeBase + default: + panic("invalid nodeBase kind") + } + + if node.prefixLen > 0 { + prefixLen := node.match(key, depth) + if prefixLen < min(node.prefixLen, maxPrefixLen) { + return arena.NullAddr, nil + } + // If node.prefixLen > maxPrefixLen, we optimistically match the prefix here. + // False positive is possible, but it's fine since we will check the full artLeaf key at last. + depth += node.prefixLen + } + + _, current = current.findChild(&t.allocator, key.charAt(int(depth)), key.valid(int(depth))) + if current.addr.IsNull() { + return arena.NullAddr, nil + } + depth++ + } +} + +// recursiveInsert returns the node address of the key. +// if insert is true, it will insert the key if not exists, unless nullAddr is returned. +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() { + t.root, _ = t.newNode4() + } + + depth := uint32(0) + prevDepth := 0 + prev := nullArtNode + current := t.root + var node *nodeBase + for { + if current.isLeaf() { + return t.expandLeaf(key, depth, prev, current) + } + + // inline: performance critical path + // get the basic node information. + switch current.kind { + case typeNode4: + node = ¤t.node4(&t.allocator).nodeBase + case typeNode16: + node = ¤t.node16(&t.allocator).nodeBase + case typeNode48: + node = ¤t.node48(&t.allocator).nodeBase + case typeNode256: + node = ¤t.node256(&t.allocator).nodeBase + default: + panic("invalid nodeBase kind") + } + + if node.prefixLen > 0 { + mismatchIdx := current.matchDeep(&t.allocator, key, depth) + if mismatchIdx < node.prefixLen { + // if the prefix doesn't match, we split the node into different prefixes. + return t.expandNode(key, depth, mismatchIdx, prev, current, node) + } + depth += node.prefixLen + } + + // search next node + valid := key.valid(int(depth)) + _, next := current.findChild(&t.allocator, key.charAt(int(depth)), valid) + if next == nullArtNode { + // insert as leaf if there is no child. + newLeaf, lf := t.newLeaf(key) + if current.addChild(&t.allocator, key.charAt(int(depth)), !key.valid(int(depth)), newLeaf) { + if prev == nullArtNode { + t.root = current + } else { + prev.swapChild(&t.allocator, key.charAt(prevDepth), current) + } + } + return newLeaf.addr, lf + } + if !valid && next.kind == typeLeaf { + // key is drained, return the leaf. + return next.addr, next.leaf(&t.allocator) + } + prev = current + current = next + prevDepth = int(depth) + depth++ + continue + } +} + +// expandLeaf expands the exist artLeaf to a node4 if the keys are different. +// it returns the addr and leaf of the given key. +func (t *ART) expandLeaf(key artKey, depth uint32, prev, current artNode) (arena.MemdbArenaAddr, *artLeaf) { + // Expand the artLeaf to a node4. + // + // ┌────────────┐ + // │ new │ + // │ node4 │ + // ┌─────────┐ └──────┬─────┘ + // │ old │ ---> │ + // │ leaf1 │ ┌────────┴────────┐ + // └─────────┘ │ │ + // ┌────▼────┐ ┌────▼────┐ + // │ old │ │ new │ + // │ leaf1 │ │ leaf2 │ + // └─────────┘ └─────────┘ + leaf1 := current.leaf(&t.allocator) + if leaf1.match(depth-1, key) { + // same key, return the artLeaf and overwrite the value. + return current.addr, leaf1 + } + prevDepth := int(depth - 1) + + leaf2Addr, leaf2 := t.newLeaf(key) + l1Key, l2Key := artKey(leaf1.GetKey()), artKey(leaf2.GetKey()) + lcp := longestCommonPrefix(l1Key, l2Key, depth) + + // calculate the common prefix length of new node. + an, n4 := t.newNode4() + n4.setPrefix(key[depth:], lcp) + depth += lcp + an.addChild(&t.allocator, l1Key.charAt(int(depth)), !l1Key.valid(int(depth)), current) + an.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 = an + } else { + prev.swapChild(&t.allocator, key.charAt(prevDepth), an) + } + return leaf2Addr.addr, leaf2 +} + +func (t *ART) expandNode(key artKey, depth, mismatchIdx uint32, prev, current artNode, currNode *nodeBase) (arena.MemdbArenaAddr, *artLeaf) { + // prefix mismatch, create a new parent node which has a shorter prefix. + // example of insert "acc" into node with "abc prefix: + // ┌────────────┐ + // │ new node4 │ + // │ prefix: a │ + // └──────┬─────┘ + // ┌─────────────┐ ┌── b ───┴── c ───┐ + // │ node4 │ ---> │ │ + // │ prefix: abc │ ┌──────▼─────┐ ┌──────▼─────┐ + // └─────────────┘ │ old node4 │ │ new leaf │ + // │ prefix: c │ │ key: acc │ + // └────────────┘ └────────────┘ + prevDepth := int(depth - 1) + + // set prefix for new node. + newArtNode, newN4 := t.newNode4() + newN4.setPrefix(key[depth:], mismatchIdx) + + // update prefix for old node and move it as a child of the new node. + if currNode.prefixLen <= maxPrefixLen { + nodeKey := currNode.prefix[mismatchIdx] + currNode.prefixLen -= mismatchIdx + 1 + copy(currNode.prefix[:], currNode.prefix[mismatchIdx+1:]) + newArtNode.addChild(&t.allocator, nodeKey, false, current) + } else { + currNode.prefixLen -= mismatchIdx + 1 + leafArtNode := minimum(&t.allocator, current) + leaf := leafArtNode.leaf(&t.allocator) + leafKey := artKey(leaf.GetKey()) + kMin := depth + mismatchIdx + 1 + kMax := depth + mismatchIdx + 1 + min(currNode.prefixLen, maxPrefixLen) + copy(currNode.prefix[:], leafKey[kMin:kMax]) + newArtNode.addChild(&t.allocator, leafKey.charAt(int(depth+mismatchIdx)), !leafKey.valid(int(depth)), current) + } + + // insert the artLeaf into new node + newLeafAddr, newLeaf := t.newLeaf(key) + newArtNode.addChild(&t.allocator, key.charAt(int(depth+mismatchIdx)), !key.valid(int(depth+mismatchIdx)), newLeafAddr) + if prev == nullArtNode { + t.root = newArtNode + } else { + prev.swapChild(&t.allocator, key.charAt(prevDepth), newArtNode) + } + return newLeafAddr.addr, newLeaf +} + +func (t *ART) newNode4() (artNode, *node4) { + addr, n4 := t.allocator.allocNode4() + return artNode{kind: typeNode4, addr: addr}, n4 +} + +func (t *ART) newLeaf(key artKey) (artNode, *artLeaf) { + addr, lf := t.allocator.allocLeaf(key) + return artNode{kind: typeLeaf, addr: addr}, lf +} + +func (t *ART) setValue(addr arena.MemdbArenaAddr, l *artLeaf, value []byte, ops []kv.FlagsOp) { + flags := l.getKeyFlags() + if flags == 0 && l.vAddr.IsNull() { + t.len++ + t.size += int(l.klen) + } + if value != nil { + flags = kv.ApplyFlagsOps(flags, append([]kv.FlagsOp{kv.DelNeedConstraintCheckInPrewrite}, ops...)...) + } else { + // an UpdateFlag operation, do not delete the NeedConstraintCheckInPrewrite flag. + flags = kv.ApplyFlagsOps(flags, ops...) + } + if flags.AndPersistent() != 0 { + t.dirty = true + } + l.setKeyFlags(flags) + if value == nil { + // value == nil means it updates flags only. + return + } + if t.trySwapValue(l.vAddr, value) { + return + } + t.size += len(value) + vAddr := t.allocator.vlogAllocator.AppendValue(addr, l.vAddr, value) + l.vAddr = vAddr +} + +// trySwapValue checks if the value can be updated in place, return true if it's updated. +func (t *ART) trySwapValue(addr arena.MemdbArenaAddr, value []byte) bool { + if addr.IsNull() { + return false + } + if len(t.stages) > 0 { + cp := t.stages[len(t.stages)-1] + if !t.allocator.vlogAllocator.CanModify(&cp, addr) { + return false + } + } + oldVal := t.allocator.vlogAllocator.GetValue(addr) + if len(oldVal) > 0 && len(oldVal) == len(value) { + copy(oldVal, value) + return true + } + t.size -= len(oldVal) + return false } func (t *ART) Dirty() bool { @@ -74,12 +365,12 @@ func (t *ART) Mem() uint64 { // Len returns the count of entries in the MemBuffer. func (t *ART) Len() int { - panic("unimplemented") + return t.len } // Size returns the size of the MemBuffer. func (t *ART) Size() int { - panic("unimplemented") + return t.size } func (t *ART) checkpoint() arena.MemDBCheckpoint { @@ -109,15 +400,13 @@ func (t *ART) Stages() []arena.MemDBCheckpoint { } func (t *ART) Staging() int { - panic("unimplemented") + return 0 } func (t *ART) Release(h int) { - panic("unimplemented") } func (t *ART) Cleanup(h int) { - panic("unimplemented") } func (t *ART) revertToCheckpoint(cp *arena.MemDBCheckpoint) { @@ -132,6 +421,18 @@ func (t *ART) truncate(snap *arena.MemDBCheckpoint) { panic("unimplemented") } +// Reset resets the MemBuffer to initial states. +func (t *ART) Reset() { + t.root = nullArtNode + t.stages = t.stages[:0] + t.dirty = false + t.vlogInvalid = false + t.size = 0 + t.len = 0 + t.allocator.nodeAllocator.Reset() + t.allocator.vlogAllocator.Reset() +} + // DiscardValues releases the memory used by all values. // NOTE: any operation need value will panic after this function. func (t *ART) DiscardValues() { diff --git a/internal/unionstore/art/art_arena.go b/internal/unionstore/art/art_arena.go index 8b29ef440..4e4f2dd70 100644 --- a/internal/unionstore/art/art_arena.go +++ b/internal/unionstore/art/art_arena.go @@ -16,6 +16,8 @@ package art import ( + "unsafe" + "github.com/tikv/client-go/v2/internal/unionstore/arena" ) @@ -33,3 +35,123 @@ type artAllocator struct { vlogAllocator arena.MemdbVlog[*artLeaf, *ART] nodeAllocator nodeArena } + +// init the allocator. +func (f *artAllocator) init() { + f.nodeAllocator.freeNode4 = make([]arena.MemdbArenaAddr, 0, 1<<4) + f.nodeAllocator.freeNode16 = make([]arena.MemdbArenaAddr, 0, 1<<3) + f.nodeAllocator.freeNode48 = make([]arena.MemdbArenaAddr, 0, 1<<2) +} + +func (f *artAllocator) allocNode4() (arena.MemdbArenaAddr, *node4) { + var ( + addr arena.MemdbArenaAddr + data []byte + ) + if len(f.nodeAllocator.freeNode4) > 0 { + addr = f.nodeAllocator.freeNode4[len(f.nodeAllocator.freeNode4)-1] + f.nodeAllocator.freeNode4 = f.nodeAllocator.freeNode4[:len(f.nodeAllocator.freeNode4)-1] + data = f.nodeAllocator.GetData(addr) + } else { + addr, data = f.nodeAllocator.Alloc(node4size, true) + } + n4 := (*node4)(unsafe.Pointer(&data[0])) + n4.init() + return addr, n4 +} + +func (f *artAllocator) freeNode4(addr arena.MemdbArenaAddr) { + f.nodeAllocator.freeNode4 = append(f.nodeAllocator.freeNode4, addr) +} + +func (f *artAllocator) getNode4(addr arena.MemdbArenaAddr) *node4 { + data := f.nodeAllocator.GetData(addr) + return (*node4)(unsafe.Pointer(&data[0])) +} + +func (f *artAllocator) allocNode16() (arena.MemdbArenaAddr, *node16) { + var ( + addr arena.MemdbArenaAddr + data []byte + ) + if len(f.nodeAllocator.freeNode16) > 0 { + addr = f.nodeAllocator.freeNode16[len(f.nodeAllocator.freeNode16)-1] + f.nodeAllocator.freeNode16 = f.nodeAllocator.freeNode16[:len(f.nodeAllocator.freeNode16)-1] + data = f.nodeAllocator.GetData(addr) + } else { + addr, data = f.nodeAllocator.Alloc(node16size, true) + } + n16 := (*node16)(unsafe.Pointer(&data[0])) + n16.init() + return addr, n16 +} + +func (f *artAllocator) freeNode16(addr arena.MemdbArenaAddr) { + f.nodeAllocator.freeNode16 = append(f.nodeAllocator.freeNode16, addr) +} + +func (f *artAllocator) getNode16(addr arena.MemdbArenaAddr) *node16 { + data := f.nodeAllocator.GetData(addr) + return (*node16)(unsafe.Pointer(&data[0])) +} + +func (f *artAllocator) allocNode48() (arena.MemdbArenaAddr, *node48) { + var ( + addr arena.MemdbArenaAddr + data []byte + ) + if len(f.nodeAllocator.freeNode48) > 0 { + addr = f.nodeAllocator.freeNode48[len(f.nodeAllocator.freeNode48)-1] + f.nodeAllocator.freeNode48 = f.nodeAllocator.freeNode48[:len(f.nodeAllocator.freeNode48)-1] + data = f.nodeAllocator.GetData(addr) + } else { + addr, data = f.nodeAllocator.Alloc(node48size, true) + } + n48 := (*node48)(unsafe.Pointer(&data[0])) + n48.init() + return addr, n48 +} + +func (f *artAllocator) freeNode48(addr arena.MemdbArenaAddr) { + f.nodeAllocator.freeNode48 = append(f.nodeAllocator.freeNode48, addr) +} + +func (f *artAllocator) getNode48(addr arena.MemdbArenaAddr) *node48 { + data := f.nodeAllocator.GetData(addr) + return (*node48)(unsafe.Pointer(&data[0])) +} + +func (f *artAllocator) allocNode256() (arena.MemdbArenaAddr, *node256) { + var ( + addr arena.MemdbArenaAddr + data []byte + ) + addr, data = f.nodeAllocator.Alloc(node256size, true) + n256 := (*node256)(unsafe.Pointer(&data[0])) + n256.init() + return addr, n256 +} + +func (f *artAllocator) getNode256(addr arena.MemdbArenaAddr) *node256 { + data := f.nodeAllocator.GetData(addr) + return (*node256)(unsafe.Pointer(&data[0])) +} + +func (f *artAllocator) allocLeaf(key []byte) (arena.MemdbArenaAddr, *artLeaf) { + size := leafSize + len(key) + addr, data := f.nodeAllocator.Alloc(size, true) + lf := (*artLeaf)(unsafe.Pointer(&data[0])) + lf.klen = uint16(len(key)) + lf.flags = 0 + lf.vAddr = arena.NullAddr + copy(data[leafSize:], key) + return addr, lf +} + +func (f *artAllocator) getLeaf(addr arena.MemdbArenaAddr) *artLeaf { + if addr.IsNull() { + return nil + } + data := f.nodeAllocator.GetData(addr) + return (*artLeaf)(unsafe.Pointer(&data[0])) +} diff --git a/internal/unionstore/art/art_node.go b/internal/unionstore/art/art_node.go index 53c568fe6..a0acb262b 100644 --- a/internal/unionstore/art/art_node.go +++ b/internal/unionstore/art/art_node.go @@ -12,47 +12,669 @@ // See the License for the specific language governing permissions and // limitations under the License. -//nolint:unused package art import ( + "bytes" "github.com/tikv/client-go/v2/internal/unionstore/arena" "github.com/tikv/client-go/v2/kv" + "math" + "math/bits" + "sort" + "unsafe" ) -type artNodeKind uint16 +type nodeKind uint16 const ( - typeARTInvalid artNodeKind = 0 - //nolint:unused - typeARTNode4 artNodeKind = 1 - typeARTNode16 artNodeKind = 2 - typeARTNode48 artNodeKind = 3 - typeARTNode256 artNodeKind = 4 - typeARTLeaf artNodeKind = 5 + typeInvalid nodeKind = 0 + typeNode4 nodeKind = 1 + typeNode16 nodeKind = 2 + typeNode48 nodeKind = 3 + typeNode256 nodeKind = 4 + typeLeaf nodeKind = 5 +) + +const ( + maxPrefixLen = 20 + node4cap = 4 + node16cap = 16 + node48cap = 48 + node256cap = 256 + inplaceIndex = -1 + notExistIndex = -2 +) + +const ( + node4size = int(unsafe.Sizeof(node4{})) + node16size = int(unsafe.Sizeof(node16{})) + node48size = int(unsafe.Sizeof(node48{})) + node256size = int(unsafe.Sizeof(node256{})) + leafSize = int(unsafe.Sizeof(artLeaf{})) +) + +var nullArtNode = artNode{kind: typeInvalid, addr: arena.NullAddr} + +var ( + // nullNode48 is used to initialize the node48 + nullNode48 = node48{} + // nullNode256 is used to initialize the node256 + nullNode256 = node256{} ) -var nullArtNode = artNode{kind: typeARTInvalid, addr: arena.NullAddr} +func init() { + for i := 0; i < node48cap; i++ { + nullNode48.children[i] = artNode{kind: typeInvalid, addr: arena.NullAddr} + } + for i := 0; i < node256cap; i++ { + nullNode256.children[i] = artNode{kind: typeInvalid, addr: arena.NullAddr} + } +} type artKey []byte type artNode struct { - kind artNodeKind + kind nodeKind addr arena.MemdbArenaAddr } +type nodeBase struct { + nodeNum uint8 + prefixLen uint32 + prefix [maxPrefixLen]byte + inplaceLeaf artNode +} + +func (n *nodeBase) init() { + // initialize basic nodeBase + n.nodeNum = 0 + n.prefixLen = 0 + n.inplaceLeaf = nullArtNode +} + +type node4 struct { + nodeBase + keys [node4cap]byte + children [node4cap]artNode +} + +type node16 struct { + nodeBase + keys [node16cap]byte + children [node16cap]artNode +} + +// n48s and n48m are used to calculate the node48's and node256's present index and bit index. +// Given idx between 0 and 255: +// +// present index = idx >> n48s +// bit index = idx % n48m +const ( + n48s = 6 // 2^n48s == n48m + n48m = 64 // it should be sizeof(node48.present[0]) +) + +type node48 struct { + nodeBase + // present is a bitmap to indicate the existence of children. + // The bitmap is divided into 4 uint64 mainly aims to speed up the iterator. + present [4]uint64 + keys [node256cap]uint8 // map byte to index + children [node48cap]artNode +} + +type node256 struct { + nodeBase + // present is a bitmap to indicate the existence of children. + // The bitmap is divided into 4 uint64 mainly aims to speed up the iterator. + present [4]uint64 + children [node256cap]artNode +} + type artLeaf struct { vAddr arena.MemdbArenaAddr klen uint16 flags uint16 } +func (an *artNode) isLeaf() bool { + return an.kind == typeLeaf +} + +func (an *artNode) leaf(a *artAllocator) *artLeaf { + return a.getLeaf(an.addr) +} + +// node gets the inner baseNode of the artNode +func (an *artNode) node(a *artAllocator) *nodeBase { + switch an.kind { + case typeNode4: + return &an.node4(a).nodeBase + case typeNode16: + return &an.node16(a).nodeBase + case typeNode48: + return &an.node48(a).nodeBase + case typeNode256: + return &an.node256(a).nodeBase + default: + panic("invalid nodeBase kind") + } +} + +// at returns the nth child of the node. +func (an *artNode) at(a *artAllocator, idx int) artNode { + switch an.kind { + case typeNode4: + return an.node4(a).children[idx] + case typeNode16: + return an.node16(a).children[idx] + case typeNode48: + n48 := an.node48(a) + return n48.children[n48.keys[idx]] + case typeNode256: + return an.node256(a).children[idx] + default: + panic("invalid nodeBase kind") + } +} + +func (n48 *node48) init() { + // initialize nodeBase + n48.nodeBase.init() + // initialize node48 + n48.present[0], n48.present[1], n48.present[2], n48.present[3] = 0, 0, 0, 0 + copy(n48.children[:], nullNode48.children[:]) +} + +// nextPresentIdx returns the next present index starting from the given index. +// Finding from present bitmap is faster than from keys. +func (n48 *node48) nextPresentIdx(start int) int { + for presentOffset := start >> n48s; presentOffset < 4; presentOffset++ { + offset := start % n48m + start = 0 + mask := math.MaxUint64 - (uint64(1) << offset) + 1 // e.g. offset=3 => 0b111...111000 + curr := n48.present[presentOffset] & mask + zeros := bits.TrailingZeros64(curr) + if zeros < n48m { + return presentOffset*n48m + zeros + } + } + return node256cap +} + +// prevPresentIdx returns the next present index starting from the given index. +// Finding from present bitmap is faster than from keys. +func (n48 *node48) prevPresentIdx(start int) int { + for presentOffset := start >> n48s; presentOffset >= 0; presentOffset-- { + offset := start % n48m + start = n48m - 1 + mask := uint64(1)<<(offset+1) - 1 // e.g. offset=3 => 0b000...0001111 + curr := n48.present[presentOffset] & mask + zeros := bits.LeadingZeros64(curr) + if zeros < n48m { + return presentOffset*n48m + n48m - (zeros + 1) + } + } + return inplaceIndex +} + +func (n256 *node256) init() { + // initialize nodeBase + n256.nodeBase.init() + // initialize node256 + copy(n256.children[:], nullNode256.children[:]) +} + +func (n256 *node256) nextPresentIdx(start int) int { + for presentOffset := start >> n48s; presentOffset < 4; presentOffset++ { + offset := start % n48m + start = 0 + mask := math.MaxUint64 - (uint64(1) << offset) + 1 // e.g. offset=3 => 0b111...111000 + curr := n256.present[presentOffset] & mask + zeros := bits.TrailingZeros64(curr) + if zeros < n48m { + return presentOffset*n48m + zeros + } + } + return node256cap +} + +func (n256 *node256) prevPresentIdx(start int) int { + for presentOffset := start >> n48s; presentOffset >= 0; presentOffset-- { + offset := start % n48m + start = n48m - 1 + mask := uint64(1)<<(offset+1) - 1 // e.g. offset=3 => 0b000...0001111 + curr := n256.present[presentOffset] & mask + zeros := bits.LeadingZeros64(curr) + if zeros < n48m { + return presentOffset*n48m + n48m - (zeros + 1) + } + } + return inplaceIndex +} + +// key methods + +func (k artKey) charAt(pos int) byte { + if pos < 0 || pos >= len(k) { + return 0 + } + return k[pos] +} + +func (k artKey) valid(pos int) bool { + return pos < len(k) +} + // GetKey gets the full key of the leaf func (l *artLeaf) GetKey() []byte { - panic("unimplemented") + base := unsafe.Add(unsafe.Pointer(l), leafSize) + return unsafe.Slice((*byte)(base), int(l.klen)) +} + +// getKeyDepth gets the partial key start from depth of the artLeaf +func (l *artLeaf) getKeyDepth(depth uint32) []byte { + base := unsafe.Add(unsafe.Pointer(l), leafSize+int(depth)) + return unsafe.Slice((*byte)(base), int(l.klen)-int(depth)) } // GetKeyFlags gets the flags of the leaf func (l *artLeaf) GetKeyFlags() kv.KeyFlags { panic("unimplemented") } + +func (l *artLeaf) match(depth uint32, key artKey) bool { + return bytes.Equal(l.getKeyDepth(depth), key[depth:]) +} + +func (l *artLeaf) setKeyFlags(flags kv.KeyFlags) arena.MemdbArenaAddr { + l.flags = uint16(flags) & flagMask + return l.vAddr +} + +func (l *artLeaf) getKeyFlags() kv.KeyFlags { + return kv.KeyFlags(l.flags & flagMask) +} + +const ( + deleteFlag uint16 = 1 << 15 + flagMask = ^deleteFlag +) + +// markDelete marks the artLeaf as deleted +func (l *artLeaf) markDelete() { + l.flags = deleteFlag +} + +//nolint:unused +func (l *artLeaf) unmarkDelete() { + l.flags &= flagMask +} + +func (l *artLeaf) isDeleted() bool { + return l.flags&deleteFlag != 0 +} + +// node methods +func (n *nodeBase) setPrefix(key artKey, prefixLen uint32) { + n.prefixLen = prefixLen + copy(n.prefix[:], key[:min(prefixLen, maxPrefixLen)]) +} + +func (n *nodeBase) match(key artKey, depth uint32) uint32 { + idx := uint32(0) + + limit := min(min(n.prefixLen, maxPrefixLen), uint32(len(key))-depth) + for ; idx < limit; idx++ { + if n.prefix[idx] != key[idx+depth] { + return idx + } + } + + return idx +} + +func (an *artNode) node4(a *artAllocator) *node4 { + return a.getNode4(an.addr) +} + +func (an *artNode) node16(a *artAllocator) *node16 { + return a.getNode16(an.addr) +} + +func (an *artNode) node48(a *artAllocator) *node48 { + return a.getNode48(an.addr) +} + +func (an *artNode) node256(a *artAllocator) *node256 { + return a.getNode256(an.addr) +} + +func (an *artNode) matchDeep(a *artAllocator, key artKey, depth uint32) uint32 /* mismatch index*/ { + n := an.node(a) + // match in-node prefix + mismatchIdx := n.match(key, depth) + if mismatchIdx < maxPrefixLen || n.prefixLen <= maxPrefixLen { + return mismatchIdx + } + // if the prefixLen is longer maxPrefixLen and mismatchIdx == maxPrefixLen, we need to match deeper with any leaf. + leafArtNode := minimum(a, *an) + lKey := leafArtNode.leaf(a).GetKey() + return longestCommonPrefix(lKey, key, depth+maxPrefixLen) + maxPrefixLen +} + +func longestCommonPrefix(l1Key, l2Key artKey, depth uint32) uint32 { + idx, limit := depth, min(uint32(len(l1Key)), uint32(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. + for ; idx < limit; idx++ { + if l1Key[idx] != l2Key[idx] { + break + } + } + return idx - depth +} + +// Find the minimum artLeaf under an artNode +func minimum(a *artAllocator, an artNode) artNode { + for { + switch an.kind { + case typeLeaf: + return an + case typeNode4: + n4 := an.node4(a) + if !n4.inplaceLeaf.addr.IsNull() { + return n4.inplaceLeaf + } + an = n4.children[0] + case typeNode16: + n16 := an.node16(a) + if !n16.inplaceLeaf.addr.IsNull() { + return n16.inplaceLeaf + } + an = n16.children[0] + case typeNode48: + n48 := an.node48(a) + if !n48.inplaceLeaf.addr.IsNull() { + return n48.inplaceLeaf + } + idx := n48.nextPresentIdx(0) + an = n48.children[n48.keys[idx]] + case typeNode256: + n256 := an.node256(a) + if !n256.inplaceLeaf.addr.IsNull() { + return n256.inplaceLeaf + } + idx := n256.nextPresentIdx(0) + an = n256.children[idx] + case typeInvalid: + return nullArtNode + } + } +} + +func (an *artNode) findChild(a *artAllocator, c byte, valid bool) (int, artNode) { + if !valid { + return inplaceIndex, an.node(a).inplaceLeaf + } + switch an.kind { + case typeNode4: + return an.findChild4(a, c) + case typeNode16: + return an.findChild16(a, c) + case typeNode48: + return an.findChild48(a, c) + case typeNode256: + return an.findChild256(a, c) + } + return notExistIndex, nullArtNode +} + +func (an *artNode) findChild4(a *artAllocator, c byte) (int, artNode) { + n4 := an.node4(a) + for idx := 0; idx < int(n4.nodeNum); idx++ { + if n4.keys[idx] == c { + return idx, n4.children[idx] + } + } + + return -2, nullArtNode +} + +func (an *artNode) findChild16(a *artAllocator, c byte) (int, artNode) { + n16 := an.node16(a) + + idx, found := sort.Find(int(n16.nodeNum), func(i int) int { + if n16.keys[i] < c { + return 1 + } + if n16.keys[i] == c { + return 0 + } + return -1 + }) + if found { + return idx, n16.children[idx] + } + return -2, nullArtNode +} + +func (an *artNode) findChild48(a *artAllocator, c byte) (int, artNode) { + n48 := an.node48(a) + if n48.present[c>>n48s]&(1<<(c%n48m)) != 0 { + return int(c), n48.children[n48.keys[c]] + } + return -2, nullArtNode +} + +func (an *artNode) findChild256(a *artAllocator, c byte) (int, artNode) { + n256 := an.node256(a) + return int(c), n256.children[c] +} + +func (an *artNode) swapChild(a *artAllocator, c byte, child artNode) { + switch an.kind { + case typeNode4: + n4 := an.node4(a) + for idx := uint8(0); idx < n4.nodeNum; idx++ { + if n4.keys[idx] == c { + n4.children[idx] = child + return + } + } + panic("swap child failed") + case typeNode16: + n16 := an.node16(a) + for idx := uint8(0); idx < n16.nodeNum; idx++ { + if n16.keys[idx] == c { + n16.children[idx] = child + return + } + } + panic("swap child failed") + case typeNode48: + n48 := an.node48(a) + if n48.present[c>>n48s]&(1<<(c%n48m)) != 0 { + n48.children[n48.keys[c]] = child + return + } + panic("swap child failed") + case typeNode256: + an.addChild256(a, c, false, child) + } +} + +// addChild adds a child to the node. +// the added index `c` should not exist. +// Return if the node is grown to higher capacity. +func (an *artNode) addChild(a *artAllocator, c byte, inplace bool, child artNode) bool { + switch an.kind { + case typeNode4: + return an.addChild4(a, c, inplace, child) + case typeNode16: + return an.addChild16(a, c, inplace, child) + case typeNode48: + return an.addChild48(a, c, inplace, child) + case typeNode256: + return an.addChild256(a, c, inplace, child) + } + return false +} + +func (an *artNode) addChild4(a *artAllocator, c byte, inplace bool, child artNode) bool { + node := an.node4(a) + + if inplace { + node.inplaceLeaf = child + return false + } + + if node.nodeNum >= node4cap { + an.grow(a) + an.addChild(a, c, inplace, child) + return true + } + + i := uint8(0) + for ; i < node.nodeNum; i++ { + if c < node.keys[i] { + break + } + } + + if i < node.nodeNum { + for j := node.nodeNum; j > i; j-- { + node.keys[j] = node.keys[j-1] + node.children[j] = node.children[j-1] + } + } + node.keys[i] = c + node.children[i] = child + node.nodeNum++ + return false +} + +func (an *artNode) addChild16(a *artAllocator, c byte, inplace bool, child artNode) bool { + node := an.node16(a) + + if inplace { + node.inplaceLeaf = child + return false + } + + if node.nodeNum >= node16cap { + an.grow(a) + an.addChild(a, c, inplace, child) + return true + } + + i, _ := sort.Find(int(node.nodeNum), func(i int) int { + if node.keys[i] < c { + return 1 + } + if node.keys[i] == c { + return 0 + } + return -1 + }) + + 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]) + } + + node.keys[i] = c + node.children[i] = child + node.nodeNum++ + return false +} + +func (an *artNode) addChild48(a *artAllocator, c byte, inplace bool, child artNode) bool { + node := an.node48(a) + + if inplace { + node.inplaceLeaf = child + return false + } + + if node.nodeNum >= node48cap { + an.grow(a) + an.addChild(a, c, inplace, child) + return true + } + + node.keys[c] = node.nodeNum + node.present[c>>n48s] |= 1 << (c % n48m) + node.children[node.nodeNum] = child + node.nodeNum++ + return false +} + +func (an *artNode) addChild256(a *artAllocator, c byte, inplace bool, child artNode) bool { + node := an.node256(a) + + if inplace { + node.inplaceLeaf = child + return false + } + + node.present[c>>n48s] |= 1 << (c % n48m) + node.children[c] = child + node.nodeNum++ + return false +} + +func (n *nodeBase) copyMeta(src *nodeBase) { + n.nodeNum = src.nodeNum + n.prefixLen = src.prefixLen + n.inplaceLeaf = src.inplaceLeaf + copy(n.prefix[:], src.prefix[:]) +} + +func (an *artNode) grow(a *artAllocator) { + switch an.kind { + case typeNode4: + n4 := an.node4(a) + newAddr, n16 := a.allocNode16() + n16.copyMeta(&n4.nodeBase) + + 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.node16(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 node16 + a.freeNode16(an.addr) + an.kind = typeNode48 + an.addr = newAddr + case typeNode48: + n48 := an.node48(a) + 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_node_test.go b/internal/unionstore/art/art_node_test.go new file mode 100644 index 000000000..1431ec1db --- /dev/null +++ b/internal/unionstore/art/art_node_test.go @@ -0,0 +1,261 @@ +package art + +import ( + "strconv" + "testing" + + "github.com/stretchr/testify/require" + "github.com/tikv/client-go/v2/internal/unionstore/arena" +) + +func TestAllocNode(t *testing.T) { + checkNodeInitialization := func(t *testing.T, n any) { + var base *nodeBase + switch n := n.(type) { + case *node4: + base = &n.nodeBase + case *node16: + base = &n.nodeBase + case *node48: + base = &n.nodeBase + require.Equal(t, [4]uint64{0, 0, 0, 0}, n.present) + case *node256: + base = &n.nodeBase + for i := 0; i < node256cap; i++ { + require.Equal(t, n.children[i], nullArtNode) + } + default: + require.Fail(t, "unknown node type") + } + require.Equal(t, uint8(0), base.nodeNum) + require.Equal(t, uint32(0), base.prefixLen) + require.Equal(t, base.inplaceLeaf, nullArtNode) + } + + var allocator artAllocator + allocator.init() + cnt := 10_000 + + // alloc node4 + n4s := make([]arena.MemdbArenaAddr, 0, cnt) + for i := 0; i < cnt; i++ { + addr, n4 := allocator.allocNode4() + require.False(t, addr.IsNull()) + require.NotNil(t, n4) + checkNodeInitialization(t, n4) + n4.nodeNum = uint8(i % 4) + n4.prefixLen = uint32(i % maxPrefixLen) + n4s = append(n4s, addr) + } + + // alloc node16 + n16s := make([]arena.MemdbArenaAddr, 0, cnt) + for i := 0; i < cnt; i++ { + addr, n16 := allocator.allocNode16() + require.False(t, addr.IsNull()) + require.NotNil(t, n16) + checkNodeInitialization(t, n16) + n16.nodeNum = uint8(i % 16) + n16.prefixLen = uint32(i % maxPrefixLen) + n16s = append(n16s, addr) + } + + // alloc node48 + n48s := make([]arena.MemdbArenaAddr, 0, cnt) + for i := 0; i < cnt; i++ { + addr, n48 := allocator.allocNode48() + require.False(t, addr.IsNull()) + require.NotNil(t, n48) + checkNodeInitialization(t, n48) + n48.nodeNum = uint8(i % 48) + n48.prefixLen = uint32(i % maxPrefixLen) + n48s = append(n48s, addr) + } + + // alloc node256 + n256s := make([]arena.MemdbArenaAddr, 0, cnt) + for i := 0; i < cnt; i++ { + addr, n256 := allocator.allocNode256() + require.False(t, addr.IsNull()) + require.NotNil(t, n256) + checkNodeInitialization(t, n256) + n256.nodeNum = uint8(i % 256) + n256.prefixLen = uint32(i % maxPrefixLen) + n256s = append(n256s, addr) + } + + // alloc leaf + leafs := make([]arena.MemdbArenaAddr, 0, cnt) + for i := 0; i < cnt; i++ { + key := []byte(strconv.Itoa(i)) + addr, leaf := allocator.allocLeaf(key) + require.False(t, addr.IsNull()) + require.NotNil(t, leaf) + require.Equal(t, key, leaf.GetKey()) + leafs = append(leafs, addr) + } + + // test memory safety by checking the assign value + for i, addr := range n4s { + n4 := allocator.getNode4(addr) + require.Equal(t, uint8(i%4), n4.nodeNum, i) + require.Equal(t, uint32(i%maxPrefixLen), n4.prefixLen, i) + } + for i, addr := range n16s { + n16 := allocator.getNode16(addr) + require.Equal(t, uint8(i%16), n16.nodeNum) + require.Equal(t, uint32(i%maxPrefixLen), n16.prefixLen, i) + } + for i, addr := range n48s { + n48 := allocator.getNode48(addr) + require.Equal(t, uint8(i%48), n48.nodeNum) + require.Equal(t, uint32(i%maxPrefixLen), n48.prefixLen, i) + } + for i, addr := range n256s { + n256 := allocator.getNode256(addr) + require.Equal(t, uint8(i%256), n256.nodeNum) + require.Equal(t, uint32(i%maxPrefixLen), n256.prefixLen, i) + } + for i, addr := range leafs { + key := []byte(strconv.Itoa(i)) + leaf := allocator.getLeaf(addr) + require.Equal(t, key, leaf.GetKey()) + } +} + +func TestNodePrefix(t *testing.T) { + var allocator artAllocator + allocator.init() + + testFn := func(an *artNode, n *nodeBase) { + // simple match + key := []byte{1, 2, 3, 4, 5, 6, 1, 2, 3, 1} + n.setPrefix([]byte{1, 2, 3, 4, 5, 66, 77, 88, 99}, 5) + + idx := n.match(key, 0) // return mismatch index + require.Equal(t, uint32(5), idx) + idx = n.match(key[:4], 0) // key is shorter than prefix + require.Equal(t, uint32(4), idx) + idx = n.match(key, 1) // key[0+depth] != prefix[0] + require.Equal(t, uint32(0), idx) + idx = n.match(key, 6) // key[3+depth] != prefix[3] + require.Equal(t, uint32(3), idx) + idx = n.match(append([]byte{1}, key...), 1) + require.Equal(t, uint32(5), idx) + + // deep match + leafKey := append(make([]byte, maxPrefixLen), []byte{1, 2, 3, 4, 5}...) + leafAddr, _ := allocator.allocLeaf(leafKey) + an.addChild(&allocator, 2, false, artNode{kind: typeLeaf, addr: leafAddr}) + // the real prefix is [0, ..., 0, 1], but the node.prefix only store [0, ..., 0] + n.setPrefix(leafKey, maxPrefixLen+1) + matchKey := append(make([]byte, maxPrefixLen), []byte{1, 22, 33, 44, 55}...) + mismatchKey := append(make([]byte, maxPrefixLen), []byte{11, 22, 33, 44, 55}...) + require.Equal(t, uint32(maxPrefixLen+1), an.matchDeep(&allocator, matchKey, 0)) + require.Equal(t, uint32(maxPrefixLen), an.matchDeep(&allocator, mismatchKey, 0)) + + // deep match with depth + leafKey = append(make([]byte, 10), leafKey...) + matchKey = append(make([]byte, 10), matchKey...) + mismatchKey = append(make([]byte, 10), mismatchKey...) + leafAddr, _ = allocator.allocLeaf(leafKey) + an.swapChild(&allocator, 2, artNode{kind: typeLeaf, addr: leafAddr}) + n.setPrefix(leafKey[10:], maxPrefixLen+1) + require.Equal(t, uint32(maxPrefixLen+1), an.matchDeep(&allocator, matchKey, 10)) + require.Equal(t, uint32(maxPrefixLen), an.matchDeep(&allocator, mismatchKey, 10)) + } + + addr, n4 := allocator.allocNode4() + testFn(&artNode{kind: typeNode4, addr: addr}, &n4.nodeBase) + addr, n16 := allocator.allocNode16() + testFn(&artNode{kind: typeNode16, addr: addr}, &n16.nodeBase) + addr, n48 := allocator.allocNode48() + testFn(&artNode{kind: typeNode48, addr: addr}, &n48.nodeBase) + addr, n256 := allocator.allocNode256() + testFn(&artNode{kind: typeNode256, addr: addr}, &n256.nodeBase) +} + +func TestOrderedChild(t *testing.T) { + // the keys of node4 and node16 should be kept ordered when adding new children. + var allocator artAllocator + allocator.init() + + testFn := func(fn func() *artNode, getKeys func(*artNode) []byte, cap int) { + // add children in ascend order + an := fn() + for i := 0; i < cap; i++ { + addr, _ := allocator.allocNode4() + an.addChild(&allocator, byte(i), false, artNode{kind: typeNode4, addr: addr}) + + node := an.node(&allocator) + require.Equal(t, node.nodeNum, uint8(i+1)) + keys := getKeys(an) + for j := 0; j <= i; j++ { + require.Equal(t, byte(j), keys[j]) + } + } + + // add children in descend order + an = fn() + for i := 0; i < cap; i++ { + addr, _ := allocator.allocNode4() + an.addChild(&allocator, byte(255-i), false, artNode{kind: typeNode4, addr: addr}) + + node := an.node(&allocator) + require.Equal(t, node.nodeNum, uint8(i+1)) + keys := getKeys(an) + for j := 0; j <= i; j++ { + require.Equal(t, byte(255-i+j), keys[j]) + } + } + } + + testFn(func() *artNode { + addr, _ := allocator.allocNode4() + return &artNode{kind: typeNode4, addr: addr} + }, func(an *artNode) []byte { + return an.node4(&allocator).keys[:] + }, 4) + + testFn(func() *artNode { + addr, _ := allocator.allocNode16() + return &artNode{kind: typeNode16, addr: addr} + }, func(an *artNode) []byte { + return an.node16(&allocator).keys[:] + }, 16) +} + +func TestNextPrevPresentIdx(t *testing.T) { + // the present bitmap is used to find the next/prev present child index in node48 and node256. + var allocator artAllocator + allocator.init() + + testFn := func(an *artNode, node interface { + nextPresentIdx(int) int + prevPresentIdx(int) int + }) { + require.Equal(t, node.nextPresentIdx(0), node256cap) + require.Equal(t, node.prevPresentIdx(0), inplaceIndex) + + mid := 128 + addr, _ := allocator.allocNode4() + an.addChild(&allocator, byte(mid), false, artNode{kind: typeNode4, addr: addr}) + for i := 0; i < node256cap; i++ { + expectedNext := mid + expectedPrev := inplaceIndex + if i > mid { + expectedNext = node256cap + } + if i >= mid { + expectedPrev = mid + } + require.Equal(t, node.nextPresentIdx(i), expectedNext) + require.Equal(t, node.prevPresentIdx(i), expectedPrev) + } + } + + addr, n48 := allocator.allocNode48() + testFn(&artNode{kind: typeNode48, addr: addr}, n48) + addr, n256 := allocator.allocNode48() + testFn(&artNode{kind: typeNode48, addr: addr}, n256) +} diff --git a/internal/unionstore/art/art_test.go b/internal/unionstore/art/art_test.go new file mode 100644 index 000000000..837f5df8e --- /dev/null +++ b/internal/unionstore/art/art_test.go @@ -0,0 +1,191 @@ +package art + +import ( + "fmt" + "testing" + + "github.com/stretchr/testify/assert" + tikverr "github.com/tikv/client-go/v2/error" + "github.com/tikv/client-go/v2/kv" +) + +func TestSimple(t *testing.T) { + tree := New() + for i := 0; i < 256; i++ { + key := []byte{byte(i)} + _, err := tree.Get(key) + assert.Equal(t, err, tikverr.ErrNotExist) + err = tree.Set(key, key) + assert.Nil(t, err) + val, err := tree.Get(key) + assert.Nil(t, err, i) + assert.Equal(t, val, key, i) + } +} + +func TestSubNode(t *testing.T) { + tree := New() + assert.Nil(t, tree.Set([]byte("a"), []byte("a"))) + assert.Nil(t, tree.Set([]byte("aa"), []byte("aa"))) + assert.Nil(t, tree.Set([]byte("aaa"), []byte("aaa"))) + v, err := tree.Get([]byte("a")) + assert.Nil(t, err) + assert.Equal(t, v, []byte("a")) + v, err = tree.Get([]byte("aa")) + assert.Nil(t, err) + assert.Equal(t, v, []byte("aa")) + v, err = tree.Get([]byte("aaa")) + assert.Nil(t, err) + assert.Equal(t, v, []byte("aaa")) +} + +func BenchmarkReadAfterWriteArt(b *testing.B) { + buf := make([][]byte, b.N) + for i := 0; i < b.N; i++ { + key := []byte{byte(i)} + buf[i] = key + } + tree := New() + b.ResetTimer() + for i := 0; i < b.N; i++ { + tree.Set(buf[i], buf[i]) + v, _ := tree.Get(buf[i]) + assert.Equal(b, v, buf[i]) + } +} + +func encodeInt(n int) []byte { + return []byte(fmt.Sprintf("%010d", n)) +} + +func TestBenchKey(t *testing.T) { + buffer := New() + cnt := 100000 + for k := 0; k < cnt; k++ { + key, value := encodeInt(k), encodeInt(k) + buffer.Set(key, value) + } + for k := 0; k < cnt; k++ { + v, err := buffer.Get(encodeInt(k)) + assert.Nil(t, err, k) + assert.Equal(t, v, encodeInt(k)) + } +} + +func TestLeafWithCommonPrefix(t *testing.T) { + tree := New() + tree.Set([]byte{1, 1, 1}, []byte{1, 1, 1}) + tree.Set([]byte{1, 1, 2}, []byte{1, 1, 2}) + v, err := tree.Get([]byte{1, 1, 1}) + assert.Nil(t, err) + assert.Equal(t, v, []byte{1, 1, 1}) + v, err = tree.Get([]byte{1, 1, 2}) + assert.Nil(t, err) + assert.Equal(t, v, []byte{1, 1, 2}) +} + +func TestUpdateInplace(t *testing.T) { + tree := New() + key := []byte{0} + for i := 0; i < 256; i++ { + val := make([]byte, 4096) + tree.Set(key, val) + assert.Equal(t, tree.allocator.vlogAllocator.Blocks(), 1) + } +} + +func TestFlag(t *testing.T) { + tree := New() + tree.Set([]byte{0}, []byte{0}, kv.SetPresumeKeyNotExists) + flags, err := tree.GetFlags([]byte{0}) + assert.Nil(t, err) + assert.True(t, flags.HasPresumeKeyNotExists()) + tree.Set([]byte{1}, []byte{1}, kv.SetKeyLocked) + flags, err = tree.GetFlags([]byte{1}) + assert.Nil(t, err) + assert.True(t, flags.HasLocked()) + // iterate can also see the flags + //it, err := tree.Iter(nil, nil) + //assert.Nil(t, err) + //assert.True(t, it.Valid()) + //assert.Equal(t, it.Key(), []byte{0}) + //assert.Equal(t, it.Value(), []byte{0}) + //assert.True(t, it.Flags().HasPresumeKeyNotExists()) + //assert.False(t, it.Flags().HasLocked()) + //assert.Nil(t, it.Next()) + //assert.True(t, it.Valid()) + //assert.Equal(t, it.Key(), []byte{1}) + //assert.Equal(t, it.Value(), []byte{1}) + //assert.True(t, it.Flags().HasLocked()) + //assert.False(t, it.Flags().HasPresumeKeyNotExists()) + //assert.Nil(t, it.Next()) + //assert.False(t, it.Valid()) +} + +func TestLongPrefix1(t *testing.T) { + key1 := []byte{109, 68, 66, 115, 0, 0, 0, 0, 0, 250, 0, 0, 0, 0, 0, 0, 0, 104, 68, 66, 58, 49, 0, 0, 0, 0, 251} + key2 := []byte{109, 68, 66, 115, 0, 0, 0, 0, 0, 250, 0, 0, 0, 0, 0, 0, 0, 105, 68, 66, 58, 49, 0, 0, 0, 0, 251} + buffer := New() + assert.Nil(t, buffer.Set(key1, []byte{1})) + assert.Nil(t, buffer.Set(key2, []byte{2})) + val, err := buffer.Get(key1) + assert.Nil(t, err) + assert.Equal(t, val, []byte{1}) + val, err = buffer.Get(key2) + assert.Nil(t, err) + assert.Equal(t, val, []byte{2}) + assert.Nil(t, buffer.Set(key2, []byte{3})) + val, err = buffer.Get(key2) + assert.Nil(t, err) + assert.Equal(t, val, []byte{3}) +} + +func TestLongPrefix2(t *testing.T) { + tree := New() + key1 := []byte{0, 97, 0, 0, 0, 0, 0, 0, 0, 248, 0, 0, 0, 0, 0, 0, 0, 108, 127, 255, 255, 255, 255, 255, 255, 255} + key2 := []byte{0, 97, 0, 0, 0, 0, 0, 0, 0, 248, 0, 0, 0, 0, 0, 0, 0, 108, 127, 255, 255, 255, 255, 255, 255, 254} + key3 := []byte{0, 97, 0, 0, 0, 0, 0, 0, 0, 248, 0, 0, 0, 0, 0, 0, 0, 108, 127, 255, 255, 255, 255, 255, 255, 253} + key4 := []byte{0, 97, 0, 0, 0, 0, 0, 0, 0, 248, 0, 0, 0, 0, 0, 0, 0, 76} + key5 := []byte{0, 97, 0, 0, 0, 0, 0, 0, 0, 248, 0, 0, 0, 0, 0, 0, 0, 108, 127, 255, 255, 255, 255, 255, 255, 252} + assert.Nil(t, tree.Set(key1, key1)) + assert.Nil(t, tree.Set(key2, key2)) + assert.Nil(t, tree.Set(key3, key3)) + assert.Nil(t, tree.Set(key4, key4)) + assert.Nil(t, tree.Set(key5, key5)) + assert.Nil(t, tree.Set(key4, key4)) + val, err := tree.Get(key1) + assert.Nil(t, err) + assert.Equal(t, val, key1) + val, err = tree.Get(key2) + assert.Nil(t, err) + assert.Equal(t, val, key2) + val, err = tree.Get(key3) + assert.Nil(t, err) + assert.Equal(t, val, key3) + val, err = tree.Get(key4) + assert.Nil(t, err) + assert.Equal(t, val, key4) + val, err = tree.Get(key5) + assert.Nil(t, err) + assert.Equal(t, val, key5) +} + +func TestFlagOnlyKey(t *testing.T) { + tree := New() + tree.Set([]byte{0}, nil, kv.SetAssertNone) + flags, err := tree.GetFlags([]byte{0}) + assert.Nil(t, err) + assert.False(t, flags.HasAssertionFlags()) + _, err = tree.Get([]byte{0}) + assert.Error(t, err) +} + +func TestSearchOptimisticMismatch(t *testing.T) { + tree := New() + prefix := make([]byte, 22) + tree.Set(append(prefix, []byte{1}...), prefix) + tree.Set(append(prefix, []byte{2}...), prefix) + // the search key is matched within maxPrefixLen, but the full key is not matched. + _, err := tree.Get(append(make([]byte, 21), []byte{1, 1}...)) + assert.NotNil(t, err) +} diff --git a/internal/unionstore/memdb_art.go b/internal/unionstore/memdb_art.go index 85abc27c6..07ace8f1d 100644 --- a/internal/unionstore/memdb_art.go +++ b/internal/unionstore/memdb_art.go @@ -50,7 +50,7 @@ func (db *artDBWithContext) set(key, value []byte, ops []kv.FlagsOp) error { db.Lock() defer db.Unlock() } - return db.ART.Set(key, value, ops) + return db.ART.Set(key, value, ops...) } func (db *artDBWithContext) Set(key, value []byte) error { diff --git a/internal/unionstore/memdb_test.go b/internal/unionstore/memdb_test.go index 7522e6cb5..8d85b87fc 100644 --- a/internal/unionstore/memdb_test.go +++ b/internal/unionstore/memdb_test.go @@ -52,6 +52,7 @@ type KeyFlags = kv.KeyFlags func TestGetSet(t *testing.T) { testGetSet(t, newRbtDBWithContext()) + testGetSet(t, newArtDBWithContext()) } func testGetSet(t *testing.T, db MemBuffer) { @@ -220,6 +221,7 @@ func testFlushOverwrite(t *testing.T, db MemBuffer) { func TestComplexUpdate(t *testing.T) { testComplexUpdate(t, newRbtDBWithContext()) + testComplexUpdate(t, newArtDBWithContext()) } func testComplexUpdate(t *testing.T, db MemBuffer) { @@ -843,8 +845,11 @@ func TestUnsetTemporaryFlag(t *testing.T) { } func TestSnapshotGetIter(t *testing.T) { + testSnapshotGetIter(t, newRbtDBWithContext()) +} + +func testSnapshotGetIter(t *testing.T, db MemBuffer) { assert := assert.New(t) - db := NewMemDB() var getters []Getter var iters []Iterator var reverseIters []Iterator From b13eb375ccc9faf8938abd863545ff508072bc6e Mon Sep 17 00:00:00 2001 From: you06 Date: Mon, 26 Aug 2024 17:37:01 +0900 Subject: [PATCH 2/7] add benchmark test Signed-off-by: you06 --- internal/unionstore/memdb_bench_test.go | 178 ++++++++++++++--------- internal/unionstore/memdb_norace_test.go | 19 ++- internal/unionstore/memdb_test.go | 2 +- 3 files changed, 120 insertions(+), 79 deletions(-) diff --git a/internal/unionstore/memdb_bench_test.go b/internal/unionstore/memdb_bench_test.go index 2b7e6e69a..249b118a5 100644 --- a/internal/unionstore/memdb_bench_test.go +++ b/internal/unionstore/memdb_bench_test.go @@ -47,116 +47,152 @@ const ( ) func BenchmarkLargeIndex(b *testing.B) { - buf := make([][valueSize]byte, 10000000) - for i := range buf { - binary.LittleEndian.PutUint32(buf[i][:], uint32(i)) - } - db := NewMemDB() - b.ResetTimer() + fn := func(b *testing.B, p MemBuffer) { + buf := make([][valueSize]byte, 10000000) + for i := range buf { + binary.LittleEndian.PutUint32(buf[i][:], uint32(i)) + } + b.ResetTimer() - for i := range buf { - db.Set(buf[i][:keySize], buf[i][:]) + for i := range buf { + p.Set(buf[i][:keySize], buf[i][:]) + } } + + b.Run("RBT", func(b *testing.B) { fn(b, newRbtDBWithContext()) }) + b.Run("ART", func(b *testing.B) { fn(b, newArtDBWithContext()) }) } func BenchmarkPut(b *testing.B) { - buf := make([][valueSize]byte, b.N) - for i := range buf { - binary.BigEndian.PutUint32(buf[i][:], uint32(i)) + fn := func(b *testing.B, p MemBuffer) { + buf := make([][valueSize]byte, b.N) + for i := range buf { + binary.BigEndian.PutUint32(buf[i][:], uint32(i)) + } + b.ResetTimer() + for i := range buf { + p.Set(buf[i][:keySize], buf[i][:]) + } } - p := NewMemDB() - b.ResetTimer() - - for i := range buf { - p.Set(buf[i][:keySize], buf[i][:]) - } + b.Run("RBT", func(b *testing.B) { fn(b, newRbtDBWithContext()) }) + b.Run("ART", func(b *testing.B) { fn(b, newArtDBWithContext()) }) } func BenchmarkPutRandom(b *testing.B) { - buf := make([][valueSize]byte, b.N) - for i := range buf { - binary.LittleEndian.PutUint32(buf[i][:], uint32(rand.Int())) + fn := func(b *testing.B, p MemBuffer) { + buf := make([][valueSize]byte, b.N) + for i := range buf { + binary.LittleEndian.PutUint32(buf[i][:], uint32(rand.Int())) + } + b.ResetTimer() + for i := range buf { + p.Set(buf[i][:keySize], buf[i][:]) + } } - p := NewMemDB() - b.ResetTimer() - - for i := range buf { - p.Set(buf[i][:keySize], buf[i][:]) - } + b.Run("RBT", func(b *testing.B) { fn(b, newRbtDBWithContext()) }) + b.Run("ART", func(b *testing.B) { fn(b, newArtDBWithContext()) }) } func BenchmarkGet(b *testing.B) { - buf := make([][valueSize]byte, b.N) - for i := range buf { - binary.BigEndian.PutUint32(buf[i][:], uint32(i)) - } + fn := func(b *testing.B, p MemBuffer) { + buf := make([][valueSize]byte, b.N) + for i := range buf { + binary.BigEndian.PutUint32(buf[i][:], uint32(i)) + } - p := NewMemDB() - for i := range buf { - p.Set(buf[i][:keySize], buf[i][:]) - } + for i := range buf { + p.Set(buf[i][:keySize], buf[i][:]) + } - ctx := context.Background() - b.ResetTimer() - for i := range buf { - p.Get(ctx, buf[i][:keySize]) + ctx := context.Background() + b.ResetTimer() + for i := range buf { + p.Get(ctx, buf[i][:keySize]) + } } + + b.Run("RBT", func(b *testing.B) { fn(b, newRbtDBWithContext()) }) + b.Run("ART", func(b *testing.B) { fn(b, newArtDBWithContext()) }) } func BenchmarkGetRandom(b *testing.B) { - buf := make([][valueSize]byte, b.N) - for i := range buf { - binary.LittleEndian.PutUint32(buf[i][:], uint32(rand.Int())) - } + fn := func(b *testing.B, p MemBuffer) { + buf := make([][valueSize]byte, b.N) + for i := range buf { + binary.LittleEndian.PutUint32(buf[i][:], uint32(rand.Int())) + } - p := NewMemDB() - for i := range buf { - p.Set(buf[i][:keySize], buf[i][:]) - } + for i := range buf { + p.Set(buf[i][:keySize], buf[i][:]) + } - ctx := context.Background() - b.ResetTimer() - for i := 0; i < b.N; i++ { - p.Get(ctx, buf[i][:keySize]) + ctx := context.Background() + b.ResetTimer() + for i := 0; i < b.N; i++ { + p.Get(ctx, buf[i][:keySize]) + } } + + b.Run("RBT", func(b *testing.B) { fn(b, newRbtDBWithContext()) }) + b.Run("ART", func(b *testing.B) { fn(b, newArtDBWithContext()) }) } var opCnt = 100000 func BenchmarkMemDbBufferSequential(b *testing.B) { - data := make([][]byte, opCnt) - for i := 0; i < opCnt; i++ { - data[i] = encodeInt(i) + fn := func(b *testing.B, buffer MemBuffer) { + data := make([][]byte, opCnt) + for i := 0; i < opCnt; i++ { + data[i] = encodeInt(i) + } + benchmarkSetGet(b, buffer, data) + b.ReportAllocs() } - buffer := NewMemDB() - benchmarkSetGet(b, buffer, data) - b.ReportAllocs() + + b.Run("RBT", func(b *testing.B) { fn(b, newRbtDBWithContext()) }) + b.Run("ART", func(b *testing.B) { fn(b, newArtDBWithContext()) }) } func BenchmarkMemDbBufferRandom(b *testing.B) { - data := make([][]byte, opCnt) - for i := 0; i < opCnt; i++ { - data[i] = encodeInt(i) + fn := func(b *testing.B, buffer MemBuffer) { + data := make([][]byte, opCnt) + for i := 0; i < opCnt; i++ { + data[i] = encodeInt(i) + } + shuffle(data) + benchmarkSetGet(b, buffer, data) + b.ReportAllocs() } - shuffle(data) - buffer := NewMemDB() - benchmarkSetGet(b, buffer, data) - b.ReportAllocs() + + b.Run("RBT", func(b *testing.B) { fn(b, newRbtDBWithContext()) }) + b.Run("ART", func(b *testing.B) { fn(b, newArtDBWithContext()) }) } func BenchmarkMemDbIter(b *testing.B) { - buffer := NewMemDB() - benchIterator(b, buffer) - b.ReportAllocs() + fn := func(b *testing.B, buffer MemBuffer) { + benchIterator(b, buffer) + b.ReportAllocs() + } + + b.Run("RBT", func(b *testing.B) { fn(b, newRbtDBWithContext()) }) + b.Run("ART", func(b *testing.B) { + b.Skip("unimplemented") + fn(b, newArtDBWithContext()) + }) } func BenchmarkMemDbCreation(b *testing.B) { - for i := 0; i < b.N; i++ { - NewMemDB() + fn := func(b *testing.B, createFn func() MemBuffer) { + for i := 0; i < b.N; i++ { + createFn() + } + b.ReportAllocs() } - b.ReportAllocs() + + b.Run("RBT", func(b *testing.B) { fn(b, func() MemBuffer { return newRbtDBWithContext() }) }) + b.Run("ART", func(b *testing.B) { fn(b, func() MemBuffer { return newArtDBWithContext() }) }) } func shuffle(slc [][]byte) { @@ -167,7 +203,7 @@ func shuffle(slc [][]byte) { slc[r], slc[i] = slc[i], slc[r] } } -func benchmarkSetGet(b *testing.B, buffer *MemDB, data [][]byte) { +func benchmarkSetGet(b *testing.B, buffer MemBuffer, data [][]byte) { ctx := context.Background() b.ResetTimer() for i := 0; i < b.N; i++ { @@ -180,7 +216,7 @@ func benchmarkSetGet(b *testing.B, buffer *MemDB, data [][]byte) { } } -func benchIterator(b *testing.B, buffer *MemDB) { +func benchIterator(b *testing.B, buffer MemBuffer) { for k := 0; k < opCnt; k++ { buffer.Set(encodeInt(k), encodeInt(k)) } diff --git a/internal/unionstore/memdb_norace_test.go b/internal/unionstore/memdb_norace_test.go index c8d8e54fe..f669ef798 100644 --- a/internal/unionstore/memdb_norace_test.go +++ b/internal/unionstore/memdb_norace_test.go @@ -59,31 +59,36 @@ func TestRandom(t *testing.T) { rand2.Read(keys[i]) } - p1 := NewMemDB() + rbtDB := newRbtDBWithContext() + artDB := newArtDBWithContext() p2 := leveldb.New(comparer.DefaultComparer, 4*1024) for _, k := range keys { - p1.Set(k, k) + rbtDB.Set(k, k) + artDB.Set(k, k) _ = p2.Put(k, k) } - require.Equal(p1.Len(), p2.Len()) - require.Equal(p1.Size(), p2.Size()) + require.Equal(rbtDB.Len(), p2.Len()) + require.Equal(rbtDB.Size(), p2.Size()) + + require.Equal(artDB.Len(), p2.Len()) + require.Equal(artDB.Size(), p2.Size()) rand.Shuffle(cnt, func(i, j int) { keys[i], keys[j] = keys[j], keys[i] }) for _, k := range keys { op := rand.Float64() if op < 0.35 { - p1.RemoveFromBuffer(k) + rbtDB.RemoveFromBuffer(k) p2.Delete(k) } else { newValue := make([]byte, rand.Intn(19)+1) rand2.Read(newValue) - p1.Set(k, newValue) + rbtDB.Set(k, newValue) _ = p2.Put(k, newValue) } } - checkConsist(t, p1, p2) + checkConsist(t, rbtDB, p2) } // The test takes too long under the race detector. diff --git a/internal/unionstore/memdb_test.go b/internal/unionstore/memdb_test.go index 8d85b87fc..62ae308d1 100644 --- a/internal/unionstore/memdb_test.go +++ b/internal/unionstore/memdb_test.go @@ -549,7 +549,7 @@ func testFlags(t *testing.T, db MemBuffer, iterWithFlags func(db MemBuffer) Iter } } -func checkConsist(t *testing.T, p1 *MemDB, p2 *leveldb.DB) { +func checkConsist(t *testing.T, p1 MemBuffer, p2 *leveldb.DB) { assert := assert.New(t) assert.Equal(p1.Len(), p2.Len()) From 4b2722d172c6f29092f3caca278ef2eb64727fea Mon Sep 17 00:00:00 2001 From: you06 Date: Tue, 27 Aug 2024 14:42:21 +0900 Subject: [PATCH 3/7] lint Signed-off-by: you06 --- internal/unionstore/art/art.go | 4 +++ internal/unionstore/art/art_arena.go | 1 - internal/unionstore/art/art_node.go | 35 +++++++++--------------- internal/unionstore/art/art_node_test.go | 4 +-- 4 files changed, 18 insertions(+), 26 deletions(-) diff --git a/internal/unionstore/art/art.go b/internal/unionstore/art/art.go index 3bfdd6252..bb7a91165 100644 --- a/internal/unionstore/art/art.go +++ b/internal/unionstore/art/art.go @@ -48,10 +48,12 @@ func New() *ART { } func (t *ART) Get(key []byte) ([]byte, error) { + // 1. search the leaf node. _, leaf := t.search(key) if leaf == nil || leaf.vAddr.IsNull() { return nil, tikverr.ErrNotExist } + // 2. get the value from the vlog. return t.allocator.vlogAllocator.GetValue(leaf.vAddr), nil } @@ -79,7 +81,9 @@ 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 exist leaf in the tree. addr, leaf := t.recursiveInsert(key) + // 2. set the value and flags. t.setValue(addr, leaf, value, ops) if uint64(t.Size()) > t.bufferSizeLimit { return &tikverr.ErrTxnTooLarge{Size: t.Size()} diff --git a/internal/unionstore/art/art_arena.go b/internal/unionstore/art/art_arena.go index 4e4f2dd70..84ca21127 100644 --- a/internal/unionstore/art/art_arena.go +++ b/internal/unionstore/art/art_arena.go @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -//nolint:unused package art import ( diff --git a/internal/unionstore/art/art_node.go b/internal/unionstore/art/art_node.go index a0acb262b..53e5e4f3f 100644 --- a/internal/unionstore/art/art_node.go +++ b/internal/unionstore/art/art_node.go @@ -16,12 +16,13 @@ package art import ( "bytes" - "github.com/tikv/client-go/v2/internal/unionstore/arena" - "github.com/tikv/client-go/v2/kv" "math" "math/bits" "sort" "unsafe" + + "github.com/tikv/client-go/v2/internal/unionstore/arena" + "github.com/tikv/client-go/v2/kv" ) type nodeKind uint16 @@ -55,22 +56,6 @@ const ( var nullArtNode = artNode{kind: typeInvalid, addr: arena.NullAddr} -var ( - // nullNode48 is used to initialize the node48 - nullNode48 = node48{} - // nullNode256 is used to initialize the node256 - nullNode256 = node256{} -) - -func init() { - for i := 0; i < node48cap; i++ { - nullNode48.children[i] = artNode{kind: typeInvalid, addr: arena.NullAddr} - } - for i := 0; i < node256cap; i++ { - nullNode256.children[i] = artNode{kind: typeInvalid, addr: arena.NullAddr} - } -} - type artKey []byte type artNode struct { @@ -161,7 +146,9 @@ func (an *artNode) node(a *artAllocator) *nodeBase { } } -// at returns the nth child of the node. +// at returns the nth child of the node, used for test. +// +//nolint:unused func (an *artNode) at(a *artAllocator, idx int) artNode { switch an.kind { case typeNode4: @@ -183,7 +170,6 @@ func (n48 *node48) init() { n48.nodeBase.init() // initialize node48 n48.present[0], n48.present[1], n48.present[2], n48.present[3] = 0, 0, 0, 0 - copy(n48.children[:], nullNode48.children[:]) } // nextPresentIdx returns the next present index starting from the given index. @@ -222,7 +208,7 @@ func (n256 *node256) init() { // initialize nodeBase n256.nodeBase.init() // initialize node256 - copy(n256.children[:], nullNode256.children[:]) + n256.present[0], n256.present[1], n256.present[2], n256.present[3] = 0, 0, 0, 0 } func (n256 *node256) nextPresentIdx(start int) int { @@ -302,6 +288,8 @@ const ( ) // markDelete marks the artLeaf as deleted +// +//nolint:unused func (l *artLeaf) markDelete() { l.flags = deleteFlag } @@ -470,7 +458,10 @@ func (an *artNode) findChild48(a *artAllocator, c byte) (int, artNode) { func (an *artNode) findChild256(a *artAllocator, c byte) (int, artNode) { n256 := an.node256(a) - return int(c), n256.children[c] + if n256.present[c>>n48s]&(1<<(c%n48m)) != 0 { + return int(c), n256.children[c] + } + return -2, nullArtNode } func (an *artNode) swapChild(a *artAllocator, c byte, child artNode) { diff --git a/internal/unionstore/art/art_node_test.go b/internal/unionstore/art/art_node_test.go index 1431ec1db..8d84080f0 100644 --- a/internal/unionstore/art/art_node_test.go +++ b/internal/unionstore/art/art_node_test.go @@ -21,9 +21,7 @@ func TestAllocNode(t *testing.T) { require.Equal(t, [4]uint64{0, 0, 0, 0}, n.present) case *node256: base = &n.nodeBase - for i := 0; i < node256cap; i++ { - require.Equal(t, n.children[i], nullArtNode) - } + require.Equal(t, [4]uint64{0, 0, 0, 0}, n.present) default: require.Fail(t, "unknown node type") } From f0a768d65369136df3a78605368b549f5b6eeabd Mon Sep 17 00:00:00 2001 From: you06 Date: Tue, 27 Aug 2024 15:15:25 +0900 Subject: [PATCH 4/7] fix large-indextest Signed-off-by: you06 --- internal/unionstore/memdb_bench_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/internal/unionstore/memdb_bench_test.go b/internal/unionstore/memdb_bench_test.go index 249b118a5..22296097b 100644 --- a/internal/unionstore/memdb_bench_test.go +++ b/internal/unionstore/memdb_bench_test.go @@ -48,14 +48,14 @@ const ( func BenchmarkLargeIndex(b *testing.B) { fn := func(b *testing.B, p MemBuffer) { - buf := make([][valueSize]byte, 10000000) + buf := make([][valueSize]byte, b.N) for i := range buf { binary.LittleEndian.PutUint32(buf[i][:], uint32(i)) } b.ResetTimer() for i := range buf { - p.Set(buf[i][:keySize], buf[i][:]) + p.Set(buf[i][:], buf[i][:]) } } From b6409381bdba2e9c5fe9ecc01fcfd51ca8f05eeb Mon Sep 17 00:00:00 2001 From: you06 Date: Wed, 4 Sep 2024 14:18:26 +0800 Subject: [PATCH 5/7] address comment Signed-off-by: you06 --- internal/unionstore/art/art.go | 54 +++++++++++++----------- internal/unionstore/art/art_node.go | 45 +++++++++++--------- internal/unionstore/art/art_node_test.go | 10 ++--- 3 files changed, 58 insertions(+), 51 deletions(-) diff --git a/internal/unionstore/art/art.go b/internal/unionstore/art/art.go index bb7a91165..a0d5e1759 100644 --- a/internal/unionstore/art/art.go +++ b/internal/unionstore/art/art.go @@ -48,7 +48,7 @@ func New() *ART { } func (t *ART) Get(key []byte) ([]byte, error) { - // 1. search the leaf node. + // 1. search the asLeaf node. _, leaf := t.search(key) if leaf == nil || leaf.vAddr.IsNull() { return nil, tikverr.ErrNotExist @@ -81,7 +81,7 @@ 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 exist leaf in the tree. + // 1. create or search the exist asLeaf in the tree. addr, leaf := t.recursiveInsert(key) // 2. set the value and flags. t.setValue(addr, leaf, value, ops) @@ -91,6 +91,9 @@ func (t *ART) Set(key artKey, value []byte, ops ...kv.FlagsOp) error { return nil } +// search looks up the asLeaf with the given key. +// It returns the address of asLeaf and asLeaf itself it there is a match asLeaf, +// returns arena.NullAddr and nil if the key is not found. func (t *ART) search(key artKey) (arena.MemdbArenaAddr, *artLeaf) { current := t.root if current == nullArtNode { @@ -100,7 +103,7 @@ func (t *ART) search(key artKey) (arena.MemdbArenaAddr, *artLeaf) { var node *nodeBase for { if current.isLeaf() { - lf := current.leaf(&t.allocator) + lf := current.asLeaf(&t.allocator) if lf.match(0, key) { return current.addr, lf } @@ -141,7 +144,7 @@ func (t *ART) search(key artKey) (arena.MemdbArenaAddr, *artLeaf) { } // recursiveInsert returns the node address of the key. -// if insert is true, it will insert the key if not exists, unless nullAddr is returned. +// It will insert the key if not exists, returns the newly inserted or exist leaf. func (t *ART) recursiveInsert(key artKey) (arena.MemdbArenaAddr, *artLeaf) { // lazy init root node and allocator. // this saves memory for read only txns. @@ -175,7 +178,7 @@ func (t *ART) recursiveInsert(key artKey) (arena.MemdbArenaAddr, *artLeaf) { } if node.prefixLen > 0 { - mismatchIdx := current.matchDeep(&t.allocator, key, depth) + mismatchIdx := node.matchDeep(&t.allocator, ¤t, key, depth) if mismatchIdx < node.prefixLen { // if the prefix doesn't match, we split the node into different prefixes. return t.expandNode(key, depth, mismatchIdx, prev, current, node) @@ -187,7 +190,7 @@ func (t *ART) recursiveInsert(key artKey) (arena.MemdbArenaAddr, *artLeaf) { valid := key.valid(int(depth)) _, next := current.findChild(&t.allocator, key.charAt(int(depth)), valid) if next == nullArtNode { - // insert as leaf if there is no child. + // insert as asLeaf if there is no child. newLeaf, lf := t.newLeaf(key) if current.addChild(&t.allocator, key.charAt(int(depth)), !key.valid(int(depth)), newLeaf) { if prev == nullArtNode { @@ -199,8 +202,8 @@ func (t *ART) recursiveInsert(key artKey) (arena.MemdbArenaAddr, *artLeaf) { return newLeaf.addr, lf } if !valid && next.kind == typeLeaf { - // key is drained, return the leaf. - return next.addr, next.leaf(&t.allocator) + // key is drained, return the asLeaf. + return next.addr, next.asLeaf(&t.allocator) } prev = current current = next @@ -210,14 +213,14 @@ func (t *ART) recursiveInsert(key artKey) (arena.MemdbArenaAddr, *artLeaf) { } } -// expandLeaf expands the exist artLeaf to a node4 if the keys are different. -// it returns the addr and leaf of the given key. +// expandLeaf expands the existing artLeaf to a node4 if the keys are different. +// it returns the addr and asLeaf of the given key. func (t *ART) expandLeaf(key artKey, depth uint32, prev, current artNode) (arena.MemdbArenaAddr, *artLeaf) { // Expand the artLeaf to a node4. // // ┌────────────┐ // │ new │ - // │ node4 │ + // │ node4 │ // ┌─────────┐ └──────┬─────┘ // │ old │ ---> │ // │ leaf1 │ ┌────────┴────────┐ @@ -226,7 +229,7 @@ func (t *ART) expandLeaf(key artKey, depth uint32, prev, current artNode) (arena // │ old │ │ new │ // │ leaf1 │ │ leaf2 │ // └─────────┘ └─────────┘ - leaf1 := current.leaf(&t.allocator) + leaf1 := current.asLeaf(&t.allocator) if leaf1.match(depth-1, key) { // same key, return the artLeaf and overwrite the value. return current.addr, leaf1 @@ -244,7 +247,7 @@ func (t *ART) expandLeaf(key artKey, depth uint32, prev, current artNode) (arena an.addChild(&t.allocator, l1Key.charAt(int(depth)), !l1Key.valid(int(depth)), current) an.addChild(&t.allocator, l2Key.charAt(int(depth)), !l2Key.valid(int(depth)), leaf2Addr) - // swap the old leaf with the new node4. + // swap the old asLeaf with the new node4. if prev == nullArtNode { t.root = an } else { @@ -263,7 +266,7 @@ func (t *ART) expandNode(key artKey, depth, mismatchIdx uint32, prev, current ar // ┌─────────────┐ ┌── b ───┴── c ───┐ // │ node4 │ ---> │ │ // │ prefix: abc │ ┌──────▼─────┐ ┌──────▼─────┐ - // └─────────────┘ │ old node4 │ │ new leaf │ + // └─────────────┘ │ old node4 │ │ new asLeaf │ // │ prefix: c │ │ key: acc │ // └────────────┘ └────────────┘ prevDepth := int(depth - 1) @@ -281,7 +284,7 @@ func (t *ART) expandNode(key artKey, depth, mismatchIdx uint32, prev, current ar } else { currNode.prefixLen -= mismatchIdx + 1 leafArtNode := minimum(&t.allocator, current) - leaf := leafArtNode.leaf(&t.allocator) + leaf := leafArtNode.asLeaf(&t.allocator) leafKey := artKey(leaf.GetKey()) kMin := depth + mismatchIdx + 1 kMax := depth + mismatchIdx + 1 + min(currNode.prefixLen, maxPrefixLen) @@ -330,32 +333,33 @@ func (t *ART) setValue(addr arena.MemdbArenaAddr, l *artLeaf, value []byte, ops // value == nil means it updates flags only. return } - if t.trySwapValue(l.vAddr, value) { + oldSize, swapper := t.trySwapValue(l.vAddr, value) + if swapper { return } - t.size += len(value) + t.size += len(value) - oldSize vAddr := t.allocator.vlogAllocator.AppendValue(addr, l.vAddr, value) l.vAddr = vAddr } -// trySwapValue checks if the value can be updated in place, return true if it's updated. -func (t *ART) trySwapValue(addr arena.MemdbArenaAddr, value []byte) bool { +// trySwapValue checks if the value can be updated in place. +// It returns 0 and true if it's updated, returns the size of old value and false if it cannot be updated in place. +func (t *ART) trySwapValue(addr arena.MemdbArenaAddr, value []byte) (int, bool) { if addr.IsNull() { - return false + return 0, false } + oldVal := t.allocator.vlogAllocator.GetValue(addr) if len(t.stages) > 0 { cp := t.stages[len(t.stages)-1] if !t.allocator.vlogAllocator.CanModify(&cp, addr) { - return false + return len(oldVal), false } } - oldVal := t.allocator.vlogAllocator.GetValue(addr) if len(oldVal) > 0 && len(oldVal) == len(value) { copy(oldVal, value) - return true + return 0, true } - t.size -= len(oldVal) - return false + return len(oldVal), false } func (t *ART) Dirty() bool { diff --git a/internal/unionstore/art/art_node.go b/internal/unionstore/art/art_node.go index 53e5e4f3f..707f22201 100644 --- a/internal/unionstore/art/art_node.go +++ b/internal/unionstore/art/art_node.go @@ -126,7 +126,7 @@ func (an *artNode) isLeaf() bool { return an.kind == typeLeaf } -func (an *artNode) leaf(a *artAllocator) *artLeaf { +func (an *artNode) asLeaf(a *artAllocator) *artLeaf { return a.getLeaf(an.addr) } @@ -252,7 +252,7 @@ func (k artKey) valid(pos int) bool { return pos < len(k) } -// GetKey gets the full key of the leaf +// GetKey gets the full key of the asLeaf func (l *artLeaf) GetKey() []byte { base := unsafe.Add(unsafe.Pointer(l), leafSize) return unsafe.Slice((*byte)(base), int(l.klen)) @@ -264,7 +264,7 @@ func (l *artLeaf) getKeyDepth(depth uint32) []byte { return unsafe.Slice((*byte)(base), int(l.klen)-int(depth)) } -// GetKeyFlags gets the flags of the leaf +// GetKeyFlags gets the flags of the asLeaf func (l *artLeaf) GetKeyFlags() kv.KeyFlags { panic("unimplemented") } @@ -273,9 +273,8 @@ func (l *artLeaf) match(depth uint32, key artKey) bool { return bytes.Equal(l.getKeyDepth(depth), key[depth:]) } -func (l *artLeaf) setKeyFlags(flags kv.KeyFlags) arena.MemdbArenaAddr { +func (l *artLeaf) setKeyFlags(flags kv.KeyFlags) { l.flags = uint16(flags) & flagMask - return l.vAddr } func (l *artLeaf) getKeyFlags() kv.KeyFlags { @@ -309,19 +308,36 @@ func (n *nodeBase) setPrefix(key artKey, prefixLen uint32) { copy(n.prefix[:], key[:min(prefixLen, maxPrefixLen)]) } -func (n *nodeBase) match(key artKey, depth uint32) uint32 { +// match returns the mismatch index of the key and the node's prefix within maxPrefixLen. +// Node if the nodeBase.prefixLen > maxPrefixLen and the returned mismatch index equals to maxPrefixLen, +// key[maxPrefixLen:] will not be checked by this function. +func (n *nodeBase) match(key artKey, depth uint32) uint32 /* mismatch index */ { idx := uint32(0) - limit := min(min(n.prefixLen, maxPrefixLen), uint32(len(key))-depth) for ; idx < limit; idx++ { if n.prefix[idx] != key[idx+depth] { return idx } } - return idx } +// matchDeep returns the mismatch index of the key and the node's prefix. +// If the key is fully match, the mismatch index is equal to the nodeBase.prefixLen. +// The nodeBase only stores prefix within maxPrefixLen, if it's not enough for matching, +// the matchDeep func looks up and match by the leaf, this function always returns the actual mismatch index. +func (n *nodeBase) matchDeep(a *artAllocator, an *artNode, key artKey, depth uint32) uint32 /* mismatch index */ { + // match in-node prefix + mismatchIdx := n.match(key, depth) + if mismatchIdx < maxPrefixLen || n.prefixLen <= maxPrefixLen { + return mismatchIdx + } + // if the prefixLen is longer maxPrefixLen and mismatchIdx == maxPrefixLen, we need to match deeper with any asLeaf. + leafArtNode := minimum(a, *an) + lKey := leafArtNode.asLeaf(a).GetKey() + return longestCommonPrefix(lKey, key, depth+maxPrefixLen) + maxPrefixLen +} + func (an *artNode) node4(a *artAllocator) *node4 { return a.getNode4(an.addr) } @@ -338,19 +354,6 @@ func (an *artNode) node256(a *artAllocator) *node256 { return a.getNode256(an.addr) } -func (an *artNode) matchDeep(a *artAllocator, key artKey, depth uint32) uint32 /* mismatch index*/ { - n := an.node(a) - // match in-node prefix - mismatchIdx := n.match(key, depth) - if mismatchIdx < maxPrefixLen || n.prefixLen <= maxPrefixLen { - return mismatchIdx - } - // if the prefixLen is longer maxPrefixLen and mismatchIdx == maxPrefixLen, we need to match deeper with any leaf. - leafArtNode := minimum(a, *an) - lKey := leafArtNode.leaf(a).GetKey() - return longestCommonPrefix(lKey, key, depth+maxPrefixLen) + maxPrefixLen -} - func longestCommonPrefix(l1Key, l2Key artKey, depth uint32) uint32 { idx, limit := depth, min(uint32(len(l1Key)), uint32(len(l2Key))) // TODO: possible optimization diff --git a/internal/unionstore/art/art_node_test.go b/internal/unionstore/art/art_node_test.go index 8d84080f0..146275061 100644 --- a/internal/unionstore/art/art_node_test.go +++ b/internal/unionstore/art/art_node_test.go @@ -82,7 +82,7 @@ func TestAllocNode(t *testing.T) { n256s = append(n256s, addr) } - // alloc leaf + // alloc asLeaf leafs := make([]arena.MemdbArenaAddr, 0, cnt) for i := 0; i < cnt; i++ { key := []byte(strconv.Itoa(i)) @@ -149,8 +149,8 @@ func TestNodePrefix(t *testing.T) { n.setPrefix(leafKey, maxPrefixLen+1) matchKey := append(make([]byte, maxPrefixLen), []byte{1, 22, 33, 44, 55}...) mismatchKey := append(make([]byte, maxPrefixLen), []byte{11, 22, 33, 44, 55}...) - require.Equal(t, uint32(maxPrefixLen+1), an.matchDeep(&allocator, matchKey, 0)) - require.Equal(t, uint32(maxPrefixLen), an.matchDeep(&allocator, mismatchKey, 0)) + require.Equal(t, uint32(maxPrefixLen+1), n.matchDeep(&allocator, an, matchKey, 0)) + require.Equal(t, uint32(maxPrefixLen), n.matchDeep(&allocator, an, mismatchKey, 0)) // deep match with depth leafKey = append(make([]byte, 10), leafKey...) @@ -159,8 +159,8 @@ func TestNodePrefix(t *testing.T) { leafAddr, _ = allocator.allocLeaf(leafKey) an.swapChild(&allocator, 2, artNode{kind: typeLeaf, addr: leafAddr}) n.setPrefix(leafKey[10:], maxPrefixLen+1) - require.Equal(t, uint32(maxPrefixLen+1), an.matchDeep(&allocator, matchKey, 10)) - require.Equal(t, uint32(maxPrefixLen), an.matchDeep(&allocator, mismatchKey, 10)) + require.Equal(t, uint32(maxPrefixLen+1), n.matchDeep(&allocator, an, matchKey, 10)) + require.Equal(t, uint32(maxPrefixLen), n.matchDeep(&allocator, an, mismatchKey, 10)) } addr, n4 := allocator.allocNode4() From 7ea496b955cb418feaeecb06a688a2bf4772941b Mon Sep 17 00:00:00 2001 From: you06 Date: Wed, 4 Sep 2024 16:13:38 +0800 Subject: [PATCH 6/7] fix comment Signed-off-by: you06 --- internal/unionstore/art/art.go | 36 ++++++------ internal/unionstore/art/art_node.go | 70 ++++++++++++------------ internal/unionstore/art/art_node_test.go | 10 ++-- 3 files changed, 58 insertions(+), 58 deletions(-) diff --git a/internal/unionstore/art/art.go b/internal/unionstore/art/art.go index a0d5e1759..343c9330e 100644 --- a/internal/unionstore/art/art.go +++ b/internal/unionstore/art/art.go @@ -48,7 +48,7 @@ func New() *ART { } func (t *ART) Get(key []byte) ([]byte, error) { - // 1. search the asLeaf node. + // 1. search the leaf node. _, leaf := t.search(key) if leaf == nil || leaf.vAddr.IsNull() { return nil, tikverr.ErrNotExist @@ -81,7 +81,7 @@ 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 exist asLeaf in the tree. + // 1. create or search the existing leaf in the tree. addr, leaf := t.recursiveInsert(key) // 2. set the value and flags. t.setValue(addr, leaf, value, ops) @@ -91,8 +91,8 @@ func (t *ART) Set(key artKey, value []byte, ops ...kv.FlagsOp) error { return nil } -// search looks up the asLeaf with the given key. -// It returns the address of asLeaf and asLeaf itself it there is a match asLeaf, +// 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) search(key artKey) (arena.MemdbArenaAddr, *artLeaf) { current := t.root @@ -114,13 +114,13 @@ func (t *ART) search(key artKey) (arena.MemdbArenaAddr, *artLeaf) { // get the basic node information. switch current.kind { case typeNode4: - node = ¤t.node4(&t.allocator).nodeBase + node = ¤t.asNode4(&t.allocator).nodeBase case typeNode16: - node = ¤t.node16(&t.allocator).nodeBase + node = ¤t.asNode16(&t.allocator).nodeBase case typeNode48: - node = ¤t.node48(&t.allocator).nodeBase + node = ¤t.asNode48(&t.allocator).nodeBase case typeNode256: - node = ¤t.node256(&t.allocator).nodeBase + node = ¤t.asNode256(&t.allocator).nodeBase default: panic("invalid nodeBase kind") } @@ -144,7 +144,7 @@ func (t *ART) search(key artKey) (arena.MemdbArenaAddr, *artLeaf) { } // recursiveInsert returns the node address of the key. -// It will insert the key if not exists, returns the newly inserted or exist leaf. +// It will insert the key if not exists, returns the newly inserted or existing leaf. func (t *ART) recursiveInsert(key artKey) (arena.MemdbArenaAddr, *artLeaf) { // lazy init root node and allocator. // this saves memory for read only txns. @@ -166,13 +166,13 @@ func (t *ART) recursiveInsert(key artKey) (arena.MemdbArenaAddr, *artLeaf) { // get the basic node information. switch current.kind { case typeNode4: - node = ¤t.node4(&t.allocator).nodeBase + node = ¤t.asNode4(&t.allocator).nodeBase case typeNode16: - node = ¤t.node16(&t.allocator).nodeBase + node = ¤t.asNode16(&t.allocator).nodeBase case typeNode48: - node = ¤t.node48(&t.allocator).nodeBase + node = ¤t.asNode48(&t.allocator).nodeBase case typeNode256: - node = ¤t.node256(&t.allocator).nodeBase + node = ¤t.asNode256(&t.allocator).nodeBase default: panic("invalid nodeBase kind") } @@ -190,7 +190,7 @@ func (t *ART) recursiveInsert(key artKey) (arena.MemdbArenaAddr, *artLeaf) { valid := key.valid(int(depth)) _, next := current.findChild(&t.allocator, key.charAt(int(depth)), valid) if next == nullArtNode { - // insert as asLeaf if there is no child. + // insert as leaf if there is no child. newLeaf, lf := t.newLeaf(key) if current.addChild(&t.allocator, key.charAt(int(depth)), !key.valid(int(depth)), newLeaf) { if prev == nullArtNode { @@ -202,7 +202,7 @@ func (t *ART) recursiveInsert(key artKey) (arena.MemdbArenaAddr, *artLeaf) { return newLeaf.addr, lf } if !valid && next.kind == typeLeaf { - // key is drained, return the asLeaf. + // key is drained, return the leaf. return next.addr, next.asLeaf(&t.allocator) } prev = current @@ -214,7 +214,7 @@ func (t *ART) recursiveInsert(key artKey) (arena.MemdbArenaAddr, *artLeaf) { } // expandLeaf expands the existing artLeaf to a node4 if the keys are different. -// it returns the addr and asLeaf of the given key. +// it returns the addr and leaf of the given key. func (t *ART) expandLeaf(key artKey, depth uint32, prev, current artNode) (arena.MemdbArenaAddr, *artLeaf) { // Expand the artLeaf to a node4. // @@ -247,7 +247,7 @@ func (t *ART) expandLeaf(key artKey, depth uint32, prev, current artNode) (arena an.addChild(&t.allocator, l1Key.charAt(int(depth)), !l1Key.valid(int(depth)), current) an.addChild(&t.allocator, l2Key.charAt(int(depth)), !l2Key.valid(int(depth)), leaf2Addr) - // swap the old asLeaf with the new node4. + // swap the old leaf with the new node4. if prev == nullArtNode { t.root = an } else { @@ -266,7 +266,7 @@ func (t *ART) expandNode(key artKey, depth, mismatchIdx uint32, prev, current ar // ┌─────────────┐ ┌── b ───┴── c ───┐ // │ node4 │ ---> │ │ // │ prefix: abc │ ┌──────▼─────┐ ┌──────▼─────┐ - // └─────────────┘ │ old node4 │ │ new asLeaf │ + // └─────────────┘ │ old node4 │ │ new leaf │ // │ prefix: c │ │ key: acc │ // └────────────┘ └────────────┘ prevDepth := int(depth - 1) diff --git a/internal/unionstore/art/art_node.go b/internal/unionstore/art/art_node.go index 707f22201..4258074c9 100644 --- a/internal/unionstore/art/art_node.go +++ b/internal/unionstore/art/art_node.go @@ -131,16 +131,16 @@ func (an *artNode) asLeaf(a *artAllocator) *artLeaf { } // node gets the inner baseNode of the artNode -func (an *artNode) node(a *artAllocator) *nodeBase { +func (an *artNode) asNode(a *artAllocator) *nodeBase { switch an.kind { case typeNode4: - return &an.node4(a).nodeBase + return &an.asNode4(a).nodeBase case typeNode16: - return &an.node16(a).nodeBase + return &an.asNode16(a).nodeBase case typeNode48: - return &an.node48(a).nodeBase + return &an.asNode48(a).nodeBase case typeNode256: - return &an.node256(a).nodeBase + return &an.asNode256(a).nodeBase default: panic("invalid nodeBase kind") } @@ -152,14 +152,14 @@ func (an *artNode) node(a *artAllocator) *nodeBase { func (an *artNode) at(a *artAllocator, idx int) artNode { switch an.kind { case typeNode4: - return an.node4(a).children[idx] + return an.asNode4(a).children[idx] case typeNode16: - return an.node16(a).children[idx] + return an.asNode16(a).children[idx] case typeNode48: - n48 := an.node48(a) + n48 := an.asNode48(a) return n48.children[n48.keys[idx]] case typeNode256: - return an.node256(a).children[idx] + return an.asNode256(a).children[idx] default: panic("invalid nodeBase kind") } @@ -252,7 +252,7 @@ func (k artKey) valid(pos int) bool { return pos < len(k) } -// GetKey gets the full key of the asLeaf +// GetKey gets the full key of the leaf func (l *artLeaf) GetKey() []byte { base := unsafe.Add(unsafe.Pointer(l), leafSize) return unsafe.Slice((*byte)(base), int(l.klen)) @@ -264,7 +264,7 @@ func (l *artLeaf) getKeyDepth(depth uint32) []byte { return unsafe.Slice((*byte)(base), int(l.klen)-int(depth)) } -// GetKeyFlags gets the flags of the asLeaf +// GetKeyFlags gets the flags of the leaf func (l *artLeaf) GetKeyFlags() kv.KeyFlags { panic("unimplemented") } @@ -332,25 +332,25 @@ func (n *nodeBase) matchDeep(a *artAllocator, an *artNode, key artKey, depth uin if mismatchIdx < maxPrefixLen || n.prefixLen <= maxPrefixLen { return mismatchIdx } - // if the prefixLen is longer maxPrefixLen and mismatchIdx == maxPrefixLen, we need to match deeper with any asLeaf. + // if the prefixLen is longer maxPrefixLen and mismatchIdx == maxPrefixLen, we need to match deeper with any leaf. leafArtNode := minimum(a, *an) lKey := leafArtNode.asLeaf(a).GetKey() return longestCommonPrefix(lKey, key, depth+maxPrefixLen) + maxPrefixLen } -func (an *artNode) node4(a *artAllocator) *node4 { +func (an *artNode) asNode4(a *artAllocator) *node4 { return a.getNode4(an.addr) } -func (an *artNode) node16(a *artAllocator) *node16 { +func (an *artNode) asNode16(a *artAllocator) *node16 { return a.getNode16(an.addr) } -func (an *artNode) node48(a *artAllocator) *node48 { +func (an *artNode) asNode48(a *artAllocator) *node48 { return a.getNode48(an.addr) } -func (an *artNode) node256(a *artAllocator) *node256 { +func (an *artNode) asNode256(a *artAllocator) *node256 { return a.getNode256(an.addr) } @@ -374,26 +374,26 @@ func minimum(a *artAllocator, an artNode) artNode { case typeLeaf: return an case typeNode4: - n4 := an.node4(a) + n4 := an.asNode4(a) if !n4.inplaceLeaf.addr.IsNull() { return n4.inplaceLeaf } an = n4.children[0] case typeNode16: - n16 := an.node16(a) + n16 := an.asNode16(a) if !n16.inplaceLeaf.addr.IsNull() { return n16.inplaceLeaf } an = n16.children[0] case typeNode48: - n48 := an.node48(a) + n48 := an.asNode48(a) if !n48.inplaceLeaf.addr.IsNull() { return n48.inplaceLeaf } idx := n48.nextPresentIdx(0) an = n48.children[n48.keys[idx]] case typeNode256: - n256 := an.node256(a) + n256 := an.asNode256(a) if !n256.inplaceLeaf.addr.IsNull() { return n256.inplaceLeaf } @@ -407,7 +407,7 @@ func minimum(a *artAllocator, an artNode) artNode { func (an *artNode) findChild(a *artAllocator, c byte, valid bool) (int, artNode) { if !valid { - return inplaceIndex, an.node(a).inplaceLeaf + return inplaceIndex, an.asNode(a).inplaceLeaf } switch an.kind { case typeNode4: @@ -423,7 +423,7 @@ func (an *artNode) findChild(a *artAllocator, c byte, valid bool) (int, artNode) } func (an *artNode) findChild4(a *artAllocator, c byte) (int, artNode) { - n4 := an.node4(a) + n4 := an.asNode4(a) for idx := 0; idx < int(n4.nodeNum); idx++ { if n4.keys[idx] == c { return idx, n4.children[idx] @@ -434,7 +434,7 @@ func (an *artNode) findChild4(a *artAllocator, c byte) (int, artNode) { } func (an *artNode) findChild16(a *artAllocator, c byte) (int, artNode) { - n16 := an.node16(a) + n16 := an.asNode16(a) idx, found := sort.Find(int(n16.nodeNum), func(i int) int { if n16.keys[i] < c { @@ -452,7 +452,7 @@ func (an *artNode) findChild16(a *artAllocator, c byte) (int, artNode) { } func (an *artNode) findChild48(a *artAllocator, c byte) (int, artNode) { - n48 := an.node48(a) + n48 := an.asNode48(a) if n48.present[c>>n48s]&(1<<(c%n48m)) != 0 { return int(c), n48.children[n48.keys[c]] } @@ -460,7 +460,7 @@ func (an *artNode) findChild48(a *artAllocator, c byte) (int, artNode) { } func (an *artNode) findChild256(a *artAllocator, c byte) (int, artNode) { - n256 := an.node256(a) + n256 := an.asNode256(a) if n256.present[c>>n48s]&(1<<(c%n48m)) != 0 { return int(c), n256.children[c] } @@ -470,7 +470,7 @@ func (an *artNode) findChild256(a *artAllocator, c byte) (int, artNode) { func (an *artNode) swapChild(a *artAllocator, c byte, child artNode) { switch an.kind { case typeNode4: - n4 := an.node4(a) + n4 := an.asNode4(a) for idx := uint8(0); idx < n4.nodeNum; idx++ { if n4.keys[idx] == c { n4.children[idx] = child @@ -479,7 +479,7 @@ func (an *artNode) swapChild(a *artAllocator, c byte, child artNode) { } panic("swap child failed") case typeNode16: - n16 := an.node16(a) + n16 := an.asNode16(a) for idx := uint8(0); idx < n16.nodeNum; idx++ { if n16.keys[idx] == c { n16.children[idx] = child @@ -488,7 +488,7 @@ func (an *artNode) swapChild(a *artAllocator, c byte, child artNode) { } panic("swap child failed") case typeNode48: - n48 := an.node48(a) + n48 := an.asNode48(a) if n48.present[c>>n48s]&(1<<(c%n48m)) != 0 { n48.children[n48.keys[c]] = child return @@ -517,7 +517,7 @@ func (an *artNode) addChild(a *artAllocator, c byte, inplace bool, child artNode } func (an *artNode) addChild4(a *artAllocator, c byte, inplace bool, child artNode) bool { - node := an.node4(a) + node := an.asNode4(a) if inplace { node.inplaceLeaf = child @@ -550,7 +550,7 @@ func (an *artNode) addChild4(a *artAllocator, c byte, inplace bool, child artNod } func (an *artNode) addChild16(a *artAllocator, c byte, inplace bool, child artNode) bool { - node := an.node16(a) + node := an.asNode16(a) if inplace { node.inplaceLeaf = child @@ -585,7 +585,7 @@ func (an *artNode) addChild16(a *artAllocator, c byte, inplace bool, child artNo } func (an *artNode) addChild48(a *artAllocator, c byte, inplace bool, child artNode) bool { - node := an.node48(a) + node := an.asNode48(a) if inplace { node.inplaceLeaf = child @@ -606,7 +606,7 @@ func (an *artNode) addChild48(a *artAllocator, c byte, inplace bool, child artNo } func (an *artNode) addChild256(a *artAllocator, c byte, inplace bool, child artNode) bool { - node := an.node256(a) + node := an.asNode256(a) if inplace { node.inplaceLeaf = child @@ -629,7 +629,7 @@ func (n *nodeBase) copyMeta(src *nodeBase) { func (an *artNode) grow(a *artAllocator) { switch an.kind { case typeNode4: - n4 := an.node4(a) + n4 := an.asNode4(a) newAddr, n16 := a.allocNode16() n16.copyMeta(&n4.nodeBase) @@ -641,7 +641,7 @@ func (an *artNode) grow(a *artAllocator) { an.kind = typeNode16 an.addr = newAddr case typeNode16: - n16 := an.node16(a) + n16 := an.asNode16(a) newAddr, n48 := a.allocNode48() n48.copyMeta(&n16.nodeBase) @@ -657,7 +657,7 @@ func (an *artNode) grow(a *artAllocator) { an.kind = typeNode48 an.addr = newAddr case typeNode48: - n48 := an.node48(a) + n48 := an.asNode48(a) newAddr, n256 := a.allocNode256() n256.copyMeta(&n48.nodeBase) diff --git a/internal/unionstore/art/art_node_test.go b/internal/unionstore/art/art_node_test.go index 146275061..88a47ce72 100644 --- a/internal/unionstore/art/art_node_test.go +++ b/internal/unionstore/art/art_node_test.go @@ -82,7 +82,7 @@ func TestAllocNode(t *testing.T) { n256s = append(n256s, addr) } - // alloc asLeaf + // alloc leaf leafs := make([]arena.MemdbArenaAddr, 0, cnt) for i := 0; i < cnt; i++ { key := []byte(strconv.Itoa(i)) @@ -185,7 +185,7 @@ func TestOrderedChild(t *testing.T) { addr, _ := allocator.allocNode4() an.addChild(&allocator, byte(i), false, artNode{kind: typeNode4, addr: addr}) - node := an.node(&allocator) + node := an.asNode(&allocator) require.Equal(t, node.nodeNum, uint8(i+1)) keys := getKeys(an) for j := 0; j <= i; j++ { @@ -199,7 +199,7 @@ func TestOrderedChild(t *testing.T) { addr, _ := allocator.allocNode4() an.addChild(&allocator, byte(255-i), false, artNode{kind: typeNode4, addr: addr}) - node := an.node(&allocator) + node := an.asNode(&allocator) require.Equal(t, node.nodeNum, uint8(i+1)) keys := getKeys(an) for j := 0; j <= i; j++ { @@ -212,14 +212,14 @@ func TestOrderedChild(t *testing.T) { addr, _ := allocator.allocNode4() return &artNode{kind: typeNode4, addr: addr} }, func(an *artNode) []byte { - return an.node4(&allocator).keys[:] + return an.asNode4(&allocator).keys[:] }, 4) testFn(func() *artNode { addr, _ := allocator.allocNode16() return &artNode{kind: typeNode16, addr: addr} }, func(an *artNode) []byte { - return an.node16(&allocator).keys[:] + return an.asNode16(&allocator).keys[:] }, 16) } From eb44430e15b619f9465715f9af15880a318feba9 Mon Sep 17 00:00:00 2001 From: you06 Date: Thu, 12 Sep 2024 14:13:02 +0800 Subject: [PATCH 7/7] address comment Signed-off-by: you06 address comment Signed-off-by: you06 --- internal/unionstore/art/art.go | 46 +++--- internal/unionstore/art/art_node.go | 166 +++++++++++---------- internal/unionstore/art/art_node_test.go | 10 +- internal/unionstore/art/art_test.go | 179 ++++++++++++++--------- 4 files changed, 237 insertions(+), 164 deletions(-) diff --git a/internal/unionstore/art/art.go b/internal/unionstore/art/art.go index 343c9330e..17e826327 100644 --- a/internal/unionstore/art/art.go +++ b/internal/unionstore/art/art.go @@ -23,6 +23,16 @@ import ( "github.com/tikv/client-go/v2/kv" ) +var testMode = false + +// ART is rollbackable Adaptive Radix Tree optimized for TiDB's transaction states buffer use scenario. +// You can think ART is a combination of two separate tree map, one for key => value and another for key => keyFlags. +// +// The value map is rollbackable, that means you can use the `Staging`, `Release` and `Cleanup` API to safely modify KVs. +// +// The flags map is not rollbackable. There are two types of flag, persistent and non-persistent. +// When discarding a newly added KV in `Cleanup`, the non-persistent flags will be cleared. +// If there are persistent flags associated with key, we will keep this key in node without value. type ART struct { allocator artAllocator root artNode @@ -135,7 +145,7 @@ func (t *ART) search(key artKey) (arena.MemdbArenaAddr, *artLeaf) { depth += node.prefixLen } - _, current = current.findChild(&t.allocator, key.charAt(int(depth)), key.valid(int(depth))) + _, current = current.findChild(&t.allocator, key.charAt(int(depth)), !key.valid(int(depth))) if current.addr.IsNull() { return arena.NullAddr, nil } @@ -188,18 +198,18 @@ func (t *ART) recursiveInsert(key artKey) (arena.MemdbArenaAddr, *artLeaf) { // search next node valid := key.valid(int(depth)) - _, next := current.findChild(&t.allocator, key.charAt(int(depth)), valid) + _, next := current.findChild(&t.allocator, key.charAt(int(depth)), !valid) if next == nullArtNode { // insert as leaf if there is no child. - newLeaf, lf := t.newLeaf(key) - if current.addChild(&t.allocator, key.charAt(int(depth)), !key.valid(int(depth)), newLeaf) { + newAn, newLeaf := t.newLeaf(key) + if current.addChild(&t.allocator, key.charAt(int(depth)), !valid, newAn) { if prev == nullArtNode { t.root = current } else { - prev.swapChild(&t.allocator, key.charAt(prevDepth), current) + prev.replaceChild(&t.allocator, key.charAt(prevDepth), current) } } - return newLeaf.addr, lf + return newAn.addr, newLeaf } if !valid && next.kind == typeLeaf { // key is drained, return the leaf. @@ -241,17 +251,17 @@ func (t *ART) expandLeaf(key artKey, depth uint32, prev, current artNode) (arena lcp := longestCommonPrefix(l1Key, l2Key, depth) // calculate the common prefix length of new node. - an, n4 := t.newNode4() - n4.setPrefix(key[depth:], lcp) + newAn, newN4 := t.newNode4() + newN4.setPrefix(key[depth:], lcp) depth += lcp - an.addChild(&t.allocator, l1Key.charAt(int(depth)), !l1Key.valid(int(depth)), current) - an.addChild(&t.allocator, l2Key.charAt(int(depth)), !l2Key.valid(int(depth)), leaf2Addr) + newAn.addChild(&t.allocator, l1Key.charAt(int(depth)), !l1Key.valid(int(depth)), current) + 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 = an + t.root = newAn } else { - prev.swapChild(&t.allocator, key.charAt(prevDepth), an) + prev.replaceChild(&t.allocator, key.charAt(prevDepth), newAn) } return leaf2Addr.addr, leaf2 } @@ -272,7 +282,7 @@ func (t *ART) expandNode(key artKey, depth, mismatchIdx uint32, prev, current ar prevDepth := int(depth - 1) // set prefix for new node. - newArtNode, newN4 := t.newNode4() + newAn, newN4 := t.newNode4() newN4.setPrefix(key[depth:], mismatchIdx) // update prefix for old node and move it as a child of the new node. @@ -280,7 +290,7 @@ func (t *ART) expandNode(key artKey, depth, mismatchIdx uint32, prev, current ar nodeKey := currNode.prefix[mismatchIdx] currNode.prefixLen -= mismatchIdx + 1 copy(currNode.prefix[:], currNode.prefix[mismatchIdx+1:]) - newArtNode.addChild(&t.allocator, nodeKey, false, current) + newAn.addChild(&t.allocator, nodeKey, false, current) } else { currNode.prefixLen -= mismatchIdx + 1 leafArtNode := minimum(&t.allocator, current) @@ -289,16 +299,16 @@ func (t *ART) expandNode(key artKey, depth, mismatchIdx uint32, prev, current ar kMin := depth + mismatchIdx + 1 kMax := depth + mismatchIdx + 1 + min(currNode.prefixLen, maxPrefixLen) copy(currNode.prefix[:], leafKey[kMin:kMax]) - newArtNode.addChild(&t.allocator, leafKey.charAt(int(depth+mismatchIdx)), !leafKey.valid(int(depth)), current) + newAn.addChild(&t.allocator, leafKey.charAt(int(depth+mismatchIdx)), !leafKey.valid(int(depth)), current) } // insert the artLeaf into new node newLeafAddr, newLeaf := t.newLeaf(key) - newArtNode.addChild(&t.allocator, key.charAt(int(depth+mismatchIdx)), !key.valid(int(depth+mismatchIdx)), newLeafAddr) + newAn.addChild(&t.allocator, key.charAt(int(depth+mismatchIdx)), !key.valid(int(depth+mismatchIdx)), newLeafAddr) if prev == nullArtNode { - t.root = newArtNode + t.root = newAn } else { - prev.swapChild(&t.allocator, key.charAt(prevDepth), newArtNode) + prev.replaceChild(&t.allocator, key.charAt(prevDepth), newAn) } return newLeafAddr.addr, newLeaf } diff --git a/internal/unionstore/art/art_node.go b/internal/unionstore/art/art_node.go index 4258074c9..0c4ee9090 100644 --- a/internal/unionstore/art/art_node.go +++ b/internal/unionstore/art/art_node.go @@ -19,6 +19,7 @@ import ( "math" "math/bits" "sort" + "testing" "unsafe" "github.com/tikv/client-go/v2/internal/unionstore/arena" @@ -354,6 +355,8 @@ func (an *artNode) asNode256(a *artAllocator) *node256 { return a.getNode256(an.addr) } +// 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))) // TODO: possible optimization @@ -405,37 +408,52 @@ func minimum(a *artAllocator, an artNode) artNode { } } -func (an *artNode) findChild(a *artAllocator, c byte, valid bool) (int, artNode) { - if !valid { +// findChild finds the child node by the given key byte, returns the index of the child and the child node. +// inplace indicates whether the target key is a leaf of the node, if valid is false, the in-place leaf will be returned. +func (an *artNode) findChild(a *artAllocator, c byte, inplace bool) (int, artNode) { + if inplace { return inplaceIndex, an.asNode(a).inplaceLeaf } switch an.kind { case typeNode4: - return an.findChild4(a, c) + n4 := an.asNode4(a) + idx := n4.findChild(c) + if idx != notExistIndex { + return idx, n4.children[idx] + } case typeNode16: - return an.findChild16(a, c) + n16 := an.asNode16(a) + idx := n16.findChild(c) + if idx != notExistIndex { + return idx, n16.children[idx] + } case typeNode48: - return an.findChild48(a, c) + n48 := an.asNode48(a) + idx := n48.findChild(c) + if idx != notExistIndex { + return idx, n48.children[idx] + } case typeNode256: - return an.findChild256(a, c) + n256 := an.asNode256(a) + idx := n256.findChild(c) + if idx != notExistIndex { + return idx, n256.children[idx] + } } return notExistIndex, nullArtNode } -func (an *artNode) findChild4(a *artAllocator, c byte) (int, artNode) { - n4 := an.asNode4(a) +func (n4 *node4) findChild(c byte) int { for idx := 0; idx < int(n4.nodeNum); idx++ { if n4.keys[idx] == c { - return idx, n4.children[idx] + return idx } } - return -2, nullArtNode + return notExistIndex } -func (an *artNode) findChild16(a *artAllocator, c byte) (int, artNode) { - n16 := an.asNode16(a) - +func (n16 *node16) findChild(c byte) int { idx, found := sort.Find(int(n16.nodeNum), func(i int) int { if n16.keys[i] < c { return 1 @@ -446,102 +464,101 @@ func (an *artNode) findChild16(a *artAllocator, c byte) (int, artNode) { return -1 }) if found { - return idx, n16.children[idx] + return idx } - return -2, nullArtNode + return notExistIndex } -func (an *artNode) findChild48(a *artAllocator, c byte) (int, artNode) { - n48 := an.asNode48(a) +func (n48 *node48) findChild(c byte) int { if n48.present[c>>n48s]&(1<<(c%n48m)) != 0 { - return int(c), n48.children[n48.keys[c]] + return int(n48.keys[c]) } - return -2, nullArtNode + return notExistIndex } -func (an *artNode) findChild256(a *artAllocator, c byte) (int, artNode) { - n256 := an.asNode256(a) +func (n256 *node256) findChild(c byte) int { if n256.present[c>>n48s]&(1<<(c%n48m)) != 0 { - return int(c), n256.children[c] + return int(c) } - return -2, nullArtNode + return notExistIndex } -func (an *artNode) swapChild(a *artAllocator, c byte, child artNode) { +func (an *artNode) replaceChild(a *artAllocator, c byte, child artNode) { switch an.kind { case typeNode4: n4 := an.asNode4(a) - for idx := uint8(0); idx < n4.nodeNum; idx++ { - if n4.keys[idx] == c { - n4.children[idx] = child - return - } + idx := n4.findChild(c) + if idx != notExistIndex { + n4.children[idx] = child + return } - panic("swap child failed") case typeNode16: n16 := an.asNode16(a) - for idx := uint8(0); idx < n16.nodeNum; idx++ { - if n16.keys[idx] == c { - n16.children[idx] = child - return - } + idx := n16.findChild(c) + if idx != notExistIndex { + n16.children[idx] = child + return } - panic("swap child failed") case typeNode48: n48 := an.asNode48(a) - if n48.present[c>>n48s]&(1<<(c%n48m)) != 0 { - n48.children[n48.keys[c]] = child + idx := n48.findChild(c) + if idx != notExistIndex { + n48.children[idx] = child return } - panic("swap child failed") case typeNode256: - an.addChild256(a, c, false, child) + n256 := an.asNode256(a) + if n256.present[c>>n48s]&(1<<(c%n48m)) != 0 { + n256.children[c] = child + return + } } + panic("replace child failed") } // addChild adds a child to the node. // the added index `c` should not exist. -// Return if the node is grown to higher capacity. +// Return true if the node is grown to higher capacity. func (an *artNode) addChild(a *artAllocator, c byte, inplace bool, child artNode) bool { + if inplace { + an.asNode(a).inplaceLeaf = child + return false + } switch an.kind { case typeNode4: - return an.addChild4(a, c, inplace, child) + return an.addChild4(a, c, child) case typeNode16: - return an.addChild16(a, c, inplace, child) + return an.addChild16(a, c, child) case typeNode48: - return an.addChild48(a, c, inplace, child) + return an.addChild48(a, c, child) case typeNode256: - return an.addChild256(a, c, inplace, child) + return an.addChild256(a, c, child) } return false } -func (an *artNode) addChild4(a *artAllocator, c byte, inplace bool, child artNode) bool { +func (an *artNode) addChild4(a *artAllocator, c byte, child artNode) bool { node := an.asNode4(a) - if inplace { - node.inplaceLeaf = child - return false - } - if node.nodeNum >= node4cap { an.grow(a) - an.addChild(a, c, inplace, child) + an.addChild(a, c, false, child) return true } i := uint8(0) for ; i < node.nodeNum; i++ { - if c < node.keys[i] { + if c <= node.keys[i] { + if testing.Testing() && c == node.keys[i] { + panic("key already exists") + } break } } if i < node.nodeNum { - for j := node.nodeNum; j > i; j-- { - node.keys[j] = node.keys[j-1] - node.children[j] = node.children[j-1] - } + 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]) } node.keys[i] = c node.children[i] = child @@ -549,21 +566,16 @@ func (an *artNode) addChild4(a *artAllocator, c byte, inplace bool, child artNod return false } -func (an *artNode) addChild16(a *artAllocator, c byte, inplace bool, child artNode) bool { +func (an *artNode) addChild16(a *artAllocator, c byte, child artNode) bool { node := an.asNode16(a) - if inplace { - node.inplaceLeaf = child - return false - } - if node.nodeNum >= node16cap { an.grow(a) - an.addChild(a, c, inplace, child) + an.addChild(a, c, false, child) return true } - i, _ := sort.Find(int(node.nodeNum), func(i int) int { + i, found := sort.Find(int(node.nodeNum), func(i int) int { if node.keys[i] < c { return 1 } @@ -573,6 +585,10 @@ func (an *artNode) addChild16(a *artAllocator, c byte, inplace bool, child artNo return -1 }) + if testing.Testing() && found { + 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]) @@ -584,20 +600,19 @@ func (an *artNode) addChild16(a *artAllocator, c byte, inplace bool, child artNo return false } -func (an *artNode) addChild48(a *artAllocator, c byte, inplace bool, child artNode) bool { +func (an *artNode) addChild48(a *artAllocator, c byte, child artNode) bool { node := an.asNode48(a) - if inplace { - node.inplaceLeaf = child - return false - } - if node.nodeNum >= node48cap { an.grow(a) - an.addChild(a, c, inplace, child) + an.addChild(a, c, false, child) return true } + if testing.Testing() && node.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 @@ -605,12 +620,11 @@ func (an *artNode) addChild48(a *artAllocator, c byte, inplace bool, child artNo return false } -func (an *artNode) addChild256(a *artAllocator, c byte, inplace bool, child artNode) bool { +func (an *artNode) addChild256(a *artAllocator, c byte, child artNode) bool { node := an.asNode256(a) - if inplace { - node.inplaceLeaf = child - return false + if testing.Testing() && node.present[c>>n48s]&(1<<(c%n48m)) != 0 { + panic("key already exists") } node.present[c>>n48s] |= 1 << (c % n48m) diff --git a/internal/unionstore/art/art_node_test.go b/internal/unionstore/art/art_node_test.go index 88a47ce72..da4ad5da0 100644 --- a/internal/unionstore/art/art_node_test.go +++ b/internal/unionstore/art/art_node_test.go @@ -157,7 +157,7 @@ func TestNodePrefix(t *testing.T) { matchKey = append(make([]byte, 10), matchKey...) mismatchKey = append(make([]byte, 10), mismatchKey...) leafAddr, _ = allocator.allocLeaf(leafKey) - an.swapChild(&allocator, 2, artNode{kind: typeLeaf, addr: leafAddr}) + an.replaceChild(&allocator, 2, artNode{kind: typeLeaf, addr: leafAddr}) n.setPrefix(leafKey[10:], maxPrefixLen+1) require.Equal(t, uint32(maxPrefixLen+1), n.matchDeep(&allocator, an, matchKey, 10)) require.Equal(t, uint32(maxPrefixLen), n.matchDeep(&allocator, an, mismatchKey, 10)) @@ -257,3 +257,11 @@ func TestNextPrevPresentIdx(t *testing.T) { addr, n256 := allocator.allocNode48() testFn(&artNode{kind: typeNode48, addr: addr}, n256) } + +func TestLCP(t *testing.T) { + k1 := []byte{1, 2, 3, 4, 5, 6, 7, 8, 9} + k2 := []byte{1, 2, 3, 4, 5, 7, 7, 8} + for i := 0; i < 6; i++ { + require.Equal(t, uint32(5-i), longestCommonPrefix(k1, k2, uint32(i))) + } +} diff --git a/internal/unionstore/art/art_test.go b/internal/unionstore/art/art_test.go index 837f5df8e..f5dc0b7d6 100644 --- a/internal/unionstore/art/art_test.go +++ b/internal/unionstore/art/art_test.go @@ -4,39 +4,43 @@ import ( "fmt" "testing" - "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" ) +func init() { + testMode = true +} + func TestSimple(t *testing.T) { tree := New() for i := 0; i < 256; i++ { key := []byte{byte(i)} _, err := tree.Get(key) - assert.Equal(t, err, tikverr.ErrNotExist) + require.Equal(t, err, tikverr.ErrNotExist) err = tree.Set(key, key) - assert.Nil(t, err) + require.Nil(t, err) val, err := tree.Get(key) - assert.Nil(t, err, i) - assert.Equal(t, val, key, i) + require.Nil(t, err, i) + require.Equal(t, val, key, i) } } func TestSubNode(t *testing.T) { tree := New() - assert.Nil(t, tree.Set([]byte("a"), []byte("a"))) - assert.Nil(t, tree.Set([]byte("aa"), []byte("aa"))) - assert.Nil(t, tree.Set([]byte("aaa"), []byte("aaa"))) + require.Nil(t, tree.Set([]byte("a"), []byte("a"))) + require.Nil(t, tree.Set([]byte("aa"), []byte("aa"))) + require.Nil(t, tree.Set([]byte("aaa"), []byte("aaa"))) v, err := tree.Get([]byte("a")) - assert.Nil(t, err) - assert.Equal(t, v, []byte("a")) + require.Nil(t, err) + require.Equal(t, v, []byte("a")) v, err = tree.Get([]byte("aa")) - assert.Nil(t, err) - assert.Equal(t, v, []byte("aa")) + require.Nil(t, err) + require.Equal(t, v, []byte("aa")) v, err = tree.Get([]byte("aaa")) - assert.Nil(t, err) - assert.Equal(t, v, []byte("aaa")) + require.Nil(t, err) + require.Equal(t, v, []byte("aaa")) } func BenchmarkReadAfterWriteArt(b *testing.B) { @@ -50,7 +54,7 @@ func BenchmarkReadAfterWriteArt(b *testing.B) { for i := 0; i < b.N; i++ { tree.Set(buf[i], buf[i]) v, _ := tree.Get(buf[i]) - assert.Equal(b, v, buf[i]) + require.Equal(b, v, buf[i]) } } @@ -67,8 +71,8 @@ func TestBenchKey(t *testing.T) { } for k := 0; k < cnt; k++ { v, err := buffer.Get(encodeInt(k)) - assert.Nil(t, err, k) - assert.Equal(t, v, encodeInt(k)) + require.Nil(t, err, k) + require.Equal(t, v, encodeInt(k)) } } @@ -77,11 +81,11 @@ func TestLeafWithCommonPrefix(t *testing.T) { tree.Set([]byte{1, 1, 1}, []byte{1, 1, 1}) tree.Set([]byte{1, 1, 2}, []byte{1, 1, 2}) v, err := tree.Get([]byte{1, 1, 1}) - assert.Nil(t, err) - assert.Equal(t, v, []byte{1, 1, 1}) + require.Nil(t, err) + require.Equal(t, v, []byte{1, 1, 1}) v, err = tree.Get([]byte{1, 1, 2}) - assert.Nil(t, err) - assert.Equal(t, v, []byte{1, 1, 2}) + require.Nil(t, err) + require.Equal(t, v, []byte{1, 1, 2}) } func TestUpdateInplace(t *testing.T) { @@ -90,7 +94,7 @@ func TestUpdateInplace(t *testing.T) { for i := 0; i < 256; i++ { val := make([]byte, 4096) tree.Set(key, val) - assert.Equal(t, tree.allocator.vlogAllocator.Blocks(), 1) + require.Equal(t, tree.allocator.vlogAllocator.Blocks(), 1) } } @@ -98,46 +102,46 @@ func TestFlag(t *testing.T) { tree := New() tree.Set([]byte{0}, []byte{0}, kv.SetPresumeKeyNotExists) flags, err := tree.GetFlags([]byte{0}) - assert.Nil(t, err) - assert.True(t, flags.HasPresumeKeyNotExists()) + require.Nil(t, err) + require.True(t, flags.HasPresumeKeyNotExists()) tree.Set([]byte{1}, []byte{1}, kv.SetKeyLocked) flags, err = tree.GetFlags([]byte{1}) - assert.Nil(t, err) - assert.True(t, flags.HasLocked()) + require.Nil(t, err) + require.True(t, flags.HasLocked()) // iterate can also see the flags //it, err := tree.Iter(nil, nil) - //assert.Nil(t, err) - //assert.True(t, it.Valid()) - //assert.Equal(t, it.Key(), []byte{0}) - //assert.Equal(t, it.Value(), []byte{0}) - //assert.True(t, it.Flags().HasPresumeKeyNotExists()) - //assert.False(t, it.Flags().HasLocked()) - //assert.Nil(t, it.Next()) - //assert.True(t, it.Valid()) - //assert.Equal(t, it.Key(), []byte{1}) - //assert.Equal(t, it.Value(), []byte{1}) - //assert.True(t, it.Flags().HasLocked()) - //assert.False(t, it.Flags().HasPresumeKeyNotExists()) - //assert.Nil(t, it.Next()) - //assert.False(t, it.Valid()) + //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) { key1 := []byte{109, 68, 66, 115, 0, 0, 0, 0, 0, 250, 0, 0, 0, 0, 0, 0, 0, 104, 68, 66, 58, 49, 0, 0, 0, 0, 251} key2 := []byte{109, 68, 66, 115, 0, 0, 0, 0, 0, 250, 0, 0, 0, 0, 0, 0, 0, 105, 68, 66, 58, 49, 0, 0, 0, 0, 251} buffer := New() - assert.Nil(t, buffer.Set(key1, []byte{1})) - assert.Nil(t, buffer.Set(key2, []byte{2})) + require.Nil(t, buffer.Set(key1, []byte{1})) + require.Nil(t, buffer.Set(key2, []byte{2})) val, err := buffer.Get(key1) - assert.Nil(t, err) - assert.Equal(t, val, []byte{1}) + require.Nil(t, err) + require.Equal(t, val, []byte{1}) val, err = buffer.Get(key2) - assert.Nil(t, err) - assert.Equal(t, val, []byte{2}) - assert.Nil(t, buffer.Set(key2, []byte{3})) + require.Nil(t, err) + require.Equal(t, val, []byte{2}) + require.Nil(t, buffer.Set(key2, []byte{3})) val, err = buffer.Get(key2) - assert.Nil(t, err) - assert.Equal(t, val, []byte{3}) + require.Nil(t, err) + require.Equal(t, val, []byte{3}) } func TestLongPrefix2(t *testing.T) { @@ -147,37 +151,37 @@ func TestLongPrefix2(t *testing.T) { key3 := []byte{0, 97, 0, 0, 0, 0, 0, 0, 0, 248, 0, 0, 0, 0, 0, 0, 0, 108, 127, 255, 255, 255, 255, 255, 255, 253} key4 := []byte{0, 97, 0, 0, 0, 0, 0, 0, 0, 248, 0, 0, 0, 0, 0, 0, 0, 76} key5 := []byte{0, 97, 0, 0, 0, 0, 0, 0, 0, 248, 0, 0, 0, 0, 0, 0, 0, 108, 127, 255, 255, 255, 255, 255, 255, 252} - assert.Nil(t, tree.Set(key1, key1)) - assert.Nil(t, tree.Set(key2, key2)) - assert.Nil(t, tree.Set(key3, key3)) - assert.Nil(t, tree.Set(key4, key4)) - assert.Nil(t, tree.Set(key5, key5)) - assert.Nil(t, tree.Set(key4, key4)) + require.Nil(t, tree.Set(key1, key1)) + require.Nil(t, tree.Set(key2, key2)) + require.Nil(t, tree.Set(key3, key3)) + require.Nil(t, tree.Set(key4, key4)) + require.Nil(t, tree.Set(key5, key5)) + require.Nil(t, tree.Set(key4, key4)) val, err := tree.Get(key1) - assert.Nil(t, err) - assert.Equal(t, val, key1) + require.Nil(t, err) + require.Equal(t, val, key1) val, err = tree.Get(key2) - assert.Nil(t, err) - assert.Equal(t, val, key2) + require.Nil(t, err) + require.Equal(t, val, key2) val, err = tree.Get(key3) - assert.Nil(t, err) - assert.Equal(t, val, key3) + require.Nil(t, err) + require.Equal(t, val, key3) val, err = tree.Get(key4) - assert.Nil(t, err) - assert.Equal(t, val, key4) + require.Nil(t, err) + require.Equal(t, val, key4) val, err = tree.Get(key5) - assert.Nil(t, err) - assert.Equal(t, val, key5) + require.Nil(t, err) + require.Equal(t, val, key5) } func TestFlagOnlyKey(t *testing.T) { tree := New() tree.Set([]byte{0}, nil, kv.SetAssertNone) flags, err := tree.GetFlags([]byte{0}) - assert.Nil(t, err) - assert.False(t, flags.HasAssertionFlags()) + require.Nil(t, err) + require.False(t, flags.HasAssertionFlags()) _, err = tree.Get([]byte{0}) - assert.Error(t, err) + require.Error(t, err) } func TestSearchOptimisticMismatch(t *testing.T) { @@ -187,5 +191,42 @@ func TestSearchOptimisticMismatch(t *testing.T) { tree.Set(append(prefix, []byte{2}...), prefix) // the search key is matched within maxPrefixLen, but the full key is not matched. _, err := tree.Get(append(make([]byte, 21), []byte{1, 1}...)) - assert.NotNil(t, err) + require.NotNil(t, err) +} + +func TestExpansion(t *testing.T) { + // expand leaf + tree := New() + prefix := make([]byte, maxPrefixLen) + tree.Set(append(prefix, []byte{1, 1, 1, 1}...), []byte{1}) + an := tree.root.asNode4(&tree.allocator).children[0] + require.Equal(t, an.kind, typeLeaf) + tree.Set(append(prefix, []byte{1, 1, 1, 2}...), []byte{2}) + an = tree.root.asNode4(&tree.allocator).children[0] + require.Equal(t, an.kind, typeNode4) + n4 := an.asNode4(&tree.allocator) + require.Equal(t, n4.nodeNum, uint8(2)) + require.Equal(t, n4.prefixLen, uint32(22)) + require.Equal(t, n4.keys[:2], []byte{1, 2}) + require.Equal(t, n4.children[0].asLeaf(&tree.allocator).GetKey(), append(prefix, []byte{1, 1, 1, 1}...)) + require.Equal(t, n4.children[1].asLeaf(&tree.allocator).GetKey(), append(prefix, []byte{1, 1, 1, 2}...)) + oldAddr := an.addr + + // expand node + tree.Set(append(prefix, []byte{1, 255, 2}...), []byte{1, 2}) + an = tree.root.asNode4(&tree.allocator).children[0] + require.Equal(t, an.kind, typeNode4) + require.NotEqual(t, an.addr, oldAddr) // the node is expanded, parent node is the newly created. + n4 = an.asNode4(&tree.allocator) + require.Equal(t, n4.nodeNum, uint8(2)) + require.Equal(t, n4.prefixLen, uint32(20)) + // the old node4 is a child of new node4 now. + oldAn := n4.children[0] + require.Equal(t, oldAn.kind, typeNode4) + require.Equal(t, oldAn.addr, oldAddr) + oldN4 := oldAn.asNode4(&tree.allocator) + require.Equal(t, oldN4.prefixLen, uint32(1)) + require.Equal(t, oldN4.prefix[:1], []byte{1}) + 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}...)) }