kv: replace memdb with a more memory efficient version

kv/memdb/arena.go

@@ -0,0 +1,119 @@
+// Copyright 2019 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 "math"
+
+type arenaAddr uint64
+
+const (
+	alignMask                 = 1<<32 - 8 // 29 bit 1 and 3 bit 0.
+	nullBlockOffset           = math.MaxUint32
+	maxBlockSize              = 128 << 20
+	nullArenaAddr   arenaAddr = 0
+)
+
+func (addr arenaAddr) blockIdx() int {
+	return int(addr>>32 - 1)
+}
+
+func (addr arenaAddr) blockOffset() uint32 {
+	return uint32(addr)
+}
+
+func newArenaAddr(blockIdx int, blockOffset uint32) arenaAddr {
+	return arenaAddr(uint64(blockIdx+1)<<32 | uint64(blockOffset))
+}
+
+type arena struct {
+	blockSize int
+	availIdx  int
+	blocks    []*arenaBlock
+}
+
+func newArenaLocator(initBlockSize int) *arena {
+	return &arena{
+		blockSize: initBlockSize,
+		blocks:    []*arenaBlock{newArenaBlock(initBlockSize)},
+	}
+}
+
+func (a *arena) getFrom(addr arenaAddr) []byte {
+	return a.blocks[addr.blockIdx()].getFrom(addr.blockOffset())
+}
+
+func (a *arena) alloc(size int) (arenaAddr, []byte) {
+	if size > a.blockSize {
+		// Use a separate block to store entry which size larger than specified block size.
+		blk := newArenaBlock(size)
+		addr := newArenaAddr(len(a.blocks), 0)
+		a.blocks = append(a.blocks, blk)
+		return addr, blk.buf
+	}
+
+	for {
+		block := a.blocks[a.availIdx]
+		blockOffset := block.alloc(size)
+		if blockOffset != nullBlockOffset {
+			return newArenaAddr(a.availIdx, blockOffset), block.buf[blockOffset : int(blockOffset)+size]
+		}
+
+		blockSize := a.blockSize << 1
+		if blockSize <= maxBlockSize {
+			a.blockSize = blockSize
+		}
+		a.blocks = append(a.blocks, newArenaBlock(a.blockSize))
+		a.availIdx = len(a.blocks) - 1
+	}
+}
+
+func (a *arena) reset() {
+	a.availIdx = 0
+	a.blockSize = len(a.blocks[0].buf)
+	a.blocks = a.blocks[:1]
+	a.blocks[0].reset()
+}
+
+type arenaBlock struct {
+	buf    []byte
+	ref    uint64
+	length int
+}
+
+func newArenaBlock(blockSize int) *arenaBlock {
+	return &arenaBlock{
+		buf: make([]byte, blockSize),
+	}
+}
+
+func (a *arenaBlock) getFrom(offset uint32) []byte {
+	return a.buf[offset:]
+}
+
+func (a *arenaBlock) alloc(size int) uint32 {
+	// The returned addr should be aligned in 8 bytes.
+	offset := (a.length + 7) & alignMask
+	a.length = offset + size
+	if a.length > len(a.buf) {
+		return nullBlockOffset
+	}
+	a.ref++
+	return uint32(offset)
+}
+
+func (a *arenaBlock) reset() {
+	a.buf = a.buf[:0]
+	a.ref = 0
+	a.length = 0
+}

kv/memdb/iterator.go

@@ -0,0 +1,102 @@
+// Copyright 2019 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 "unsafe"
+
+// Iterator iterates the entries in the DB.
+type Iterator struct {
+	db   *DB
+	curr *node
+	key  []byte
+	val  []byte
+}
+
+// NewIterator returns a new Iterator for the lock store.
+func (db *DB) NewIterator() Iterator {
+	return Iterator{
+		db: db,
+	}
+}
+
+// Valid returns true iff the iterator is positioned at a valid node.
+func (it *Iterator) Valid() bool { return it.curr != nil }
+
+// Key returns the key at the current position.
+func (it *Iterator) Key() []byte {
+	return it.key
+}
+
+// Value returns value.
+func (it *Iterator) Value() []byte {
+	return it.val
+}
+
+// Next moves the iterator to the next entry.
+func (it *Iterator) Next() {
+	it.changeToAddr(it.curr.getNextAddr(0))
+}
+
+// Prev moves the iterator to the previous entry.
+func (it *Iterator) Prev() {
+	it.changeToAddr(it.curr.getPrevAddr())
+}
+
+// Seek locates the iterator to the first entry with a key >= seekKey.
+func (it *Iterator) Seek(seekKey []byte) {
+	node, nodeData, _ := it.db.findGreater(seekKey, true) // find >=.
+	it.updateState(node, nodeData)
+}
+
+// SeekForPrev locates the iterator to the last entry with key <= target.
+func (it *Iterator) SeekForPrev(target []byte) {
+	node, nodeData, _ := it.db.findLess(target, true) // find <=.
+	it.updateState(node, nodeData)
+}
+
+// SeekForExclusivePrev locates the iterator to the last entry with key < target.
+func (it *Iterator) SeekForExclusivePrev(target []byte) {
+	node, nodeData, _ := it.db.findLess(target, false)
+	it.updateState(node, nodeData)
+}
+
+// SeekToFirst locates the iterator to the first entry.
+func (it *Iterator) SeekToFirst() {
+	node, nodeData := it.db.getNext(it.db.head.node, 0)
+	it.updateState(node, nodeData)
+}
+
+// SeekToLast locates the iterator to the last entry.
+func (it *Iterator) SeekToLast() {
+	node, nodeData := it.db.findLast()
+	it.updateState(node, nodeData)
+}
+
+func (it *Iterator) updateState(node *node, nodeData []byte) {
+	it.curr = node
+	if node != nil {
+		it.key = node.getKey(nodeData)
+		it.val = node.getValue(nodeData)
+	}
+}
+
+func (it *Iterator) changeToAddr(addr arenaAddr) {
+	var data []byte
+	var n *node
+	if addr != nullArenaAddr {
+		data = it.db.getArena().getFrom(addr)
+		n = (*node)(unsafe.Pointer(&data[0]))
+	}
+	it.updateState(n, data)
+}