Line data    Source code

       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           1 : func NewArena(buf []byte) *Arena {
      46           1 :         if len(buf) > constants.MaxUint32OrInt {
      47           0 :                 if invariants.Enabled {
      48           0 :                         panic(errors.AssertionFailedf("attempting to create arena of size %d", len(buf)))
      49             :                 }
      50           0 :                 buf = buf[:constants.MaxUint32OrInt]
      51             :         }
      52           1 :         a := &Arena{
      53           1 :                 buf: buf,
      54           1 :         }
      55           1 :         // We don't store data at position 0 in order to reserve offset=0 as a kind of
      56           1 :         // nil pointer.
      57           1 :         a.n.Store(1)
      58           1 :         return a
      59             : }
      60             : 
      61             : // Size returns the number of bytes allocated by the arena.
      62           1 : func (a *Arena) Size() uint32 {
      63           1 :         s := a.n.Load()
      64           1 :         if s > constants.MaxUint32OrInt {
      65           1 :                 // The last failed allocation can push the size higher than len(a.buf).
      66           1 :                 // Saturate at the maximum representable offset.
      67           1 :                 return constants.MaxUint32OrInt
      68           1 :         }
      69           1 :         return uint32(s)
      70             : }
      71             : 
      72             : // Capacity returns the capacity of the arena.
      73           1 : func (a *Arena) Capacity() uint32 {
      74           1 :         return uint32(len(a.buf))
      75           1 : }
      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           1 : func (a *Arena) alloc(size, alignment, overflow uint32) (uint32, uint32, error) {
      84           1 :         if invariants.Enabled && (alignment&(alignment-1)) != 0 {
      85           0 :                 panic(errors.AssertionFailedf("invalid alignment %d", alignment))
      86             :         }
      87             :         // Verify that the arena isn't already full.
      88           1 :         origSize := a.n.Load()
      89           1 :         if int(origSize) > len(a.buf) {
      90           1 :                 return 0, 0, ErrArenaFull
      91           1 :         }
      92             : 
      93             :         // Pad the allocation with enough bytes to ensure the requested alignment.
      94           1 :         padded := uint64(size) + uint64(alignment) - 1
      95           1 : 
      96           1 :         newSize := a.n.Add(padded)
      97           1 :         if newSize+uint64(overflow) > uint64(len(a.buf)) {
      98           1 :                 return 0, 0, ErrArenaFull
      99           1 :         }
     100             : 
     101             :         // Return the aligned offset.
     102           1 :         offset := (uint32(newSize) - size) & ^(alignment - 1)
     103           1 :         return offset, uint32(padded), nil
     104             : }
     105             : 
     106           1 : func (a *Arena) getBytes(offset uint32, size uint32) []byte {
     107           1 :         if offset == 0 {
     108           0 :                 return nil
     109           0 :         }
     110           1 :         return a.buf[offset : offset+size : offset+size]
     111             : }
     112             : 
     113           1 : func (a *Arena) getPointer(offset uint32) unsafe.Pointer {
     114           1 :         if offset == 0 {
     115           0 :                 return nil
     116           0 :         }
     117           1 :         return unsafe.Pointer(&a.buf[offset])
     118             : }
     119             : 
     120           1 : func (a *Arena) getPointerOffset(ptr unsafe.Pointer) uint32 {
     121           1 :         if ptr == nil {
     122           0 :                 return 0
     123           0 :         }
     124           1 :         return uint32(uintptr(ptr) - uintptr(unsafe.Pointer(&a.buf[0])))
     125             : }