github.com/pingcap/badger@v1.5.1-0.20230103063557-828f39b09b6d/table/memtable/arena.go

/*
 * Copyright 2017 Dgraph Labs, Inc. and Contributors
 *
 * 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 memtable

import (
	"encoding/binary"
	"sync/atomic"
	"unsafe"

	"github.com/pingcap/badger/y"
)

const (
	offsetSize = int(unsafe.Sizeof(uint32(0)))

	// Always align nodes on 64-bit boundaries, even on 32-bit architectures,
	// so that the node.value field is 64-bit aligned. This is necessary because
	// node.getValueAddr uses atomic.LoadUint64, which expects its input
	// pointer to be 64-bit aligned.
	nodeAlign = int(unsafe.Sizeof(uint64(0))) - 1
)

// Arena should be lock-free.
type arena struct {
	n   uint32
	buf []byte
}

// newArena returns a new arena.
func newArena(n int64) *arena {
	// Don't store data at position 0 in order to reserve offset=0 as a kind
	// of nil pointer.
	out := &arena{
		n:   1,
		buf: make([]byte, n),
	}
	return out
}

func (s *arena) size() int64 {
	return int64(atomic.LoadUint32(&s.n))
}

func (s *arena) reset() {
	atomic.StoreUint32(&s.n, 0)
}

// putNode allocates a node in the arena. The node is aligned on a pointer-sized
// boundary. The arena offset of the node is returned.
func (s *arena) putNode(height int) uint32 {
	// Compute the amount of the tower that will never be used, since the height
	// is less than maxHeight.
	unusedSize := (maxHeight - height) * offsetSize

	// Pad the allocation with enough bytes to ensure pointer alignment.
	l := uint32(MaxNodeSize - unusedSize + nodeAlign)
	n := atomic.AddUint32(&s.n, l)
	y.Assert(int(n) <= len(s.buf))

	// Return the aligned offset.
	m := (n - l + uint32(nodeAlign)) & ^uint32(nodeAlign)
	return m
}

// Put will *copy* val into arena. To make better use of this, reuse your input
// val buffer. Returns an offset into buf. User is responsible for remembering
// size of val. We could also store this size inside arena but the encoding and
// decoding will incur some overhead.
func (s *arena) putVal(v y.ValueStruct) uint32 {
	l := v.EncodedSize()
	n := atomic.AddUint32(&s.n, l)
	y.Assert(int(n) <= len(s.buf))
	m := n - l
	v.Encode(s.buf[m:])
	return m
}

func (s *arena) putKey(key []byte) uint32 {
	l := uint32(len(key))
	n := atomic.AddUint32(&s.n, l)
	y.Assert(int(n) <= len(s.buf))
	m := n - l
	copy(s.buf[m:], key)
	return m
}

// getNode returns a pointer to the node located at offset. If the offset is
// zero, then the nil node pointer is returned.
func (s *arena) getNode(offset uint32) *node {
	if offset == 0 {
		return nil
	}

	return (*node)(unsafe.Pointer(&s.buf[offset]))
}

// getKey returns byte slice at offset.
func (s *arena) getKey(offset uint32, size uint16) (k []byte) {
	return s.buf[offset : offset+uint32(size)]
}

// getVal returns byte slice at offset. The given size should be just the value
// size and should NOT include the meta bytes.
func (s *arena) getVal(offset uint32, size uint32) (ret y.ValueStruct) {
	ret.Decode(s.buf[offset : offset+size])
	return
}

func (s *arena) fillVal(vs *y.ValueStruct, offset uint32, size uint32) {
	vs.Decode(s.buf[offset : offset+size])
}

// getNodeOffset returns the offset of node in the arena. If the node pointer is
// nil, then the zero offset is returned.
func (s *arena) getNodeOffset(nd *node) uint32 {
	if nd == nil {
		return 0
	}

	return uint32(uintptr(unsafe.Pointer(nd)) - uintptr(unsafe.Pointer(&s.buf[0])))
}

const valueNodeSize = uint32(unsafe.Sizeof(valueNode{}))

// valueNode is used to store multiple versions of a key.
// We enforce that the newly put value must have a greater version, so the single linked list is in descending order
// by version.
type valueNode struct {
	valAddr     uint64
	nextValAddr uint64
}

func (vn valueNode) encode(b []byte) {
	binary.LittleEndian.PutUint64(b, vn.valAddr)
	binary.LittleEndian.PutUint64(b[8:], vn.nextValAddr)
}

func (vn *valueNode) decode(b []byte) {
	vn.valAddr = binary.LittleEndian.Uint64(b)
	vn.nextValAddr = binary.LittleEndian.Uint64(b[8:])
}

func (s *arena) putValueNode(vn valueNode) uint32 {
	n := atomic.AddUint32(&s.n, valueNodeSize)
	y.Assert(int(n) <= len(s.buf))
	m := n - valueNodeSize
	vn.encode(s.buf[m:])
	return m
}

func (s *arena) getValueNode(offset uint32) valueNode {
	var vl valueNode
	vl.decode(s.buf[offset : offset+valueNodeSize])
	return vl
}