Merge pull request #7939 from onflow/leo/refactor-insert-chunk-data-pack

[Storage] Refactor insert chunk data pack

Author: zhangchiqing

engine/execution/pruner/core_test.go

@@ -75,7 +75,9 @@ func TestLoopPruneExecutionDataFromRootToLatestSealed(t *testing.T) {
 			require.NoError(t, err)
 			require.NoError(t, results.Store(chunk.Result))
 			require.NoError(t, results.Index(chunk.Result.BlockID, chunk.Result.ID()))
-			require.NoError(t, chunkDataPacks.Store([]*flow.ChunkDataPack{chunk.ChunkDataPack}))
+			require.NoError(t, unittest.WithLock(t, lockManager, storage.LockInsertChunkDataPack, func(lctx lockctx.Context) error {
+				return chunkDataPacks.StoreByChunkID(lctx, []*flow.ChunkDataPack{chunk.ChunkDataPack})
+			}))
 			_, storeErr := collections.Store(chunk.ChunkDataPack.Collection)
 			require.NoError(t, storeErr)
 			// verify that chunk data pack fixture can be found by the result

engine/execution/state/state.go

@@ -397,6 +397,8 @@ func (s *state) SaveExecutionResults(
 	return nil
 }
 
+// saveExecutionResults saves all data related to the execution of a block.
+// It is concurrent-safe
 func (s *state) saveExecutionResults(
 	ctx context.Context,
 	result *execution.ComputationResult,
@@ -408,73 +410,72 @@ func (s *state) saveExecutionResults(
 		return fmt.Errorf("can not retrieve chunk data packs: %w", err)
 	}
 
-	err = s.chunkDataPacks.Store(chunks)
-	if err != nil {
-		return fmt.Errorf("can not store multiple chunk data pack: %w", err)
-	}
+	// Acquire both locks to ensure it's concurrent safe when inserting the execution results and chunk data packs.
+	return storage.WithLocks(s.lockManager, []string{storage.LockInsertOwnReceipt, storage.LockInsertChunkDataPack}, func(lctx lockctx.Context) error {
+		err := s.chunkDataPacks.StoreByChunkID(lctx, chunks)
+		if err != nil {
+			return fmt.Errorf("can not store multiple chunk data pack: %w", err)
+		}
 
-	lctx := s.lockManager.NewContext()
-	defer lctx.Release()
-	err = lctx.AcquireLock(storage.LockInsertOwnReceipt)
-	if err != nil {
-		return err
-	}
+		// Save entire execution result (including all chunk data packs) within one batch to minimize
+		// the number of database interactions.
+		return s.db.WithReaderBatchWriter(func(batch storage.ReaderBatchWriter) error {
+			batch.AddCallback(func(err error) {
+				// Rollback if an error occurs during batch operations
+				// Chunk data packs are saved in a separate database, there is a chance
+				// that execution result was failed to save, but chunk data packs was saved and
+				// didnt get removed.
+				// TODO(leo): when retrieving chunk data packs, we need to add a check to ensure the block
+				// has been executed before returning chunk data packs
+				if err != nil {
+					chunkIDs := make([]flow.Identifier, 0, len(chunks))
+					for _, chunk := range chunks {
+						chunkIDs = append(chunkIDs, chunk.ChunkID)
+					}
+					_ = s.chunkDataPacks.Remove(chunkIDs)
+				}
+			})
 
-	// Save entire execution result (including all chunk data packs) within one batch to minimize
-	// the number of database interactions. This is a large batch of data, which might not be
-	// committed within a single operation (e.g. if using Badger DB as storage backend, which has
-	// a size limit for its transactions).
-	return s.db.WithReaderBatchWriter(func(batch storage.ReaderBatchWriter) error {
-		batch.AddCallback(func(err error) {
-			// Rollback if an error occurs during batch operations
+			err = s.events.BatchStore(blockID, []flow.EventsList{result.AllEvents()}, batch)
 			if err != nil {
-				chunkIDs := make([]flow.Identifier, 0, len(chunks))
-				for _, chunk := range chunks {
-					chunkIDs = append(chunkIDs, chunk.ChunkID)
-				}
-				_ = s.chunkDataPacks.Remove(chunkIDs)
+				return fmt.Errorf("cannot store events: %w", err)
 			}
-		})
 
-		err = s.events.BatchStore(blockID, []flow.EventsList{result.AllEvents()}, batch)
-		if err != nil {
-			return fmt.Errorf("cannot store events: %w", err)
-		}
-
-		err = s.serviceEvents.BatchStore(blockID, result.AllServiceEvents(), batch)
-		if err != nil {
-			return fmt.Errorf("cannot store service events: %w", err)
-		}
+			err = s.serviceEvents.BatchStore(blockID, result.AllServiceEvents(), batch)
+			if err != nil {
+				return fmt.Errorf("cannot store service events: %w", err)
+			}
 
-		err = s.transactionResults.BatchStore(
-			blockID,
-			result.AllTransactionResults(),
-			batch)
-		if err != nil {
-			return fmt.Errorf("cannot store transaction result: %w", err)
-		}
+			err = s.transactionResults.BatchStore(
+				blockID,
+				result.AllTransactionResults(),
+				batch)
+			if err != nil {
+				return fmt.Errorf("cannot store transaction result: %w", err)
+			}
 
-		executionResult := &result.ExecutionReceipt.ExecutionResult
-		// saving my receipts will also save the execution result
-		err = s.myReceipts.BatchStoreMyReceipt(lctx, result.ExecutionReceipt, batch)
-		if err != nil {
-			return fmt.Errorf("could not persist execution result: %w", err)
-		}
+			executionResult := &result.ExecutionReceipt.ExecutionResult
+			// saving my receipts will also save the execution result
+			err = s.myReceipts.BatchStoreMyReceipt(lctx, result.ExecutionReceipt, batch)
+			if err != nil {
+				return fmt.Errorf("could not persist execution result: %w", err)
+			}
 
-		err = s.results.BatchIndex(blockID, executionResult.ID(), batch)
-		if err != nil {
-			return fmt.Errorf("cannot index execution result: %w", err)
-		}
+			err = s.results.BatchIndex(blockID, executionResult.ID(), batch)
+			if err != nil {
+				return fmt.Errorf("cannot index execution result: %w", err)
+			}
 
-		// the state commitment is the last data item to be stored, so that
-		// IsBlockExecuted can be implemented by checking whether state commitment exists
-		// in the database
-		err = s.commits.BatchStore(lctx, blockID, result.CurrentEndState(), batch)
-		if err != nil {
-			return fmt.Errorf("cannot store state commitment: %w", err)
-		}
+			// the state commitment is the last data item to be stored, so that
+			// IsBlockExecuted can be implemented by checking whether state commitment exists
+			// in the database
+			err = s.commits.BatchStore(lctx, blockID, result.CurrentEndState(), batch)
+			if err != nil {
+				return fmt.Errorf("cannot store state commitment: %w", err)
+			}
 
-		return nil
+			return nil
+		})
 	})
 }
 

engine/execution/state/state_storehouse_test.go

@@ -61,7 +61,7 @@ func prepareStorehouseTest(f func(t *testing.T, es state.ExecutionState, l *ledg
 			txResults := storagemock.NewTransactionResults(t)
 			txResults.On("BatchStore", mock.Anything, mock.Anything, mock.Anything).Return(nil)
 			chunkDataPacks := storagemock.NewChunkDataPacks(t)
-			chunkDataPacks.On("Store", mock.Anything).Return(nil)
+			chunkDataPacks.On("StoreByChunkID", mock.Anything, mock.Anything).Return(nil)
 			results := storagemock.NewExecutionResults(t)
 			results.On("BatchIndex", mock.Anything, mock.Anything, mock.Anything).Return(nil)
 			myReceipts := storagemock.NewMyExecutionReceipts(t)

model/flow/chunk.go

@@ -362,3 +362,23 @@ func stringsToCids(strs []string) ([]cid.Cid, error) {
 	}
 	return cids, nil
 }
+
+// Equals returns true if and only if receiver BlockExecutionDataRoot is equal to the `other`.
+func (b BlockExecutionDataRoot) Equals(other BlockExecutionDataRoot) bool {
+	// Compare BlockID fields
+	if b.BlockID != other.BlockID {
+		return false
+	}
+
+	// Compare ChunkExecutionDataIDs slices
+	if len(b.ChunkExecutionDataIDs) != len(other.ChunkExecutionDataIDs) {
+		return false
+	}
+	for i, cid := range b.ChunkExecutionDataIDs {
+		if !cid.Equals(other.ChunkExecutionDataIDs[i]) {
+			return false
+		}
+	}
+
+	return true
+}

module/block_iterator/executor/executor_test.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/cockroachdb/pebble/v2"
+	"github.com/jordanschalm/lockctx"
 	"github.com/stretchr/testify/require"
 
 	"github.com/onflow/flow-go/model/flow"
@@ -23,6 +24,7 @@ import (
 // verify the executor is able to iterate through all blocks from the iterator.
 func TestExecute(t *testing.T) {
 	unittest.RunWithPebbleDB(t, func(db *pebble.DB) {
+		lockManager := storage.NewTestingLockManager()
 		blockCount := 10
 
 		// prepare data
@@ -39,8 +41,10 @@ func TestExecute(t *testing.T) {
 		// store the chunk data packs to be pruned later
 		for _, cdp := range cdps {
 			sc := storage.ToStoredChunkDataPack(cdp)
-			require.NoError(t, pdb.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-				return operation.InsertChunkDataPack(rw.Writer(), sc)
+			require.NoError(t, unittest.WithLock(t, lockManager, storage.LockInsertChunkDataPack, func(lctx lockctx.Context) error {
+				return pdb.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
+					return operation.InsertChunkDataPack(lctx, rw, sc)
+				})
 			}))
 		}
 
@@ -72,6 +76,7 @@ func TestExecute(t *testing.T) {
 // verify the pruning can be interrupted and resumed
 func TestExecuteCanBeResumed(t *testing.T) {
 	unittest.RunWithPebbleDB(t, func(db *pebble.DB) {
+		lockManager := storage.NewTestingLockManager()
 		blockCount := 10
 
 		cdps := make([]*flow.ChunkDataPack, 0, blockCount)
@@ -87,8 +92,10 @@ func TestExecuteCanBeResumed(t *testing.T) {
 		// store the chunk data packs to be pruned later
 		for _, cdp := range cdps {
 			sc := storage.ToStoredChunkDataPack(cdp)
-			require.NoError(t, pdb.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-				return operation.InsertChunkDataPack(rw.Writer(), sc)
+			require.NoError(t, unittest.WithLock(t, lockManager, storage.LockInsertChunkDataPack, func(lctx lockctx.Context) error {
+				return pdb.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
+					return operation.InsertChunkDataPack(lctx, rw, sc)
+				})
 			}))
 		}
 

module/pruner/pruners/chunk_data_pack_test.go

@@ -6,6 +6,7 @@ import (
 	"testing"
 
 	"github.com/cockroachdb/pebble/v2"
+	"github.com/jordanschalm/lockctx"
 	"github.com/stretchr/testify/require"
 
 	"github.com/onflow/flow-go/module/metrics"
@@ -18,6 +19,7 @@ import (
 func TestChunkDataPackPruner(t *testing.T) {
 
 	unittest.RunWithPebbleDB(t, func(pebbleDB *pebble.DB) {
+		lockManager := storage.NewTestingLockManager()
 		m := metrics.NewNoopCollector()
 		db := pebbleimpl.ToDB(pebbleDB)
 		results := store.NewExecutionResults(m, db)
@@ -29,7 +31,9 @@ func TestChunkDataPackPruner(t *testing.T) {
 		// store the chunks
 		cdp1, result1 := unittest.ChunkDataPacksFixtureAndResult()
 		require.NoError(t, results.Store(result1))
-		require.NoError(t, chunks.Store(cdp1))
+		require.NoError(t, unittest.WithLock(t, lockManager, storage.LockInsertChunkDataPack, func(lctx lockctx.Context) error {
+			return chunks.StoreByChunkID(lctx, cdp1)
+		}))
 
 		pruner := NewChunkDataPackPruner(chunks, results)
 

storage/chunk_data_packs.go

@@ -1,6 +1,10 @@
 package storage
 
 import (
+	"bytes"
+
+	"github.com/jordanschalm/lockctx"
+
 	"github.com/onflow/flow-go/model/flow"
 )
 
@@ -9,7 +13,7 @@ type ChunkDataPacks interface {
 
 	// Store stores multiple ChunkDataPacks cs keyed by their ChunkIDs in a batch.
 	// No errors are expected during normal operation, but it may return generic error
-	Store(cs []*flow.ChunkDataPack) error
+	StoreByChunkID(lctx lockctx.Proof, cs []*flow.ChunkDataPack) error
 
 	// Remove removes multiple ChunkDataPacks cs keyed by their ChunkIDs in a batch.
 	// No errors are expected during normal operation, but it may return generic error
@@ -54,3 +58,25 @@ func ToStoredChunkDataPack(c *flow.ChunkDataPack) *StoredChunkDataPack {
 
 	return sc
 }
+
+func (c StoredChunkDataPack) Equals(other StoredChunkDataPack) error {
+	if c.ChunkID != other.ChunkID {
+		return ErrDataMismatch
+	}
+	if c.StartState != other.StartState {
+		return ErrDataMismatch
+	}
+	if !c.ExecutionDataRoot.Equals(other.ExecutionDataRoot) {
+		return ErrDataMismatch
+	}
+	if c.SystemChunk != other.SystemChunk {
+		return ErrDataMismatch
+	}
+	if !bytes.Equal(c.Proof, other.Proof) {
+		return ErrDataMismatch
+	}
+	if c.CollectionID != other.CollectionID {
+		return ErrDataMismatch
+	}
+	return nil
+}

storage/locks.go

@@ -26,6 +26,8 @@ const (
 	LockInsertCollection = "lock_insert_collection"
 	// LockBootstrapping protects data that is *exclusively* written during bootstrapping.
 	LockBootstrapping = "lock_bootstrapping"
+	// LockInsertChunkDataPack protects the insertion of chunk data packs (not yet used anywhere
+	LockInsertChunkDataPack = "lock_insert_chunk_data_pack"
 )
 
 // Locks returns a list of all named locks used by the storage layer.
@@ -38,6 +40,7 @@ func Locks() []string {
 		LockInsertOwnReceipt,
 		LockInsertCollection,
 		LockBootstrapping,
+		LockInsertChunkDataPack,
 	}
 }
 
@@ -61,6 +64,7 @@ func makeLockPolicy() lockctx.Policy {
 	return lockctx.NewDAGPolicyBuilder().
 		Add(LockInsertBlock, LockFinalizeBlock).
 		Add(LockFinalizeBlock, LockBootstrapping).
+		Add(LockInsertOwnReceipt, LockInsertChunkDataPack).
 		Build()
 }
 
@@ -110,3 +114,24 @@ func MakeSingletonLockManager() lockctx.Manager {
 func NewTestingLockManager() lockctx.Manager {
 	return lockctx.NewManager(Locks(), makeLockPolicy())
 }
+
+// WithLock is a helper function that creates a new lock context, acquires the specified lock,
+// and executes the provided function within that context.
+// This function passes through any errors returned by fn.
+func WithLock(manager lockctx.Manager, lockID string, fn func(lctx lockctx.Context) error) error {
+	return WithLocks(manager, []string{lockID}, fn)
+}
+
+// WithLocks is a helper function that creates a new lock context, acquires the specified locks,
+// and executes the provided function within that context.
+// This function passes through any errors returned by fn.
+func WithLocks(manager lockctx.Manager, lockIDs []string, fn func(lctx lockctx.Context) error) error {
+	lctx := manager.NewContext()
+	defer lctx.Release()
+	for _, lockID := range lockIDs {
+		if err := lctx.AcquireLock(lockID); err != nil {
+			return err
+		}
+	}
+	return fn(lctx)
+}