github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/storage/badger/batch.go (about)

     1  package badger
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/dgraph-io/badger/v2"
     7  )
     8  
     9  type BatchBuilder interface {
    10  	NewWriteBatch() *badger.WriteBatch
    11  }
    12  
    13  type Batch struct {
    14  	writer *badger.WriteBatch
    15  
    16  	lock      sync.RWMutex
    17  	callbacks []func()
    18  }
    19  
    20  func NewBatch(db BatchBuilder) *Batch {
    21  	batch := db.NewWriteBatch()
    22  	return &Batch{
    23  		writer:    batch,
    24  		callbacks: make([]func(), 0),
    25  	}
    26  }
    27  
    28  func (b *Batch) GetWriter() *badger.WriteBatch {
    29  	return b.writer
    30  }
    31  
    32  // OnSucceed adds a callback to execute after the batch has
    33  // been successfully flushed.
    34  // useful for implementing the cache where we will only cache
    35  // after the batch has been successfully flushed
    36  func (b *Batch) OnSucceed(callback func()) {
    37  	b.lock.Lock()
    38  	defer b.lock.Unlock()
    39  	b.callbacks = append(b.callbacks, callback)
    40  }
    41  
    42  // Flush will call the badger Batch's Flush method, in
    43  // addition, it will call the callbacks added by
    44  // OnSucceed
    45  func (b *Batch) Flush() error {
    46  	err := b.writer.Flush()
    47  	if err != nil {
    48  		return err
    49  	}
    50  
    51  	b.lock.RLock()
    52  	defer b.lock.RUnlock()
    53  	for _, callback := range b.callbacks {
    54  		callback()
    55  	}
    56  	return nil
    57  }