// Copyright 2020 PingCAP, Inc. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // See the License for the specific language governing permissions and // limitations under the License. package memdb import ( "bytes" "math" "unsafe" "github.com/pingcap/tidb/util/fastrand" ) // KeyFlags are metadata associated with key type KeyFlags uint8 const ( flagNeedLocked KeyFlags = 1 << iota maxHeight = 16 nodeHeaderSize = int(unsafe.Sizeof(nodeHeader{})) initBlockSize = 4 * 1024 ) // HasNeedLocked returns whether the associated key required lock. func (m KeyFlags) HasNeedLocked() bool { return m&flagNeedLocked != 0 } // Sandbox is a space to keep pending kvs. type Sandbox struct { frozen bool done bool head headNode parent *Sandbox arena *arena height int length int size int arenaSnap arenaSnapshot } // NewSandbox create a new Sandbox. func NewSandbox() *Sandbox { arena := newArenaLocator() return &Sandbox{ height: 1, arena: arena, arenaSnap: arena.snapshot(), } } // Get returns value for key in this sandbox's space. func (sb *Sandbox) Get(key []byte) []byte { node, data, match := sb.findGreaterEqual(key) if !match { return nil } return node.getValue(data) } // GetFlags returns flag for key in this sandbox's space. func (sb *Sandbox) GetFlags(key []byte) KeyFlags { node, _, match := sb.findGreaterEqual(key) if !match { return 0 } return node.flags } // PutWithFlags inserts kv into this sandbox with flags. func (sb *Sandbox) PutWithFlags(key []byte, flags KeyFlags, value []byte) { if sb.frozen { panic("cannot write to a sandbox when it has forked a new sandbox") } head := sb.getHead() arena := sb.arena lsHeight := sb.height var prev [maxHeight + 1]nodeWithAddr var next [maxHeight + 1]nodeWithAddr prev[lsHeight] = head var exists bool for i := lsHeight - 1; i >= 0; i-- { // Use higher level to speed up for current level. prev[i], next[i], exists = sb.findSpliceForLevel(key, prev[i+1], i) } var height int if !exists { height = sb.randomHeight() } else { height = sb.prepareOverwrite(next[:]) } x, addr := arena.newNode(key, value, height) if height > lsHeight { sb.height = height } x.flags |= flags // We always insert from the base level and up. After you add a node in base level, we cannot // create a node in the level above because it would have discovered the node in the base level. for i := 0; i < height; i++ { x.setNexts(i, next[i].addr) if prev[i].node == nil { prev[i] = head } prev[i].setNexts(i, addr) } x.prev = prev[0].addr if next[0].node != nil { next[0].prev = addr } sb.length++ sb.size += len(key) + len(value) } // Put inserts kv into this sandbox. func (sb *Sandbox) Put(key, value []byte) { sb.PutWithFlags(key, 0, value) } // PutWithNeedLock inserts kv into this sandbox and mark it with flagNeedLocked. func (sb *Sandbox) PutWithNeedLock(key, value []byte) { sb.PutWithFlags(key, flagNeedLocked, value) } // Derive derive a new sandbox to buffer a batch of modifactions. func (sb *Sandbox) Derive() *Sandbox { if sb.frozen { panic("cannot start second sandbox") } sb.frozen = true new := &Sandbox{ parent: sb, height: 1, arena: sb.arena, arenaSnap: sb.arena.snapshot(), } return new } // Flush flushes all kvs into parent sandbox. func (sb *Sandbox) Flush() int { if sb.parent == nil || sb.done { return 0 } if !sb.parent.frozen { panic("the parent sandbox must be freezed when doing flush") } sb.parent.frozen = false sb.done = true return sb.parent.merge(sb) } // GetParent returns the parent sandbox. func (sb *Sandbox) GetParent() *Sandbox { return sb.parent } // Discard discards all kvs in this sandbox. // It is safe to discard a flushed sandbox, and it is recommend to // call discard using defer to maintain correct state of parent. func (sb *Sandbox) Discard() { if sb.done { return } if sb.parent != nil { if !sb.parent.frozen { panic("the parent sandbox must be freezed when doing discard") } sb.parent.frozen = false sb.done = true } else if sb.frozen { panic("root sandbox is freezed") } sb.head = headNode{} sb.height = 1 sb.length = 0 sb.size = 0 sb.arena.revert(sb.arenaSnap) if sb.parent != nil { // nil out arena to pervent data corruption by accident. sb.arena = nil } } // Len returns the number of entries in the DB. func (sb *Sandbox) Len() int { return sb.length } // Size returns sum of keys and values length. Note that deleted // key/value will not be accounted for, but it will still consume // the buffer, since the buffer is append only. func (sb *Sandbox) Size() int { return sb.size } // The pointers in findSpliceForLevel may point to the node which going to be overwrite, // prepareOverwrite update them to point to the next node, so we can link new node with the list correctly. func (sb *Sandbox) prepareOverwrite(next []nodeWithAddr) int { old := next[0] // Update necessary states. sb.size -= int(old.valLen) + int(old.keyLen) sb.length-- height := int(old.height) for i := 0; i < height; i++ { if next[i].addr == old.addr { next[i].addr = old.nexts(i) if !next[i].addr.isNull() { data := sb.arena.getFrom(next[i].addr) next[i].node = (*node)(unsafe.Pointer(&data[0])) } } } return height } func (sb *Sandbox) getHead() nodeWithAddr { head := (*node)(unsafe.Pointer(&sb.head)) return nodeWithAddr{node: head} } func (sb *Sandbox) randomHeight() int { h := 1 for h < maxHeight && fastrand.Uint32() < uint32(math.MaxUint32)/4 { h++ } return h } // findSpliceForLevel returns (outBefore, outAfter) with outBefore.key < key <= outAfter.key. // The input "before" tells us where to start looking. // If we found a node with the same key, then we return true. func (sb *Sandbox) findSpliceForLevel(key []byte, before nodeWithAddr, level int) (nodeWithAddr, nodeWithAddr, bool) { arena := sb.arena for { // Assume before.key < key. nextAddr := before.nexts(level) if nextAddr.isNull() { return before, nodeWithAddr{}, false } data := arena.getFrom(nextAddr) next := nodeWithAddr{(*node)(unsafe.Pointer(&data[0])), nextAddr} nextKey := next.getKey(data) cmp := bytes.Compare(nextKey, key) if cmp >= 0 { // before.key < key < next.key. We are done for this level. return before, next, cmp == 0 } before = next // Keep moving right on this level. } } func (sb *Sandbox) findGreaterEqual(key []byte) (*node, []byte, bool) { head := sb.getHead() prev := head.node level := sb.height - 1 arena := sb.arena for { var nextData []byte var next *node addr := prev.nexts(level) if !addr.isNull() { nextData = arena.getFrom(addr) next = (*node)(unsafe.Pointer(&nextData[0])) nextKey := next.getKey(nextData) cmp := bytes.Compare(nextKey, key) if cmp < 0 { // next key is still smaller, keep moving. prev = next continue } if cmp == 0 { // prev.key < key == next.key. return next, nextData, true } } // next is greater than key or next is nil. go to the lower level. if level > 0 { level-- continue } return next, nextData, false } } func (sb *Sandbox) findLess(key []byte, allowEqual bool) (*node, []byte, bool) { var prevData []byte head := sb.getHead() prev := head.node level := sb.height - 1 arena := sb.arena for { next, nextData := prev.getNext(arena, level) if next != nil { cmp := bytes.Compare(key, next.getKey(nextData)) if cmp > 0 { // prev.key < next.key < key. We can continue to move right. prev = next prevData = nextData continue } if cmp == 0 && allowEqual { // prev.key < key == next.key. return next, nextData, true } } // get closer to the key in the lower level. if level > 0 { level-- continue } break } // We are not going to return head. if prev == head.node { return nil, nil, false } return prev, prevData, false } // findLast returns the last element. If head (empty db), we return nil. All the find functions // will NEVER return the head nodes. func (sb *Sandbox) findLast() (*node, []byte) { var nodeData []byte head := sb.getHead() node := head.node level := sb.height - 1 arena := sb.arena for { next, nextData := node.getNext(arena, level) if next != nil { node = next nodeData = nextData continue } if level == 0 { if node == head.node { return nil, nil } return node, nodeData } level-- } } func (sb *Sandbox) merge(new *Sandbox) int { var ms mergeState arena := sb.arena if sb.head.nexts[0].isNull() { // current skip-list is empty, overwite head node using the new list's head. sb.head = new.head sb.height = new.height sb.length = new.length sb.size = new.size return new.length } var ( newNode *node nextNode *node newNodeAddr arenaAddr nextNodeAddr arenaAddr newNodeData []byte nextNodeData []byte ) head := new.getHead() newNodeAddr = head.nexts(0) newNode, newNodeData = head.getNext(arena, 0) for newNode != nil { key := newNode.getKey(newNodeData) recomputeHeight := ms.calculateRecomputeHeight(key, sb) nextNodeAddr = newNode.nexts(0) nextNode, nextNodeData = newNode.getNext(arena, 0) var exists bool if recomputeHeight > 0 { for i := recomputeHeight - 1; i >= 0; i-- { ms.prev[i], ms.next[i], exists = sb.findSpliceForLevel(key, ms.prev[i+1], i) } } height := int(newNode.height) if exists { newNode.flags = newNode.flags | ms.next[0].flags height = sb.prepareOverwrite(ms.next[:]) if height > int(newNode.height) { // The space is not enough, we have to create a new node. k := newNode.getKey(newNodeData) v := newNode.getValue(newNodeData) flags := newNode.flags newNode, newNodeAddr = arena.newNode(k, v, height) newNode.flags = flags } } if height > sb.height { sb.height = height } for i := 0; i < height; i++ { newNode.setNexts(i, ms.next[i].addr) if ms.prev[i].node == nil { ms.prev[i] = head } ms.prev[i].setNexts(i, newNodeAddr) } newNode.prev = ms.prev[0].addr if ms.next[0].node != nil { ms.next[0].prev = newNodeAddr } newNode, newNodeAddr, newNodeData = nextNode, nextNodeAddr, nextNodeData } sb.length += new.length sb.size += new.size return new.length } type mergeState struct { height int // hitHeight is used to reduce cost of calculateRecomputeHeight. // For random workload, comparing hint keys from bottom up is wasted work. // So we record the hit height of the last operation, only grow recompute height from near that height. hitHeight int prev [maxHeight + 1]nodeWithAddr next [maxHeight + 1]nodeWithAddr } func (ms *mergeState) calculateRecomputeHeight(key []byte, sb *Sandbox) int { head := sb.getHead() listHeight := sb.height arena := sb.arena if ms.height < listHeight { // Either splice is never used or list height has grown, we recompute all. ms.prev[listHeight] = head ms.next[listHeight] = nodeWithAddr{} ms.height = listHeight ms.hitHeight = ms.height return listHeight } recomputeHeight := ms.hitHeight - 2 if recomputeHeight < 0 { recomputeHeight = 0 } for recomputeHeight < listHeight { prev := ms.prev[recomputeHeight] next := ms.next[recomputeHeight] prevNext := prev.nexts(recomputeHeight) if prevNext != next.addr { recomputeHeight++ continue } if prev.addr != head.addr && !prev.addr.isNull() && bytes.Compare(key, prev.getKey(arena.getFrom(prev.addr))) <= 0 { // Key is before splice. for prev.addr == ms.prev[recomputeHeight].addr { recomputeHeight++ } continue } if !next.addr.isNull() && bytes.Compare(key, next.getKey(arena.getFrom(next.addr))) > 0 { // Key is after splice. for next == ms.next[recomputeHeight] { recomputeHeight++ } continue } break } ms.hitHeight = recomputeHeight return recomputeHeight } type nodeHeader struct { flags KeyFlags height uint8 keyLen uint16 valLen uint32 } type node struct { nodeHeader // Addr of previous node at base level. prev arenaAddr // node is a variable length struct. // The nextsBase is the first element of nexts slice, // it act as the base pointer we do pointer arithmetic in `next` and `setNext`. nextsBase arenaAddr } type headNode struct { nodeHeader // Addr of previous node at base level. prev arenaAddr nexts [maxHeight]arenaAddr } type nodeWithAddr struct { *node addr arenaAddr } func (n *node) nodeLen() int { return int(n.height)*8 + 8 + nodeHeaderSize } func (n *node) getKey(buf []byte) []byte { nodeLen := n.nodeLen() return buf[nodeLen : nodeLen+int(n.keyLen)] } func (n *node) getValue(buf []byte) []byte { nodeLenKeyLen := n.nodeLen() + int(n.keyLen) return buf[nodeLenKeyLen : nodeLenKeyLen+int(n.valLen)] } func (n *node) nexts(level int) arenaAddr { return *n.nextsAddr(level) } func (n *node) setNexts(level int, val arenaAddr) { *n.nextsAddr(level) = val } func (n *node) nextsAddr(idx int) *arenaAddr { offset := uintptr(idx) * unsafe.Sizeof(n.nextsBase) return (*arenaAddr)(unsafe.Pointer(uintptr(unsafe.Pointer(&n.nextsBase)) + offset)) } func (n *node) getNext(arena *arena, level int) (*node, []byte) { addr := n.nexts(level) if addr.isNull() { return nil, nil } data := arena.getFrom(addr) node := (*node)(unsafe.Pointer(&data[0])) return node, data }