kv: replace memdb with a more memory efficient version (#11807)

kv/memdb/arena.go

@@ -0,0 +1,131 @@
+// 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 struct {
+	blockIdx    uint32
+	blockOffset uint32
+}
+
+func (addr arenaAddr) isNull() bool {
+	return addr.blockIdx == 0 && addr.blockOffset == 0
+}
+
+func newArenaAddr(idx int, offset uint32) arenaAddr {
+	return arenaAddr{
+		blockIdx:    uint32(idx) + 1,
+		blockOffset: offset,
+	}
+}
+
+const (
+	nullBlockOffset = math.MaxUint32
+	maxBlockSize    = 128 << 20
+)
+
+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-1].getFrom(addr.blockOffset)
+}
+
+func (a *arena) alloc(size int) (arenaAddr, []byte) {
+	if size >= maxBlockSize {
+		// Use a separate block to store entry which size larger than specified block size.
+		blk := newArenaBlock(size)
+		blk.length = size
+		a.blocks = append(a.blocks, blk)
+
+		addr := newArenaAddr(len(a.blocks)-1, 0)
+		return addr, blk.buf
+	}
+
+	addr, data := a.allocInBlock(a.availIdx, size)
+	if !addr.isNull() {
+		return addr, data
+	}
+
+	a.enlarge(size)
+	return a.allocInBlock(a.availIdx, size)
+}
+
+func (a *arena) enlarge(size int) {
+	a.blockSize <<= 1
+	for a.blockSize <= size {
+		a.blockSize <<= 1
+	}
+	// Size always less than maxBlockSize.
+	if a.blockSize > maxBlockSize {
+		a.blockSize = maxBlockSize
+	}
+	a.blocks = append(a.blocks, newArenaBlock(a.blockSize))
+	a.availIdx = int(uint32(len(a.blocks) - 1))
+}
+
+func (a *arena) allocInBlock(idx, size int) (arenaAddr, []byte) {
+	offset, data := a.blocks[idx].alloc(size)
+	if offset == nullBlockOffset {
+		return arenaAddr{}, nil
+	}
+	return newArenaAddr(idx, offset), data
+}
+
+func (a *arena) reset() {
+	a.availIdx = 0
+	a.blockSize = len(a.blocks[0].buf)
+	a.blocks = []arenaBlock{a.blocks[0]}
+	a.blocks[0].reset()
+}
+
+type arenaBlock struct {
+	buf    []byte
+	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, []byte) {
+	offset := a.length
+	newLen := offset + size
+	if newLen > len(a.buf) {
+		return nullBlockOffset, nil
+	}
+	a.length = newLen
+	return uint32(offset), a.buf[offset : offset+size]
+}
+
+func (a *arenaBlock) reset() {
+	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 if 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.nexts[0])
+}
+
+// Prev moves the iterator to the previous entry.
+func (it *Iterator) Prev() {
+	it.changeToAddr(it.curr.prev)
+}
+
+// Seek locates the iterator to the first entry with a key >= seekKey.
+func (it *Iterator) Seek(seekKey []byte) {
+	node, nodeData, _ := it.db.findGreaterEqual(seekKey) // 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.isNull() {
+		data = it.db.arena.getFrom(addr)
+		n = (*node)(unsafe.Pointer(&data[0]))
+	}
+	it.updateState(n, data)
+}