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

     1  package badger
     2  
     3  import (
     4  	"github.com/dgraph-io/badger/v2"
     5  
     6  	"github.com/onflow/flow-go/model/flow"
     7  	"github.com/onflow/flow-go/module"
     8  	"github.com/onflow/flow-go/module/metrics"
     9  	"github.com/onflow/flow-go/storage/badger/operation"
    10  	"github.com/onflow/flow-go/storage/badger/transaction"
    11  )
    12  
    13  // Guarantees implements persistent storage for collection guarantees.
    14  type Guarantees struct {
    15  	db    *badger.DB
    16  	cache *Cache[flow.Identifier, *flow.CollectionGuarantee]
    17  }
    18  
    19  func NewGuarantees(collector module.CacheMetrics, db *badger.DB, cacheSize uint) *Guarantees {
    20  
    21  	store := func(collID flow.Identifier, guarantee *flow.CollectionGuarantee) func(*transaction.Tx) error {
    22  		return transaction.WithTx(operation.SkipDuplicates(operation.InsertGuarantee(collID, guarantee)))
    23  	}
    24  
    25  	retrieve := func(collID flow.Identifier) func(*badger.Txn) (*flow.CollectionGuarantee, error) {
    26  		var guarantee flow.CollectionGuarantee
    27  		return func(tx *badger.Txn) (*flow.CollectionGuarantee, error) {
    28  			err := operation.RetrieveGuarantee(collID, &guarantee)(tx)
    29  			return &guarantee, err
    30  		}
    31  	}
    32  
    33  	g := &Guarantees{
    34  		db: db,
    35  		cache: newCache[flow.Identifier, *flow.CollectionGuarantee](collector, metrics.ResourceGuarantee,
    36  			withLimit[flow.Identifier, *flow.CollectionGuarantee](cacheSize),
    37  			withStore(store),
    38  			withRetrieve(retrieve)),
    39  	}
    40  
    41  	return g
    42  }
    43  
    44  func (g *Guarantees) storeTx(guarantee *flow.CollectionGuarantee) func(*transaction.Tx) error {
    45  	return g.cache.PutTx(guarantee.ID(), guarantee)
    46  }
    47  
    48  func (g *Guarantees) retrieveTx(collID flow.Identifier) func(*badger.Txn) (*flow.CollectionGuarantee, error) {
    49  	return func(tx *badger.Txn) (*flow.CollectionGuarantee, error) {
    50  		val, err := g.cache.Get(collID)(tx)
    51  		if err != nil {
    52  			return nil, err
    53  		}
    54  		return val, nil
    55  	}
    56  }
    57  
    58  func (g *Guarantees) Store(guarantee *flow.CollectionGuarantee) error {
    59  	return operation.RetryOnConflictTx(g.db, transaction.Update, g.storeTx(guarantee))
    60  }
    61  
    62  func (g *Guarantees) ByCollectionID(collID flow.Identifier) (*flow.CollectionGuarantee, error) {
    63  	tx := g.db.NewTransaction(false)
    64  	defer tx.Discard()
    65  	return g.retrieveTx(collID)(tx)
    66  }