github.com/MetalBlockchain/metalgo@v1.11.9/database/pebbledb/batch.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package pebbledb
     5  
     6  import (
     7  	"fmt"
     8  
     9  	"github.com/cockroachdb/pebble"
    10  
    11  	"github.com/MetalBlockchain/metalgo/database"
    12  )
    13  
    14  var _ database.Batch = (*batch)(nil)
    15  
    16  // Not safe for concurrent use.
    17  type batch struct {
    18  	batch *pebble.Batch
    19  	db    *Database
    20  	size  int
    21  
    22  	// True iff [batch] has been written to the database
    23  	// since the last time [Reset] was called.
    24  	written bool
    25  }
    26  
    27  func (db *Database) NewBatch() database.Batch {
    28  	return &batch{
    29  		db:    db,
    30  		batch: db.pebbleDB.NewBatch(),
    31  	}
    32  }
    33  
    34  func (b *batch) Put(key, value []byte) error {
    35  	b.size += len(key) + len(value) + pebbleByteOverHead
    36  	return b.batch.Set(key, value, pebble.Sync)
    37  }
    38  
    39  func (b *batch) Delete(key []byte) error {
    40  	b.size += len(key) + pebbleByteOverHead
    41  	return b.batch.Delete(key, pebble.Sync)
    42  }
    43  
    44  func (b *batch) Size() int {
    45  	return b.size
    46  }
    47  
    48  // Assumes [b.db.lock] is not held.
    49  func (b *batch) Write() error {
    50  	b.db.lock.RLock()
    51  	defer b.db.lock.RUnlock()
    52  
    53  	// Committing to a closed database makes pebble panic
    54  	// so make sure [b.db] isn't closed.
    55  	if b.db.closed {
    56  		return database.ErrClosed
    57  	}
    58  
    59  	if b.written {
    60  		// pebble doesn't support writing a batch twice so we have to clone the
    61  		// batch before writing it.
    62  		newBatch := b.db.pebbleDB.NewBatch()
    63  		if err := newBatch.Apply(b.batch, nil); err != nil {
    64  			return err
    65  		}
    66  		b.batch = newBatch
    67  	}
    68  
    69  	b.written = true
    70  	return updateError(b.batch.Commit(pebble.Sync))
    71  }
    72  
    73  func (b *batch) Reset() {
    74  	b.batch.Reset()
    75  	b.written = false
    76  	b.size = 0
    77  }
    78  
    79  func (b *batch) Replay(w database.KeyValueWriterDeleter) error {
    80  	reader := b.batch.Reader()
    81  	for {
    82  		kind, k, v, ok := reader.Next()
    83  		if !ok {
    84  			return nil
    85  		}
    86  		switch kind {
    87  		case pebble.InternalKeyKindSet:
    88  			if err := w.Put(k, v); err != nil {
    89  				return err
    90  			}
    91  		case pebble.InternalKeyKindDelete:
    92  			if err := w.Delete(k); err != nil {
    93  				return err
    94  			}
    95  		default:
    96  			return fmt.Errorf("%w: %v", errInvalidOperation, kind)
    97  		}
    98  	}
    99  }
   100  
   101  func (b *batch) Inner() database.Batch {
   102  	return b
   103  }