github.com/koko1123/flow-go-1@v0.29.6/storage/badger/batch.go (about)

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