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 2 : func NewArena(buf []byte) *Arena { 46 2 : 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 2 : a := &Arena{ 53 2 : buf: buf, 54 2 : } 55 2 : // We don't store data at position 0 in order to reserve offset=0 as a kind of 56 2 : // nil pointer. 57 2 : a.n.Store(1) 58 2 : return a 59 : } 60 : 61 : // Size returns the number of bytes allocated by the arena. 62 2 : func (a *Arena) Size() uint32 { 63 2 : s := a.n.Load() 64 2 : 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 2 : return uint32(s) 70 : } 71 : 72 : // Capacity returns the capacity of the arena. 73 2 : func (a *Arena) Capacity() uint32 { 74 2 : return uint32(len(a.buf)) 75 2 : } 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 2 : func (a *Arena) alloc(size, alignment, overflow uint32) (uint32, uint32, error) { 84 2 : 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 2 : origSize := a.n.Load() 89 2 : 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 2 : padded := uint64(size) + uint64(alignment) - 1 95 2 : 96 2 : newSize := a.n.Add(padded) 97 2 : if newSize+uint64(overflow) > uint64(len(a.buf)) { 98 1 : return 0, 0, ErrArenaFull 99 1 : } 100 : 101 : // Return the aligned offset. 102 2 : offset := (uint32(newSize) - size) & ^(alignment - 1) 103 2 : return offset, uint32(padded), nil 104 : } 105 : 106 2 : func (a *Arena) getBytes(offset uint32, size uint32) []byte { 107 2 : if offset == 0 { 108 0 : return nil 109 0 : } 110 2 : return a.buf[offset : offset+size : offset+size] 111 : } 112 : 113 2 : func (a *Arena) getPointer(offset uint32) unsafe.Pointer { 114 2 : if offset == 0 { 115 0 : return nil 116 0 : } 117 2 : return unsafe.Pointer(&a.buf[offset]) 118 : } 119 : 120 2 : func (a *Arena) getPointerOffset(ptr unsafe.Pointer) uint32 { 121 2 : if ptr == nil { 122 0 : return 0 123 0 : } 124 2 : return uint32(uintptr(ptr) - uintptr(unsafe.Pointer(&a.buf[0]))) 125 : }