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

     1  package badger
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"github.com/dgraph-io/badger/v2"
     7  
     8  	"github.com/onflow/flow-go/storage/badger/operation"
     9  )
    10  
    11  type ConsumerProgress struct {
    12  	db       *badger.DB
    13  	consumer string // to distinguish the consume progress between different consumers
    14  }
    15  
    16  func NewConsumerProgress(db *badger.DB, consumer string) *ConsumerProgress {
    17  	return &ConsumerProgress{
    18  		db:       db,
    19  		consumer: consumer,
    20  	}
    21  }
    22  
    23  func (cp *ConsumerProgress) ProcessedIndex() (uint64, error) {
    24  	var processed uint64
    25  	err := cp.db.View(operation.RetrieveProcessedIndex(cp.consumer, &processed))
    26  	if err != nil {
    27  		return 0, fmt.Errorf("failed to retrieve processed index: %w", err)
    28  	}
    29  	return processed, nil
    30  }
    31  
    32  // InitProcessedIndex insert the default processed index to the storage layer, can only be done once.
    33  // initialize for the second time will return storage.ErrAlreadyExists
    34  func (cp *ConsumerProgress) InitProcessedIndex(defaultIndex uint64) error {
    35  	err := operation.RetryOnConflict(cp.db.Update, operation.InsertProcessedIndex(cp.consumer, defaultIndex))
    36  	if err != nil {
    37  		return fmt.Errorf("could not update processed index: %w", err)
    38  	}
    39  
    40  	return nil
    41  }
    42  
    43  func (cp *ConsumerProgress) SetProcessedIndex(processed uint64) error {
    44  	err := operation.RetryOnConflict(cp.db.Update, operation.SetProcessedIndex(cp.consumer, processed))
    45  	if err != nil {
    46  		return fmt.Errorf("could not update processed index: %w", err)
    47  	}
    48  
    49  	return nil
    50  }