github.com/matrixorigin/matrixone@v1.2.0/pkg/common/arenaskl/arena.go (about)

     1  /*
     2   * Copyright 2017 Dgraph Labs, Inc. and Contributors
     3   * Modifications copyright (C) 2017 Andy Kimball and Contributors
     4   * and copyright (C) 2024 MatrixOrigin Inc.
     5   *
     6   * Licensed under the Apache License, Version 2.0 (the "License");
     7   * you may not use this file except in compliance with the License.
     8   * You may obtain a copy of the License at
     9   *
    10   *     http://www.apache.org/licenses/LICENSE-2.0
    11   *
    12   * Unless required by applicable law or agreed to in writing, software
    13   * distributed under the License is distributed on an "AS IS" BASIS,
    14   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    15   * See the License for the specific language governing permissions and
    16   * limitations under the License.
    17   */
    18  
    19  package arenaskl
    20  
    21  import (
    22  	"fmt"
    23  	"math"
    24  	"sync/atomic"
    25  	"unsafe"
    26  
    27  	"github.com/matrixorigin/matrixone/pkg/common/moerr"
    28  )
    29  
    30  const (
    31  	maxArenaSize = math.MaxUint32
    32  )
    33  
    34  // Arena is lock-free.
    35  type Arena struct {
    36  	n   atomic.Uint64
    37  	buf []byte
    38  }
    39  
    40  const nodeAlignment = 4
    41  
    42  var (
    43  	// ErrArenaFull indicates that the arena is full and cannot perform any more
    44  	// allocations.
    45  	ErrArenaFull = moerr.NewArenaFullNoCtx()
    46  )
    47  
    48  // NewArena allocates a new arena using the specified buffer as the backing
    49  // store.
    50  func NewArena(buf []byte) *Arena {
    51  	if len(buf) > maxArenaSize {
    52  		panic(fmt.Sprintf("attempting to create arena of size %d", len(buf)))
    53  	}
    54  	a := &Arena{
    55  		buf: buf,
    56  	}
    57  	// We don't store data at position 0 in order to reserve offset=0 as a kind of
    58  	// nil pointer.
    59  	a.n.Store(1)
    60  	return a
    61  }
    62  
    63  // Size returns the number of bytes allocated by the arena.
    64  func (a *Arena) Size() uint32 {
    65  	s := a.n.Load()
    66  	if s > maxArenaSize {
    67  		// The last failed allocation can push the size higher than len(a.buf).
    68  		// Saturate at the maximum representable offset.
    69  		return maxArenaSize
    70  	}
    71  	return uint32(s)
    72  }
    73  
    74  // Capacity returns the capacity of the arena.
    75  func (a *Arena) Capacity() uint32 {
    76  	return uint32(len(a.buf))
    77  }
    78  
    79  // alloc allocates a buffer of the given size and with the given alignment
    80  // (which must be a power of 2).
    81  //
    82  // If overflow is not 0, it also ensures that many bytes after the buffer are
    83  // inside the arena (this is used for structures that are larger than the
    84  // requested size but don't use those extra bytes).
    85  func (a *Arena) alloc(size, alignment, overflow uint32) (uint32, uint32, error) {
    86  	if (alignment & (alignment - 1)) != 0 {
    87  		panic(fmt.Sprintf("invalid alignment %d", alignment))
    88  	}
    89  	// Verify that the arena isn't already full.
    90  	origSize := a.n.Load()
    91  	if int(origSize) > len(a.buf) {
    92  		return 0, 0, ErrArenaFull
    93  	}
    94  
    95  	// Pad the allocation with enough bytes to ensure the requested alignment.
    96  	padded := uint64(size) + uint64(alignment) - 1
    97  
    98  	newSize := a.n.Add(padded)
    99  	if newSize+uint64(overflow) > uint64(len(a.buf)) {
   100  		return 0, 0, ErrArenaFull
   101  	}
   102  
   103  	// Return the aligned offset.
   104  	offset := (uint32(newSize) - size) & ^(alignment - 1)
   105  	return offset, uint32(padded), nil
   106  }
   107  
   108  func (a *Arena) getBytes(offset uint32, size uint32) []byte {
   109  	if offset == 0 {
   110  		return nil
   111  	}
   112  	return a.buf[offset : offset+size : offset+size]
   113  }
   114  
   115  func (a *Arena) getPointer(offset uint32) unsafe.Pointer {
   116  	if offset == 0 {
   117  		return nil
   118  	}
   119  	return unsafe.Pointer(&a.buf[offset])
   120  }
   121  
   122  func (a *Arena) getPointerOffset(ptr unsafe.Pointer) uint32 {
   123  	if ptr == nil {
   124  		return 0
   125  	}
   126  	return uint32(uintptr(ptr) - uintptr(unsafe.Pointer(&a.buf[0])))
   127  }