adding test cases

Author: zhangchiqing

module/finalizer/collection/finalizer_pebble_test.go

@@ -361,6 +361,9 @@ func TestFinalizerPebble(t *testing.T) {
 			})
 		})
 	})
+
+	t.Run("concurrency safety", func(t *testing.T) {
+	})
 }
 
 // assertClusterBlocksIndexedByReferenceHeightPebble checks the given cluster blocks have

module/finalizer/consensus/finalizer_pebble_test.go

@@ -1,7 +1,10 @@
 package consensus
 
 import (
+	"context"
+	"fmt"
 	"math/rand"
+	"sync"
 	"testing"
 
 	"github.com/cockroachdb/pebble"
@@ -13,6 +16,7 @@ import (
 	"github.com/onflow/flow-go/module/metrics"
 	"github.com/onflow/flow-go/module/trace"
 	mockprot "github.com/onflow/flow-go/state/protocol/mock"
+	protocolstorage "github.com/onflow/flow-go/storage"
 	mockstor "github.com/onflow/flow-go/storage/mock"
 	storage "github.com/onflow/flow-go/storage/pebble"
 	"github.com/onflow/flow-go/storage/pebble/operation"
@@ -90,11 +94,12 @@ func TestMakeFinalValidChainPebble(t *testing.T) {
 		// initialize the finalizer with the dependencies and make the call
 		metrics := metrics.NewNoopCollector()
 		fin := FinalizerPebble{
-			db:      db,
-			headers: storage.NewHeaders(metrics, db),
-			state:   state,
-			tracer:  trace.NewNoopTracer(),
-			cleanup: LogCleanup(&list),
+			db:         db,
+			headers:    storage.NewHeaders(metrics, db),
+			state:      state,
+			tracer:     trace.NewNoopTracer(),
+			cleanup:    LogCleanup(&list),
+			finalizing: new(sync.Mutex),
 		}
 		err = fin.MakeFinal(lastID)
 		require.NoError(t, err)
@@ -146,11 +151,12 @@ func TestMakeFinalInvalidHeightPebble(t *testing.T) {
 		// initialize the finalizer with the dependencies and make the call
 		metrics := metrics.NewNoopCollector()
 		fin := FinalizerPebble{
-			db:      db,
-			headers: storage.NewHeaders(metrics, db),
-			state:   state,
-			tracer:  trace.NewNoopTracer(),
-			cleanup: LogCleanup(&list),
+			db:         db,
+			headers:    storage.NewHeaders(metrics, db),
+			state:      state,
+			tracer:     trace.NewNoopTracer(),
+			cleanup:    LogCleanup(&list),
+			finalizing: new(sync.Mutex),
 		}
 		err = fin.MakeFinal(pending.ID())
 		require.Error(t, err)
@@ -194,11 +200,12 @@ func TestMakeFinalDuplicatePebble(t *testing.T) {
 		// initialize the finalizer with the dependencies and make the call
 		metrics := metrics.NewNoopCollector()
 		fin := FinalizerPebble{
-			db:      db,
-			headers: storage.NewHeaders(metrics, db),
-			state:   state,
-			tracer:  trace.NewNoopTracer(),
-			cleanup: LogCleanup(&list),
+			db:         db,
+			headers:    storage.NewHeaders(metrics, db),
+			state:      state,
+			tracer:     trace.NewNoopTracer(),
+			cleanup:    LogCleanup(&list),
+			finalizing: new(sync.Mutex),
 		}
 		err = fin.MakeFinal(final.ID())
 		require.NoError(t, err)
@@ -210,3 +217,92 @@ func TestMakeFinalDuplicatePebble(t *testing.T) {
 	// make sure no cleanup was done
 	assert.Empty(t, list)
 }
+
+// create a chain of 10 blocks, calling MakeFinal(1), MakeFinal(2), ..., MakeFinal(10) concurrently
+// expect 10 is finalized in the end
+func TestMakeFinalConcurrencySafe(t *testing.T) {
+	genesis := unittest.BlockHeaderFixture()
+	blocks := unittest.ChainFixtureFrom(10, genesis)
+
+	blockLookup := make(map[flow.Identifier]*flow.Block)
+	for _, block := range blocks {
+		blockLookup[block.Header.ID()] = block
+	}
+
+	var list []flow.Identifier
+
+	unittest.RunWithPebbleDB(t, func(db *pebble.DB) {
+		// create a mock protocol state to check finalize calls
+		state := mockprot.NewFollowerState(t)
+		state.On("Finalize", mock.Anything, mock.Anything).Return(
+			func(ctx context.Context, blockID flow.Identifier) error {
+				block, ok := blockLookup[blockID]
+				if !ok {
+					return fmt.Errorf("block %s not found", blockID)
+				}
+
+				header := block.Header
+
+				return operation.WithReaderBatchWriter(db, func(rw protocolstorage.PebbleReaderBatchWriter) error {
+					_, tx := rw.ReaderWriter()
+					err := operation.IndexBlockHeight(header.Height, header.ID())(tx)
+					if err != nil {
+						return err
+					}
+					return operation.UpdateFinalizedHeight(header.Height)(tx)
+				})
+			})
+
+		// insert the latest finalized height
+		err := operation.InsertFinalizedHeight(genesis.Height)(db)
+		require.NoError(t, err)
+
+		// map the finalized height to the finalized block ID
+		err = operation.IndexBlockHeight(genesis.Height, genesis.ID())(db)
+		require.NoError(t, err)
+
+		// insert the finalized block header into the DB
+		err = operation.InsertHeader(genesis.ID(), genesis)(db)
+		require.NoError(t, err)
+
+		// insert all of the pending blocks into the DB
+		for _, block := range blocks {
+			header := block.Header
+			err = operation.InsertHeader(header.ID(), header)(db)
+			require.NoError(t, err)
+		}
+
+		// initialize the finalizer with the dependencies and make the call
+		metrics := metrics.NewNoopCollector()
+		fin := FinalizerPebble{
+			db:         db,
+			headers:    storage.NewHeaders(metrics, db),
+			state:      state,
+			tracer:     trace.NewNoopTracer(),
+			cleanup:    LogCleanup(&list),
+			finalizing: new(sync.Mutex),
+		}
+
+		// Concurrently finalize blocks[0] to blocks[9]
+		var wg sync.WaitGroup
+		for _, block := range blocks {
+			wg.Add(1)
+			go func(block *flow.Block) {
+				defer wg.Done()
+				err := fin.MakeFinal(block.Header.ID())
+				require.NoError(t, err)
+			}(block)
+		}
+
+		// Wait for all finalization operations to complete
+		wg.Wait()
+
+		var finalized uint64
+		require.NoError(t, operation.RetrieveFinalizedHeight(&finalized)(db))
+
+		require.Equal(t, blocks[len(blocks)-1].Header.Height, finalized)
+
+		// make sure that nothing was finalized
+		state.AssertExpectations(t)
+	})
+}