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

     1  package badger
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  
     7  	"github.com/dgraph-io/badger/v2"
     8  
     9  	"github.com/onflow/flow-go/model/chunks"
    10  	"github.com/onflow/flow-go/model/flow"
    11  	"github.com/onflow/flow-go/storage"
    12  	"github.com/onflow/flow-go/storage/badger/operation"
    13  )
    14  
    15  // ChunksQueue stores a queue of chunk locators that assigned to me to verify.
    16  // Job consumers can read the locators as job from the queue by index.
    17  // Chunk locators stored in this queue are unique.
    18  type ChunksQueue struct {
    19  	db *badger.DB
    20  }
    21  
    22  const JobQueueChunksQueue = "JobQueueChunksQueue"
    23  
    24  // NewChunkQueue will initialize the underlying badger database of chunk locator queue.
    25  func NewChunkQueue(db *badger.DB) *ChunksQueue {
    26  	return &ChunksQueue{
    27  		db: db,
    28  	}
    29  }
    30  
    31  // Init initializes chunk queue's latest index with the given default index.
    32  func (q *ChunksQueue) Init(defaultIndex uint64) (bool, error) {
    33  	_, err := q.LatestIndex()
    34  	if errors.Is(err, storage.ErrNotFound) {
    35  		err = q.db.Update(operation.InitJobLatestIndex(JobQueueChunksQueue, defaultIndex))
    36  		if err != nil {
    37  			return false, fmt.Errorf("could not init chunk locator queue with default index %v: %w", defaultIndex, err)
    38  		}
    39  		return true, nil
    40  	}
    41  	if err != nil {
    42  		return false, fmt.Errorf("could not get latest index: %w", err)
    43  	}
    44  
    45  	return false, nil
    46  }
    47  
    48  // StoreChunkLocator stores a new chunk locator that assigned to me to the job queue.
    49  // A true will be returned, if the locator was new.
    50  // A false will be returned, if the locator was duplicate.
    51  func (q *ChunksQueue) StoreChunkLocator(locator *chunks.Locator) (bool, error) {
    52  	err := operation.RetryOnConflict(q.db.Update, func(tx *badger.Txn) error {
    53  		// make sure the chunk locator is unique
    54  		err := operation.InsertChunkLocator(locator)(tx)
    55  		if err != nil {
    56  			return fmt.Errorf("failed to insert chunk locator: %w", err)
    57  		}
    58  
    59  		// read the latest index
    60  		var latest uint64
    61  		err = operation.RetrieveJobLatestIndex(JobQueueChunksQueue, &latest)(tx)
    62  		if err != nil {
    63  			return fmt.Errorf("failed to retrieve job index for chunk locator queue: %w", err)
    64  		}
    65  
    66  		// insert to the next index
    67  		next := latest + 1
    68  		err = operation.InsertJobAtIndex(JobQueueChunksQueue, next, locator.ID())(tx)
    69  		if err != nil {
    70  			return fmt.Errorf("failed to set job index for chunk locator queue at index %v: %w", next, err)
    71  		}
    72  
    73  		// update the next index as the latest index
    74  		err = operation.SetJobLatestIndex(JobQueueChunksQueue, next)(tx)
    75  		if err != nil {
    76  			return fmt.Errorf("failed to update latest index %v: %w", next, err)
    77  		}
    78  
    79  		return nil
    80  	})
    81  
    82  	// was trying to store a duplicate locator
    83  	if errors.Is(err, storage.ErrAlreadyExists) {
    84  		return false, nil
    85  	}
    86  	if err != nil {
    87  		return false, fmt.Errorf("failed to store chunk locator: %w", err)
    88  	}
    89  	return true, nil
    90  }
    91  
    92  // LatestIndex returns the index of the latest chunk locator stored in the queue.
    93  func (q *ChunksQueue) LatestIndex() (uint64, error) {
    94  	var latest uint64
    95  	err := q.db.View(operation.RetrieveJobLatestIndex(JobQueueChunksQueue, &latest))
    96  	if err != nil {
    97  		return 0, fmt.Errorf("could not retrieve latest index for chunks queue: %w", err)
    98  	}
    99  	return latest, nil
   100  }
   101  
   102  // AtIndex returns the chunk locator stored at the given index in the queue.
   103  func (q *ChunksQueue) AtIndex(index uint64) (*chunks.Locator, error) {
   104  	var locatorID flow.Identifier
   105  	err := q.db.View(operation.RetrieveJobAtIndex(JobQueueChunksQueue, index, &locatorID))
   106  	if err != nil {
   107  		return nil, fmt.Errorf("could not retrieve chunk locator in queue: %w", err)
   108  	}
   109  
   110  	var locator chunks.Locator
   111  	err = q.db.View(operation.RetrieveChunkLocator(locatorID, &locator))
   112  	if err != nil {
   113  		return nil, fmt.Errorf("could not retrieve locator for chunk id %v: %w", locatorID, err)
   114  	}
   115  
   116  	return &locator, nil
   117  }