Line data    Source code

       1             : // Copyright 2023 The LevelDB-Go and Pebble Authors. All rights reserved. Use
       2             : // of this source code is governed by a BSD-style license that can be found in
       3             : // the LICENSE file.
       4             : 
       5             : package block
       6             : 
       7             : import (
       8             :         "github.com/cockroachdb/errors"
       9             :         "github.com/cockroachdb/pebble/internal/base"
      10             :         "github.com/cockroachdb/pebble/internal/cache"
      11             : )
      12             : 
      13             : // Alloc allocates a new Value for a block of length n (excluding the block
      14             : // trailer). If bufferPool is non-nil, Alloc allocates the buffer from the pool.
      15             : // Otherwise it allocates it from the block cache.
      16           1 : func Alloc(n int, p *BufferPool) Value {
      17           1 :         if p != nil {
      18           1 :                 return Value{buf: p.Alloc(int(n))}
      19           1 :         }
      20           1 :         return Value{v: cache.Alloc(int(n))}
      21             : }
      22             : 
      23             : // Value is a block buffer, either backed by the block cache or a BufferPool.
      24             : type Value struct {
      25             :         // buf.Valid() returns true if backed by a BufferPool.
      26             :         buf Buf
      27             :         // v is non-nil if backed by the block cache.
      28             :         v *cache.Value
      29             : }
      30             : 
      31             : // Get returns the underlying block's byte slice.
      32           1 : func (b Value) Get() []byte {
      33           1 :         if b.buf.Valid() {
      34           1 :                 return b.buf.p.pool[b.buf.i].b
      35           1 :         }
      36           1 :         return b.v.Buf()
      37             : }
      38             : 
      39             : // MakeHandle constructs a BufferHandle from the Value. If the Value is not
      40             : // backed by a buffer pool, MakeHandle inserts the value into the block cache,
      41             : // returning a handle to the now resident value.
      42             : func (b Value) MakeHandle(
      43             :         c *cache.Cache, cacheID uint64, fileNum base.DiskFileNum, offset uint64,
      44           1 : ) BufferHandle {
      45           1 :         if b.buf.Valid() {
      46           1 :                 return BufferHandle{b: b.buf}
      47           1 :         }
      48           1 :         return BufferHandle{h: c.Set(cacheID, fileNum, offset, b.v)}
      49             : }
      50             : 
      51             : // Release releases the handle.
      52           1 : func (b Value) Release() {
      53           1 :         if b.buf.Valid() {
      54           1 :                 b.buf.Release()
      55           1 :         } else {
      56           1 :                 cache.Free(b.v)
      57           1 :         }
      58             : }
      59             : 
      60             : // Truncate truncates the block to n bytes.
      61           1 : func (b Value) Truncate(n int) {
      62           1 :         if b.buf.Valid() {
      63           1 :                 b.buf.p.pool[b.buf.i].b = b.buf.p.pool[b.buf.i].b[:n]
      64           1 :         } else {
      65           1 :                 b.v.Truncate(n)
      66           1 :         }
      67             : }
      68             : 
      69             : // A BufferHandle is a handle to manually-managed memory. The handle may point
      70             : // to a block in the block cache (h.Get() != nil), or a buffer that exists
      71             : // outside the block cache allocated from a BufferPool (b.Valid()).
      72             : type BufferHandle struct {
      73             :         h cache.Handle
      74             :         b Buf
      75             : }
      76             : 
      77             : // CacheBufferHandle constructs a BufferHandle from a block cache Handle.
      78           1 : func CacheBufferHandle(h cache.Handle) BufferHandle {
      79           1 :         return BufferHandle{h: h}
      80           1 : }
      81             : 
      82             : // Get retrieves the underlying buffer referenced by the handle.
      83           1 : func (bh BufferHandle) Get() []byte {
      84           1 :         if v := bh.h.Get(); v != nil {
      85           1 :                 return v
      86           1 :         } else if bh.b.p != nil {
      87           1 :                 return bh.b.p.pool[bh.b.i].b
      88           1 :         }
      89           1 :         return nil
      90             : }
      91             : 
      92             : // Release releases the buffer, either back to the block cache or BufferPool.
      93           1 : func (bh BufferHandle) Release() {
      94           1 :         bh.h.Release()
      95           1 :         bh.b.Release()
      96           1 : }
      97             : 
      98             : // A BufferPool holds a pool of buffers for holding sstable blocks. An initial
      99             : // size of the pool is provided on Init, but a BufferPool will grow to meet the
     100             : // largest working set size. It'll never shrink. When a buffer is released, the
     101             : // BufferPool recycles the buffer for future allocations.
     102             : //
     103             : // A BufferPool should only be used for short-lived allocations with
     104             : // well-understood working set sizes to avoid excessive memory consumption.
     105             : //
     106             : // BufferPool is not thread-safe.
     107             : type BufferPool struct {
     108             :         // pool contains all the buffers held by the pool, including buffers that
     109             :         // are in-use. For every i < len(pool): pool[i].v is non-nil.
     110             :         pool []AllocedBuffer
     111             : }
     112             : 
     113             : // AllocedBuffer is an allocated memory buffer.
     114             : type AllocedBuffer struct {
     115             :         v *cache.Value
     116             :         // b holds the current byte slice. It's backed by v, but may be a subslice
     117             :         // of v's memory while the buffer is in-use [ len(b) ≤ len(v.Buf()) ].
     118             :         //
     119             :         // If the buffer is not currently in-use, b is nil. When being recycled, the
     120             :         // BufferPool.Alloc will reset b to be a subslice of v.Buf().
     121             :         b []byte
     122             : }
     123             : 
     124             : // Init initializes the pool with an initial working set buffer size of
     125             : // `initialSize`.
     126           1 : func (p *BufferPool) Init(initialSize int) {
     127           1 :         *p = BufferPool{
     128           1 :                 pool: make([]AllocedBuffer, 0, initialSize),
     129           1 :         }
     130           1 : }
     131             : 
     132             : // InitPreallocated is like Init but for internal sstable package use in
     133             : // instances where a pre-allocated slice of []allocedBuffer already exists. It's
     134             : // used to avoid an extra allocation initializing BufferPool.pool.
     135           1 : func (p *BufferPool) InitPreallocated(pool []AllocedBuffer) {
     136           1 :         *p = BufferPool{
     137           1 :                 pool: pool[:0],
     138           1 :         }
     139           1 : }
     140             : 
     141             : // Release releases all buffers held by the pool and resets the pool to an
     142             : // uninitialized state.
     143           1 : func (p *BufferPool) Release() {
     144           1 :         for i := range p.pool {
     145           1 :                 if p.pool[i].b != nil {
     146           0 :                         panic(errors.AssertionFailedf("Release called on a BufferPool with in-use buffers"))
     147             :                 }
     148           1 :                 cache.Free(p.pool[i].v)
     149             :         }
     150           1 :         *p = BufferPool{}
     151             : }
     152             : 
     153             : // Alloc allocates a new buffer of size n. If the pool already holds a buffer at
     154             : // least as large as n, the pooled buffer is used instead.
     155             : //
     156             : // Alloc is O(MAX(N,M)) where N is the largest number of concurrently in-use
     157             : // buffers allocated and M is the initialSize passed to Init.
     158           1 : func (p *BufferPool) Alloc(n int) Buf {
     159           1 :         unusableBufferIdx := -1
     160           1 :         for i := 0; i < len(p.pool); i++ {
     161           1 :                 if p.pool[i].b == nil {
     162           1 :                         if len(p.pool[i].v.Buf()) >= n {
     163           1 :                                 p.pool[i].b = p.pool[i].v.Buf()[:n]
     164           1 :                                 return Buf{p: p, i: i}
     165           1 :                         }
     166           1 :                         unusableBufferIdx = i
     167             :                 }
     168             :         }
     169             : 
     170             :         // If we would need to grow the size of the pool to allocate another buffer,
     171             :         // but there was a slot available occupied by a buffer that's just too
     172             :         // small, replace the too-small buffer.
     173           1 :         if len(p.pool) == cap(p.pool) && unusableBufferIdx >= 0 {
     174           1 :                 i := unusableBufferIdx
     175           1 :                 cache.Free(p.pool[i].v)
     176           1 :                 p.pool[i].v = cache.Alloc(n)
     177           1 :                 p.pool[i].b = p.pool[i].v.Buf()
     178           1 :                 return Buf{p: p, i: i}
     179           1 :         }
     180             : 
     181             :         // Allocate a new buffer.
     182           1 :         v := cache.Alloc(n)
     183           1 :         p.pool = append(p.pool, AllocedBuffer{v: v, b: v.Buf()[:n]})
     184           1 :         return Buf{p: p, i: len(p.pool) - 1}
     185             : }
     186             : 
     187             : // A Buf holds a reference to a manually-managed, pooled byte buffer.
     188             : type Buf struct {
     189             :         p *BufferPool
     190             :         // i holds the index into p.pool where the buffer may be found. This scheme
     191             :         // avoids needing to allocate the handle to the buffer on the heap at the
     192             :         // cost of copying two words instead of one.
     193             :         i int
     194             : }
     195             : 
     196             : // Valid returns true if the buf holds a valid buffer.
     197           1 : func (b Buf) Valid() bool {
     198           1 :         return b.p != nil
     199           1 : }
     200             : 
     201             : // Release releases the buffer back to the pool.
     202           1 : func (b *Buf) Release() {
     203           1 :         if b.p == nil {
     204           1 :                 return
     205           1 :         }
     206             :         // Clear the allocedBuffer's byte slice. This signals the allocated buffer
     207             :         // is no longer in use and a future call to BufferPool.Alloc may reuse this
     208             :         // buffer.
     209           1 :         b.p.pool[b.i].b = nil
     210           1 :         b.p = nil
     211             : }