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

     1  package transaction
     2  
     3  import (
     4  	dbbadger "github.com/dgraph-io/badger/v3"
     5  
     6  	ioutils "github.com/koko1123/flow-go-1/utils/io"
     7  )
     8  
     9  type Tx struct {
    10  	DBTxn     *dbbadger.Txn
    11  	callbacks []func()
    12  }
    13  
    14  // OnSucceed adds a callback to execute after the batch has
    15  // been successfully flushed.
    16  // useful for implementing the cache where we will only cache
    17  // after the batch has been successfully flushed
    18  func (b *Tx) OnSucceed(callback func()) {
    19  	b.callbacks = append(b.callbacks, callback)
    20  }
    21  
    22  // Update creates a badger transaction, passing it to a chain of functions,
    23  // if all succeed. Useful to use callback to update cache in order to ensure data
    24  // in badgerDB and cache are consistent.
    25  func Update(db *dbbadger.DB, f func(*Tx) error) error {
    26  	dbTxn := db.NewTransaction(true)
    27  	defer dbTxn.Discard()
    28  
    29  	tx := &Tx{DBTxn: dbTxn}
    30  	err := f(tx)
    31  	if err != nil {
    32  		return err
    33  	}
    34  
    35  	err = dbTxn.Commit()
    36  	if err != nil {
    37  		return ioutils.TerminateOnFullDisk(err)
    38  	}
    39  
    40  	for _, callback := range tx.callbacks {
    41  		callback()
    42  	}
    43  	return nil
    44  }
    45  
    46  // WithTx is useful when transaction is used without adding callback.
    47  func WithTx(f func(*dbbadger.Txn) error) func(*Tx) error {
    48  	return func(tx *Tx) error {
    49  		return f(tx.DBTxn)
    50  	}
    51  }