You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

578 lines
13 KiB

// 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
}