github.com/zuoyebang/bitalostable@v1.0.1-0.20240229032404-e3b99a834294/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  	"math"
    22  	"sync/atomic"
    23  	"unsafe"
    24  
    25  	"github.com/cockroachdb/errors"
    26  )
    27  
    28  // Arena is lock-free.
    29  type Arena struct {
    30  	n   uint64
    31  	buf []byte
    32  }
    33  
    34  const (
    35  	align4 = 3
    36  )
    37  
    38  var (
    39  	// ErrArenaFull indicates that the arena is full and cannot perform any more
    40  	// allocations.
    41  	ErrArenaFull = errors.New("allocation failed because arena is full")
    42  )
    43  
    44  // NewArena allocates a new arena using the specified buffer as the backing
    45  // store.
    46  func NewArena(buf []byte) *Arena {
    47  	// Don't store data at position 0 in order to reserve offset=0 as a kind
    48  	// of nil pointer.
    49  	return &Arena{
    50  		n:   1,
    51  		buf: buf,
    52  	}
    53  }
    54  
    55  // Size returns the number of bytes allocated by the arena.
    56  func (a *Arena) Size() uint32 {
    57  	s := atomic.LoadUint64(&a.n)
    58  	if s > math.MaxUint32 {
    59  		// Saturate at MaxUint32.
    60  		return math.MaxUint32
    61  	}
    62  	return uint32(s)
    63  }
    64  
    65  // Capacity returns the capacity of the arena.
    66  func (a *Arena) Capacity() uint32 {
    67  	return uint32(len(a.buf))
    68  }
    69  
    70  func (a *Arena) alloc(size, align, overflow uint32) (uint32, uint32, error) {
    71  	// Verify that the arena isn't already full.
    72  	origSize := atomic.LoadUint64(&a.n)
    73  	if int(origSize) > len(a.buf) {
    74  		return 0, 0, ErrArenaFull
    75  	}
    76  
    77  	// Pad the allocation with enough bytes to ensure the requested alignment.
    78  	padded := uint32(size) + align
    79  
    80  	newSize := atomic.AddUint64(&a.n, uint64(padded))
    81  	if int(newSize)+int(overflow) > len(a.buf) {
    82  		return 0, 0, ErrArenaFull
    83  	}
    84  
    85  	// Return the aligned offset.
    86  	offset := (uint32(newSize) - padded + align) & ^align
    87  	return offset, padded, nil
    88  }
    89  
    90  func (a *Arena) getBytes(offset uint32, size uint32) []byte {
    91  	if offset == 0 {
    92  		return nil
    93  	}
    94  	return a.buf[offset : offset+size : offset+size]
    95  }
    96  
    97  func (a *Arena) getPointer(offset uint32) unsafe.Pointer {
    98  	if offset == 0 {
    99  		return nil
   100  	}
   101  	return unsafe.Pointer(&a.buf[offset])
   102  }
   103  
   104  func (a *Arena) getPointerOffset(ptr unsafe.Pointer) uint32 {
   105  	if ptr == nil {
   106  		return 0
   107  	}
   108  	return uint32(uintptr(ptr) - uintptr(unsafe.Pointer(&a.buf[0])))
   109  }