client-go/internal/unionstore/rbt/rbt_arena.go

102 lines
2.9 KiB
Go

// 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 rbt
import (
"unsafe"
"github.com/tikv/client-go/v2/internal/unionstore/arena"
)
type nodeAllocator struct {
arena.MemdbArena
// Dummy node, so that we can make X.left.up = X.
// We then use this instead of NULL to mean the top or bottom
// end of the rb tree. It is a black node.
nullNode memdbNode
}
func (a *nodeAllocator) init() {
a.nullNode = memdbNode{
up: arena.NullAddr,
left: arena.NullAddr,
right: arena.NullAddr,
vptr: arena.NullAddr,
}
}
func (a *nodeAllocator) getNode(addr arena.MemdbArenaAddr) *memdbNode {
if addr.IsNull() {
return &a.nullNode
}
data := a.GetData(addr)
return (*memdbNode)(unsafe.Pointer(&data[0]))
}
const memdbNodeSize = int(unsafe.Sizeof(memdbNode{}))
func (a *nodeAllocator) allocNode(key []byte) (arena.MemdbArenaAddr, *memdbNode) {
nodeSize := memdbNodeSize + len(key)
prevBlocks := a.Blocks()
addr, mem := a.Alloc(nodeSize, true)
n := (*memdbNode)(unsafe.Pointer(&mem[0]))
n.vptr = arena.NullAddr
n.klen = uint16(len(key))
copy(n.getKey(), key)
if prevBlocks != a.Blocks() {
a.OnMemChange()
}
return addr, n
}
func (a *nodeAllocator) freeNode(addr arena.MemdbArenaAddr) {
if testMode {
// Make it easier for debug.
n := a.getNode(addr)
n.left = arena.BadAddr
n.right = arena.BadAddr
n.up = arena.BadAddr
n.vptr = arena.BadAddr
return
}
// TODO: reuse freed nodes. Need to fix lastTraversedNode when implementing this.
}
func (a *nodeAllocator) reset() {
a.MemdbArena.Reset()
a.init()
}