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

     1  package badger
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"github.com/dgraph-io/badger/v3"
     7  
     8  	"github.com/koko1123/flow-go-1/model/flow"
     9  	"github.com/koko1123/flow-go-1/module"
    10  	"github.com/koko1123/flow-go-1/module/metrics"
    11  	"github.com/koko1123/flow-go-1/storage"
    12  	badgermodel "github.com/koko1123/flow-go-1/storage/badger/model"
    13  	"github.com/koko1123/flow-go-1/storage/badger/operation"
    14  	"github.com/koko1123/flow-go-1/storage/badger/transaction"
    15  )
    16  
    17  type ChunkDataPacks struct {
    18  	db             *badger.DB
    19  	collections    storage.Collections
    20  	byChunkIDCache *Cache
    21  }
    22  
    23  func NewChunkDataPacks(collector module.CacheMetrics, db *badger.DB, collections storage.Collections, byChunkIDCacheSize uint) *ChunkDataPacks {
    24  
    25  	store := func(key interface{}, val interface{}) func(*transaction.Tx) error {
    26  		chdp := val.(*badgermodel.StoredChunkDataPack)
    27  		return transaction.WithTx(operation.SkipDuplicates(operation.InsertChunkDataPack(chdp)))
    28  	}
    29  
    30  	retrieve := func(key interface{}) func(tx *badger.Txn) (interface{}, error) {
    31  		chunkID := key.(flow.Identifier)
    32  
    33  		var c badgermodel.StoredChunkDataPack
    34  		return func(tx *badger.Txn) (interface{}, error) {
    35  			err := operation.RetrieveChunkDataPack(chunkID, &c)(tx)
    36  			return &c, err
    37  		}
    38  	}
    39  
    40  	cache := newCache(collector, metrics.ResourceChunkDataPack,
    41  		withLimit(byChunkIDCacheSize),
    42  		withStore(store),
    43  		withRetrieve(retrieve),
    44  	)
    45  
    46  	ch := ChunkDataPacks{
    47  		db:             db,
    48  		byChunkIDCache: cache,
    49  		collections:    collections,
    50  	}
    51  	return &ch
    52  }
    53  
    54  func (ch *ChunkDataPacks) Store(c *flow.ChunkDataPack) error {
    55  	sc := toStoredChunkDataPack(c)
    56  	err := operation.RetryOnConflictTx(ch.db, transaction.Update, ch.byChunkIDCache.PutTx(sc.ChunkID, sc))
    57  	if err != nil {
    58  		return fmt.Errorf("could not store chunk datapack: %w", err)
    59  	}
    60  	return nil
    61  }
    62  
    63  func (ch *ChunkDataPacks) Remove(chunkID flow.Identifier) error {
    64  	err := operation.RetryOnConflict(ch.db.Update, operation.RemoveChunkDataPack(chunkID))
    65  	if err != nil {
    66  		return fmt.Errorf("could not remove chunk datapack: %w", err)
    67  	}
    68  	// TODO Integrate cache removal in a similar way as storage/retrieval is
    69  	ch.byChunkIDCache.Remove(chunkID)
    70  	return nil
    71  }
    72  
    73  // BatchStore stores ChunkDataPack c keyed by its ChunkID in provided batch.
    74  // No errors are expected during normal operation, but it may return generic error
    75  // if entity is not serializable or Badger unexpectedly fails to process request
    76  func (ch *ChunkDataPacks) BatchStore(c *flow.ChunkDataPack, batch storage.BatchStorage) error {
    77  	sc := toStoredChunkDataPack(c)
    78  	writeBatch := batch.GetWriter()
    79  	batch.OnSucceed(func() {
    80  		ch.byChunkIDCache.Insert(sc.ChunkID, sc)
    81  	})
    82  	return operation.BatchInsertChunkDataPack(sc)(writeBatch)
    83  }
    84  
    85  // BatchRemove removes ChunkDataPack c keyed by its ChunkID in provided batch
    86  // No errors are expected during normal operation, even if no entries are matched.
    87  // If Badger unexpectedly fails to process the request, the error is wrapped in a generic error and returned.
    88  func (ch *ChunkDataPacks) BatchRemove(chunkID flow.Identifier, batch storage.BatchStorage) error {
    89  	writeBatch := batch.GetWriter()
    90  	batch.OnSucceed(func() {
    91  		ch.byChunkIDCache.Remove(chunkID)
    92  	})
    93  	return operation.BatchRemoveChunkDataPack(chunkID)(writeBatch)
    94  }
    95  
    96  func (ch *ChunkDataPacks) ByChunkID(chunkID flow.Identifier) (*flow.ChunkDataPack, error) {
    97  	schdp, err := ch.byChunkID(chunkID)
    98  	if err != nil {
    99  		return nil, err
   100  	}
   101  
   102  	chdp := &flow.ChunkDataPack{
   103  		ChunkID:    schdp.ChunkID,
   104  		StartState: schdp.StartState,
   105  		Proof:      schdp.Proof,
   106  	}
   107  
   108  	if !schdp.SystemChunk {
   109  		collection, err := ch.collections.ByID(schdp.CollectionID)
   110  		if err != nil {
   111  			return nil, fmt.Errorf("could not retrive collection (id: %x) for stored chunk data pack: %w", schdp.CollectionID, err)
   112  		}
   113  
   114  		chdp.Collection = collection
   115  	}
   116  
   117  	return chdp, nil
   118  }
   119  
   120  func (ch *ChunkDataPacks) byChunkID(chunkID flow.Identifier) (*badgermodel.StoredChunkDataPack, error) {
   121  	tx := ch.db.NewTransaction(false)
   122  	defer tx.Discard()
   123  
   124  	schdp, err := ch.retrieveCHDP(chunkID)(tx)
   125  	if err != nil {
   126  		return nil, fmt.Errorf("could not retrive stored chunk data pack: %w", err)
   127  	}
   128  
   129  	return schdp, nil
   130  }
   131  
   132  func (ch *ChunkDataPacks) retrieveCHDP(chunkID flow.Identifier) func(*badger.Txn) (*badgermodel.StoredChunkDataPack, error) {
   133  	return func(tx *badger.Txn) (*badgermodel.StoredChunkDataPack, error) {
   134  		val, err := ch.byChunkIDCache.Get(chunkID)(tx)
   135  		if err != nil {
   136  			return nil, err
   137  		}
   138  		return val.(*badgermodel.StoredChunkDataPack), nil
   139  	}
   140  }
   141  
   142  func toStoredChunkDataPack(c *flow.ChunkDataPack) *badgermodel.StoredChunkDataPack {
   143  	sc := &badgermodel.StoredChunkDataPack{
   144  		ChunkID:     c.ChunkID,
   145  		StartState:  c.StartState,
   146  		Proof:       c.Proof,
   147  		SystemChunk: false,
   148  	}
   149  
   150  	if c.Collection != nil {
   151  		// non system chunk
   152  		sc.CollectionID = c.Collection.ID()
   153  	} else {
   154  		sc.SystemChunk = true
   155  	}
   156  
   157  	return sc
   158  }