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