github.com/cockroachdb/pebble@v1.1.1-0.20240513155919-3622ade60459/internal/arenaskl/node.go (about)

     1  /*
     2   * Copyright 2017 Dgraph Labs, Inc. and Contributors
     3   * Modifications copyright (C) 2017 Andy Kimball and Contributors
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   */
    17  
    18  package arenaskl
    19  
    20  import (
    21  	"math"
    22  	"sync/atomic"
    23  
    24  	"github.com/cockroachdb/pebble/internal/base"
    25  )
    26  
    27  // MaxNodeSize returns the maximum space needed for a node with the specified
    28  // key and value sizes. This could overflow a uint32, which is why a uint64
    29  // is used here. If a key/value overflows a uint32, it should not be added to
    30  // the skiplist.
    31  func MaxNodeSize(keySize, valueSize uint32) uint64 {
    32  	const maxPadding = nodeAlignment - 1
    33  	return uint64(maxNodeSize) + uint64(keySize) + uint64(valueSize) + maxPadding
    34  }
    35  
    36  type links struct {
    37  	nextOffset atomic.Uint32
    38  	prevOffset atomic.Uint32
    39  }
    40  
    41  func (l *links) init(prevOffset, nextOffset uint32) {
    42  	l.nextOffset.Store(nextOffset)
    43  	l.prevOffset.Store(prevOffset)
    44  }
    45  
    46  type node struct {
    47  	// Immutable fields, so no need to lock to access key.
    48  	keyOffset  uint32
    49  	keySize    uint32
    50  	keyTrailer uint64
    51  	valueSize  uint32
    52  	allocSize  uint32
    53  
    54  	// Most nodes do not need to use the full height of the tower, since the
    55  	// probability of each successive level decreases exponentially. Because
    56  	// these elements are never accessed, they do not need to be allocated.
    57  	// Therefore, when a node is allocated in the arena, its memory footprint
    58  	// is deliberately truncated to not include unneeded tower elements.
    59  	//
    60  	// All accesses to elements should use CAS operations, with no need to lock.
    61  	tower [maxHeight]links
    62  }
    63  
    64  func newNode(
    65  	arena *Arena, height uint32, key base.InternalKey, value []byte,
    66  ) (nd *node, err error) {
    67  	if height < 1 || height > maxHeight {
    68  		panic("height cannot be less than one or greater than the max height")
    69  	}
    70  	keySize := len(key.UserKey)
    71  	if int64(keySize) > math.MaxUint32 {
    72  		panic("key is too large")
    73  	}
    74  	valueSize := len(value)
    75  	if int64(len(value)) > math.MaxUint32 {
    76  		panic("value is too large")
    77  	}
    78  	if int64(len(value))+int64(keySize)+int64(maxNodeSize) > math.MaxUint32 {
    79  		panic("combined key and value size is too large")
    80  	}
    81  
    82  	nd, err = newRawNode(arena, height, uint32(keySize), uint32(valueSize))
    83  	if err != nil {
    84  		return
    85  	}
    86  	nd.keyTrailer = key.Trailer
    87  	copy(nd.getKeyBytes(arena), key.UserKey)
    88  	copy(nd.getValue(arena), value)
    89  	return
    90  }
    91  
    92  func newRawNode(arena *Arena, height uint32, keySize, valueSize uint32) (nd *node, err error) {
    93  	// Compute the amount of the tower that will never be used, since the height
    94  	// is less than maxHeight.
    95  	unusedSize := uint32((maxHeight - int(height)) * linksSize)
    96  	nodeSize := uint32(maxNodeSize) - unusedSize
    97  
    98  	nodeOffset, allocSize, err := arena.alloc(nodeSize+keySize+valueSize, nodeAlignment, unusedSize)
    99  	if err != nil {
   100  		return
   101  	}
   102  
   103  	nd = (*node)(arena.getPointer(nodeOffset))
   104  	nd.keyOffset = nodeOffset + nodeSize
   105  	nd.keySize = keySize
   106  	nd.valueSize = valueSize
   107  	nd.allocSize = allocSize
   108  	return
   109  }
   110  
   111  func (n *node) getKeyBytes(arena *Arena) []byte {
   112  	return arena.getBytes(n.keyOffset, n.keySize)
   113  }
   114  
   115  func (n *node) getValue(arena *Arena) []byte {
   116  	return arena.getBytes(n.keyOffset+n.keySize, uint32(n.valueSize))
   117  }
   118  
   119  func (n *node) nextOffset(h int) uint32 {
   120  	return n.tower[h].nextOffset.Load()
   121  }
   122  
   123  func (n *node) prevOffset(h int) uint32 {
   124  	return n.tower[h].prevOffset.Load()
   125  }
   126  
   127  func (n *node) casNextOffset(h int, old, val uint32) bool {
   128  	return n.tower[h].nextOffset.CompareAndSwap(old, val)
   129  }
   130  
   131  func (n *node) casPrevOffset(h int, old, val uint32) bool {
   132  	return n.tower[h].prevOffset.CompareAndSwap(old, val)
   133  }