github.com/cockroachdb/pebble@v1.1.1-0.20240513155919-3622ade60459/internal/arenaskl/arena.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  	"sync/atomic"
    22  	"unsafe"
    23  
    24  	"github.com/cockroachdb/errors"
    25  	"github.com/cockroachdb/pebble/internal/constants"
    26  	"github.com/cockroachdb/pebble/internal/invariants"
    27  )
    28  
    29  // Arena is lock-free.
    30  type Arena struct {
    31  	n   atomic.Uint64
    32  	buf []byte
    33  }
    34  
    35  const nodeAlignment = 4
    36  
    37  var (
    38  	// ErrArenaFull indicates that the arena is full and cannot perform any more
    39  	// allocations.
    40  	ErrArenaFull = errors.New("allocation failed because arena is full")
    41  )
    42  
    43  // NewArena allocates a new arena using the specified buffer as the backing
    44  // store.
    45  func NewArena(buf []byte) *Arena {
    46  	if len(buf) > constants.MaxUint32OrInt {
    47  		if invariants.Enabled {
    48  			panic(errors.AssertionFailedf("attempting to create arena of size %d", len(buf)))
    49  		}
    50  		buf = buf[:constants.MaxUint32OrInt]
    51  	}
    52  	a := &Arena{
    53  		buf: buf,
    54  	}
    55  	// We don't store data at position 0 in order to reserve offset=0 as a kind of
    56  	// nil pointer.
    57  	a.n.Store(1)
    58  	return a
    59  }
    60  
    61  // Size returns the number of bytes allocated by the arena.
    62  func (a *Arena) Size() uint32 {
    63  	s := a.n.Load()
    64  	if s > constants.MaxUint32OrInt {
    65  		// The last failed allocation can push the size higher than len(a.buf).
    66  		// Saturate at the maximum representable offset.
    67  		return constants.MaxUint32OrInt
    68  	}
    69  	return uint32(s)
    70  }
    71  
    72  // Capacity returns the capacity of the arena.
    73  func (a *Arena) Capacity() uint32 {
    74  	return uint32(len(a.buf))
    75  }
    76  
    77  // alloc allocates a buffer of the given size and with the given alignment
    78  // (which must be a power of 2).
    79  //
    80  // If overflow is not 0, it also ensures that many bytes after the buffer are
    81  // inside the arena (this is used for structures that are larger than the
    82  // requested size but don't use those extra bytes).
    83  func (a *Arena) alloc(size, alignment, overflow uint32) (uint32, uint32, error) {
    84  	if invariants.Enabled && (alignment&(alignment-1)) != 0 {
    85  		panic(errors.AssertionFailedf("invalid alignment %d", alignment))
    86  	}
    87  	// Verify that the arena isn't already full.
    88  	origSize := a.n.Load()
    89  	if int(origSize) > len(a.buf) {
    90  		return 0, 0, ErrArenaFull
    91  	}
    92  
    93  	// Pad the allocation with enough bytes to ensure the requested alignment.
    94  	padded := uint64(size) + uint64(alignment) - 1
    95  
    96  	newSize := a.n.Add(padded)
    97  	if newSize+uint64(overflow) > uint64(len(a.buf)) {
    98  		return 0, 0, ErrArenaFull
    99  	}
   100  
   101  	// Return the aligned offset.
   102  	offset := (uint32(newSize) - size) & ^(alignment - 1)
   103  	return offset, uint32(padded), nil
   104  }
   105  
   106  func (a *Arena) getBytes(offset uint32, size uint32) []byte {
   107  	if offset == 0 {
   108  		return nil
   109  	}
   110  	return a.buf[offset : offset+size : offset+size]
   111  }
   112  
   113  func (a *Arena) getPointer(offset uint32) unsafe.Pointer {
   114  	if offset == 0 {
   115  		return nil
   116  	}
   117  	return unsafe.Pointer(&a.buf[offset])
   118  }
   119  
   120  func (a *Arena) getPointerOffset(ptr unsafe.Pointer) uint32 {
   121  	if ptr == nil {
   122  		return 0
   123  	}
   124  	return uint32(uintptr(ptr) - uintptr(unsafe.Pointer(&a.buf[0])))
   125  }