membuffer: refactor the memdb to support multi implementations (#1426)

ref pingcap/tidb#55287

Signed-off-by: you06 <you1474600@gmail.com>

internal/unionstore/arena/arena_test.go

@@ -0,0 +1,79 @@
+// Copyright 2021 TiKV Authors
+//
+// 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,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// NOTE: The code in this file is based on code from the
+// TiDB project, licensed under the Apache License v 2.0
+//
+// https://github.com/pingcap/tidb/tree/cc5e161ac06827589c4966674597c137cc9e809c/store/tikv/unionstore/memdb_arena.go
+//
+
+// 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,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package arena
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+type dummyMemDB struct{}
+
+func (m *dummyMemDB) RevertNode(hdr *MemdbVlogHdr) {}
+func (m *dummyMemDB) InspectNode(addr MemdbArenaAddr) (KeyFlagsGetter, MemdbArenaAddr) {
+	return nil, NullAddr
+}
+
+func TestBigValue(t *testing.T) {
+	assert := assert.New(t)
+
+	var vlog MemdbVlog[KeyFlagsGetter, *dummyMemDB]
+	vlog.AppendValue(MemdbArenaAddr{0, 0}, NullAddr, make([]byte, 80<<20))
+	assert.Equal(vlog.blockSize, maxBlockSize)
+	assert.Equal(len(vlog.blocks), 1)
+
+	cp := vlog.Checkpoint()
+	vlog.AppendValue(MemdbArenaAddr{0, 1}, NullAddr, make([]byte, 127<<20))
+	vlog.RevertToCheckpoint(&dummyMemDB{}, &cp)
+
+	assert.Equal(vlog.blockSize, maxBlockSize)
+	assert.Equal(len(vlog.blocks), 2)
+	assert.PanicsWithValue("alloc size is larger than max block size", func() {
+		vlog.AppendValue(MemdbArenaAddr{0, 2}, NullAddr, make([]byte, maxBlockSize+1))
+	})
+}
+
+func TestValueLargeThanBlock(t *testing.T) {
+	assert := assert.New(t)
+	var vlog MemdbVlog[KeyFlagsGetter, *dummyMemDB]
+	vlog.AppendValue(MemdbArenaAddr{0, 0}, NullAddr, make([]byte, 1))
+	vlog.AppendValue(MemdbArenaAddr{0, 1}, NullAddr, make([]byte, 4096))
+	assert.Equal(len(vlog.blocks), 2)
+	vAddr := vlog.AppendValue(MemdbArenaAddr{0, 2}, NullAddr, make([]byte, 3000))
+	assert.Equal(len(vlog.blocks), 2)
+	val := vlog.GetValue(vAddr)
+	assert.Equal(len(val), 3000)
+}

internal/unionstore/art/art.go

@@ -0,0 +1,176 @@
+// Copyright 2024 TiKV Authors
+//
+// 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,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//nolint:unused
+package art
+
+import (
+	"math"
+
+	"github.com/tikv/client-go/v2/internal/unionstore/arena"
+	"github.com/tikv/client-go/v2/kv"
+)
+
+type ART struct {
+	allocator       artAllocator
+	root            artNode
+	stages          []arena.MemDBCheckpoint
+	vlogInvalid     bool
+	dirty           bool
+	entrySizeLimit  uint64
+	bufferSizeLimit uint64
+	len             int
+	size            int
+}
+
+func New() *ART {
+	var t ART
+	t.root = nullArtNode
+	t.stages = make([]arena.MemDBCheckpoint, 0, 2)
+	t.entrySizeLimit = math.MaxUint64
+	t.bufferSizeLimit = math.MaxUint64
+	t.allocator.nodeAllocator.freeNode4 = make([]arena.MemdbArenaAddr, 0, 1<<4)
+	t.allocator.nodeAllocator.freeNode16 = make([]arena.MemdbArenaAddr, 0, 1<<3)
+	t.allocator.nodeAllocator.freeNode48 = make([]arena.MemdbArenaAddr, 0, 1<<2)
+	return &t
+}
+
+func (t *ART) Get(key []byte) ([]byte, error) {
+	panic("unimplemented")
+}
+
+// GetFlags returns the latest flags associated with key.
+func (t *ART) GetFlags(key []byte) (kv.KeyFlags, error) {
+	panic("unimplemented")
+}
+
+func (t *ART) Set(key artKey, value []byte, ops []kv.FlagsOp) error {
+	panic("unimplemented")
+}
+
+func (t *ART) search(key artKey) (arena.MemdbArenaAddr, *artLeaf) {
+	panic("unimplemented")
+}
+
+func (t *ART) Dirty() bool {
+	panic("unimplemented")
+}
+
+// Mem returns the memory usage of MemBuffer.
+func (t *ART) Mem() uint64 {
+	panic("unimplemented")
+}
+
+// Len returns the count of entries in the MemBuffer.
+func (t *ART) Len() int {
+	panic("unimplemented")
+}
+
+// Size returns the size of the MemBuffer.
+func (t *ART) Size() int {
+	panic("unimplemented")
+}
+
+func (t *ART) checkpoint() arena.MemDBCheckpoint {
+	panic("unimplemented")
+}
+
+func (t *ART) RevertNode(hdr *arena.MemdbVlogHdr) {
+	panic("unimplemented")
+}
+
+func (t *ART) InspectNode(addr arena.MemdbArenaAddr) (*artLeaf, arena.MemdbArenaAddr) {
+	panic("unimplemented")
+}
+
+// Checkpoint returns a checkpoint of ART.
+func (t *ART) Checkpoint() *arena.MemDBCheckpoint {
+	panic("unimplemented")
+}
+
+// RevertToCheckpoint reverts the ART to the checkpoint.
+func (t *ART) RevertToCheckpoint(cp *arena.MemDBCheckpoint) {
+	panic("unimplemented")
+}
+
+func (t *ART) Stages() []arena.MemDBCheckpoint {
+	panic("unimplemented")
+}
+
+func (t *ART) Staging() int {
+	panic("unimplemented")
+}
+
+func (t *ART) Release(h int) {
+	panic("unimplemented")
+}
+
+func (t *ART) Cleanup(h int) {
+	panic("unimplemented")
+}
+
+func (t *ART) revertToCheckpoint(cp *arena.MemDBCheckpoint) {
+	panic("unimplemented")
+}
+
+func (t *ART) moveBackCursor(cursor *arena.MemDBCheckpoint, hdr *arena.MemdbVlogHdr) {
+	panic("unimplemented")
+}
+
+func (t *ART) truncate(snap *arena.MemDBCheckpoint) {
+	panic("unimplemented")
+}
+
+// DiscardValues releases the memory used by all values.
+// NOTE: any operation need value will panic after this function.
+func (t *ART) DiscardValues() {
+	panic("unimplemented")
+}
+
+// InspectStage used to inspect the value updates in the given stage.
+func (t *ART) InspectStage(handle int, f func([]byte, kv.KeyFlags, []byte)) {
+	panic("unimplemented")
+}
+
+// SelectValueHistory select the latest value which makes `predicate` returns true from the modification history.
+func (t *ART) SelectValueHistory(key []byte, predicate func(value []byte) bool) ([]byte, error) {
+	panic("unimplemented")
+}
+
+func (t *ART) SetMemoryFootprintChangeHook(fn func(uint64)) {
+	panic("unimplemented")
+}
+
+// MemHookSet implements the MemBuffer interface.
+func (t *ART) MemHookSet() bool {
+	panic("unimplemented")
+}
+
+// GetKeyByHandle returns key by handle.
+func (t *ART) GetKeyByHandle(handle arena.MemKeyHandle) []byte {
+	panic("unimplemented")
+}
+
+// GetValueByHandle returns value by handle.
+func (t *ART) GetValueByHandle(handle arena.MemKeyHandle) ([]byte, bool) {
+	panic("unimplemented")
+}
+
+func (t *ART) SetEntrySizeLimit(entryLimit, bufferLimit uint64) {
+	panic("unimplemented")
+}
+
+func (t *ART) RemoveFromBuffer(key []byte) {
+	panic("unimplemented")
+}

internal/unionstore/art/art_arena.go

@@ -0,0 +1,35 @@
+// Copyright 2024 TiKV Authors
+//
+// 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,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//nolint:unused
+package art
+
+import (
+	"github.com/tikv/client-go/v2/internal/unionstore/arena"
+)
+
+// fixedSizeArena is a fixed size arena allocator.
+// because the size of each type of node is fixed, the discarded nodes can be reused.
+// reusing blocks reduces the memory pieces.
+type nodeArena struct {
+	arena.MemdbArena
+	freeNode4  []arena.MemdbArenaAddr
+	freeNode16 []arena.MemdbArenaAddr
+	freeNode48 []arena.MemdbArenaAddr
+}
+
+type artAllocator struct {
+	vlogAllocator arena.MemdbVlog[*artLeaf, *ART]
+	nodeAllocator nodeArena
+}

internal/unionstore/art/art_iterator.go

@@ -0,0 +1,31 @@
+// Copyright 2024 TiKV Authors
+//
+// 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,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package art
+
+func (*ART) Iter([]byte, []byte) (*Iterator, error) {
+	panic("unimplemented")
+}
+
+func (*ART) IterReverse([]byte, []byte) (*Iterator, error) {
+	panic("unimplemented")
+}
+
+type Iterator struct{}
+
+func (i *Iterator) Valid() bool   { panic("unimplemented") }
+func (i *Iterator) Key() []byte   { panic("unimplemented") }
+func (i *Iterator) Value() []byte { panic("unimplemented") }
+func (i *Iterator) Next() error   { panic("unimplemented") }
+func (i *Iterator) Close()        { panic("unimplemented") }

internal/unionstore/art/art_node.go

@@ -0,0 +1,58 @@
+// Copyright 2024 TiKV Authors
+//
+// 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,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//nolint:unused
+package art
+
+import (
+	"github.com/tikv/client-go/v2/internal/unionstore/arena"
+	"github.com/tikv/client-go/v2/kv"
+)
+
+type artNodeKind uint16
+
+const (
+	typeARTInvalid artNodeKind = 0
+	//nolint:unused
+	typeARTNode4   artNodeKind = 1
+	typeARTNode16  artNodeKind = 2
+	typeARTNode48  artNodeKind = 3
+	typeARTNode256 artNodeKind = 4
+	typeARTLeaf    artNodeKind = 5
+)
+
+var nullArtNode = artNode{kind: typeARTInvalid, addr: arena.NullAddr}
+
+type artKey []byte
+
+type artNode struct {
+	kind artNodeKind
+	addr arena.MemdbArenaAddr
+}
+
+type artLeaf struct {
+	vAddr arena.MemdbArenaAddr
+	klen  uint16
+	flags uint16
+}
+
+// GetKey gets the full key of the leaf
+func (l *artLeaf) GetKey() []byte {
+	panic("unimplemented")
+}
+
+// GetKeyFlags gets the flags of the leaf
+func (l *artLeaf) GetKeyFlags() kv.KeyFlags {
+	panic("unimplemented")
+}