[Testing] Fix races and bugs found by flaky tests

module/executiondatasync/optimistic_sync/pipeline/pipeline.go

@@ -129,7 +129,12 @@ func (p *Pipeline) Run(ctx context.Context, core optimistic_sync.Core, parentSta
 		return irrecoverable.NewExceptionf("pipeline has been already started, it is not designed to be run again")
 	}
 	p.core = core
-	p.parentStateCache.Store(int32(parentState))
+	// Initialize the parent state cache with the provided initial state.
+	// CompareAndSwap ensures that a state update delivered concurrently via OnParentStateUpdated
+	// (which is always at least as current as the caller-provided initial state) is not
+	// overwritten and lost. Otherwise, the pipeline could deadlock, e.g. never observing the
+	// parent reaching StateComplete, and hence never persisting.
+	p.parentStateCache.CompareAndSwap(int32(optimistic_sync.StatePending), int32(parentState))
 	// run the main event loop by calling p.loop. any error returned from it needs to be propagated to the caller.
 	// IMPORTANT: after the main loop has exited we need to ensure that worker goroutine has also finished
 	// because we need to ensure that it can report any error that has happened during the execution of detached operation.
@@ -279,12 +284,13 @@ func (p *Pipeline) SetSealed() {
 }
 
 // OnParentStateUpdated updates the pipeline's state based on the provided parent state.
-// If the parent state has changed, it will notify the state consumer and trigger a state change notification.
+// It will notify the state consumer and trigger a state change notification.
 func (p *Pipeline) OnParentStateUpdated(parentState optimistic_sync.State) {
-	oldState := p.parentStateCache.Load()
-	if p.parentStateCache.CompareAndSwap(oldState, int32(parentState)) {
-		p.stateChangedNotifier.Notify()
-	}
+	// Note: an unconditional store is used, so that an update can never be silently dropped
+	// (a Load-CompareAndSwap sequence could fail if it races with the initialization in Run,
+	// losing the update). Spurious notifications are cheap: the notifier merges them.
+	p.parentStateCache.Store(int32(parentState))
+	p.stateChangedNotifier.Notify()
 }
 
 // Abandon marks the pipeline as abandoned

module/executiondatasync/optimistic_sync/pipeline/pipeline_functional_test.go

@@ -2,6 +2,7 @@ package pipeline
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"os"
 	"testing"
@@ -239,10 +240,14 @@ func (p *PipelineFunctionalSuite) TestPipelineCompletesSuccessfully() {
 	p.txResultErrMsgsRequester.On("Request", mock.Anything).Return(p.expectedTxResultErrMsgs, nil).Once()
 
 	p.WithRunningPipeline(func(pipeline optimistic_sync.Pipeline, updateChan chan optimistic_sync.State, errChan chan error, cancel context.CancelFunc) {
-		// Check for errors in a separate goroutine
+		// Check for errors in a separate goroutine.
+		// Note: WithRunningPipeline cancels the pipeline context when the test function returns,
+		// which can race with Run returning after the pipeline has already completed all of its
+		// work. A context.Canceled error is therefore an expected teardown artifact and must not
+		// fail the test.
 		go func() {
 			err := <-errChan
-			if err != nil {
+			if err != nil && !errors.Is(err, context.Canceled) {
 				p.T().Errorf("Pipeline error: %v", err)
 			}
 		}()

module/executiondatasync/optimistic_sync/pipeline/pipeline_test.go

@@ -68,12 +68,14 @@ func TestPipelineParentDependentTransitions(t *testing.T) {
 		assertNoUpdate(t, pipeline, updateChan, optimistic_sync.StatePending)
 
 		// 2. Update parent to downloading
-		parent.UpdateState(optimistic_sync.StateProcessing, pipeline)
-
 		// Pipeline should now call Download and Index within the processing state, then progress to
-		// WaitingPersist and stop
+		// WaitingPersist and stop.
+		// Note: the mocked calls must be registered BEFORE the parent state update: the pipeline's
+		// event loop may invoke them immediately after the update is delivered, racing the mock
+		// registration (mocks panic on unexpected calls).
 		mockCore.On("Download", mock.Anything).Return(nil)
 		mockCore.On("Index").Return(nil)
+		parent.UpdateState(optimistic_sync.StateProcessing, pipeline)
 		for _, expected := range []optimistic_sync.State{optimistic_sync.StateProcessing, optimistic_sync.StateWaitingPersist} {
 			synctest.Wait()
 			assertUpdate(t, updateChan, expected)
@@ -92,8 +94,9 @@ func TestPipelineParentDependentTransitions(t *testing.T) {
 		assertNoUpdate(t, pipeline, updateChan, optimistic_sync.StateWaitingPersist)
 
 		// 4. Mark the execution result as sealed, this should allow the pipeline to progress to Complete state
-		pipeline.SetSealed()
+		// (the mocked Persist must be registered BEFORE SetSealed, see the note above)
 		mockCore.On("Persist").Return(nil)
+		pipeline.SetSealed()
 
 		// Wait for pipeline to complete
 		synctest.Wait()

module/executiondatasync/optimistic_sync/pipeline/pipeline_test_utils.go

@@ -62,35 +62,28 @@ func (m *mockStateConsumer) OnStateUpdated(state optimistic_sync.State) {
 	m.updateChan <- state
 }
 
-// waitForStateUpdates waits for a sequence of state updates to occur or timeout after 500ms.
+// waitForStateUpdates waits for a sequence of state updates to occur or timeout after 5s.
 // updates must be received in the correct order or the test will fail.
-func waitForStateUpdates(t *testing.T, updateChan <-chan optimistic_sync.State, errChan <-chan error, expectedStates ...optimistic_sync.State) {
+//
+// Note: this function deliberately does NOT consume from errChan. The pipeline always queues a
+// state update (buffered updateChan) before offering an error (unbuffered errChan), but a select
+// listening on both channels picks randomly among ready cases: it could consume a (potentially
+// expected) error before a pending state update, failing the test spuriously and stealing the
+// error from a subsequent waitForError call. If the pipeline errors instead of emitting the
+// expected state updates, this function fails via its timeout.
+func waitForStateUpdates(t *testing.T, updateChan <-chan optimistic_sync.State, _ <-chan error, expectedStates ...optimistic_sync.State) {
 	done := make(chan struct{})
 	unittest.RequireReturnsBefore(t, func() {
 		for _, expected := range expectedStates {
-			// Prefer consuming pending state updates over errors: the pipeline always reports a state
-			// update before emitting an error, but both may already be queued by the time we read
-			// them (e.g. in tests expecting a state transition immediately followed by an expected
-			// error). A single select would pick one of the ready channels at random, potentially
-			// failing on the error before consuming the preceding state update.
-			select {
-			case update := <-updateChan:
-				assert.Equalf(t, expected, update, "expected pipeline to transition to %s, but got %s", expected, update)
-				continue
-			default:
-			}
-
 			select {
 			case <-done:
 				return
-			case err := <-errChan:
-				require.NoError(t, err, "pipeline returned error")
 			case update := <-updateChan:
 				assert.Equalf(t, expected, update, "expected pipeline to transition to %s, but got %s", expected, update)
 			}
 		}
-	}, 500*time.Millisecond, "Timeout waiting for state update")
-	close(done) // make sure function exists after timeout
+	}, 5*time.Second, "Timeout waiting for state update")
+	close(done) // make sure function exits after timeout
 }
 
 // waitForStateUpdatesAndNoError behaves like waitForStateUpdates, but additionally requires that
@@ -120,7 +113,7 @@ func waitForErrorWithCustomCheckers(t *testing.T, errChan <-chan error, errorChe
 				checker(err)
 			}
 		}
-	}, 500*time.Millisecond, "Timeout waiting for error")
+	}, 5*time.Second, "Timeout waiting for error")
 }
 
 // waitForError waits for an error from the errChan within 500ms and asserts it matches the expected error.
@@ -132,7 +125,7 @@ func waitForError(t *testing.T, errChan <-chan error, expectedErr error) {
 		} else {
 			assert.ErrorIs(t, err, expectedErr)
 		}
-	}, 500*time.Millisecond, "Timeout waiting for error")
+	}, 5*time.Second, "Timeout waiting for error")
 }
 
 // createPipeline initializes and returns a pipeline instance with its mock dependencies.

module/grpcserver/server.go

@@ -1,6 +1,7 @@
 package grpcserver
 
 import (
+	"errors"
 	"net"
 	"sync"
 
@@ -89,7 +90,12 @@ func (g *GrpcServer) serveGRPCWorker(ctx irrecoverable.SignalerContext, ready co
 	ready()
 
 	err = g.server.Serve(l) // blocking call
-	if err != nil {
+	if err != nil && !errors.Is(err, grpc.ErrServerStopped) {
+		// Serve returns nil when the server is stopped via Stop or GracefulStop while serving.
+		// ErrServerStopped is only returned when the server was stopped BEFORE Serve was called,
+		// which happens when the component is shut down in the short window between ready() above
+		// and the Serve call: the shutdownWorker may complete GracefulStop first. This is a normal
+		// shutdown, not an exception, hence it must not be thrown as an irrecoverable error.
 		g.log.Err(err).Msg("fatal error in grpc server")
 		ctx.Throw(err)
 	}

module/mempool/herocache/backdata/cache.go

@@ -367,9 +367,12 @@ func (c *Cache[V]) get(key flow.Identifier) (value V, bckIndex bucketIndex, sltI
 
 		id, linkedValue, linked := c.linkedValueOf(b, s)
 		if !linked {
-			// no linked entity for this (bucketIndex, slotIndex) pair.
-			c.collector.OnKeyGetFailure()
-			return value, 0, 0, false
+			// The slot's value is no longer in the underlying entities pool (it was ejected), but
+			// the slot still carries the ejected value's 32-bit id prefix, which happened to match
+			// the queried key's prefix. This does NOT imply the queried key is absent: a different
+			// slot in this bucket may hold the queried key (32-bit prefixes of distinct keys can
+			// collide), so we must continue scanning the remaining slots.
+			continue
 		}
 
 		if id != key {

module/mempool/herocache/backdata/cache_test.go

@@ -357,6 +357,55 @@ func TestArrayBackData_LRU_Ejection(t *testing.T) {
 	testRetrievableFrom(t, bd, entities, 900_000)
 }
 
+// TestArrayBackData_StalePrefixCollisionOnEjectedSlot is a regression test for a bug where Get
+// would miss an existing entity: when a slot's value is ejected from the underlying pool, the
+// slot retains the ejected value's 32-bit id prefix. If a later slot in the same bucket holds an
+// entity whose id prefix collides with the stale slot's prefix, Get would stop scanning at the
+// stale slot and erroneously report the entity as missing.
+func TestArrayBackData_StalePrefixCollisionOnEjectedSlot(t *testing.T) {
+	limit := uint32(8)
+	bd := NewCache[*unittest.MockEntity](limit,
+		8,
+		heropool.LRUEjection,
+		unittest.Logger(),
+		metrics.NewNoopCollector())
+
+	// two distinct identifiers sharing the same bucket (selected by bytes [0:8]) and the same
+	// 32-bit in-memory prefix (bytes [8:12]), differing only in the remaining bytes.
+	var idA, idB flow.Identifier
+	idA[8] = 0xaa
+	idB[8] = 0xaa
+	idA[12] = 0x01
+	idB[12] = 0x02
+
+	entityA := &unittest.MockEntity{Identifier: idA}
+	entityB := &unittest.MockEntity{Identifier: idB}
+
+	require.True(t, bd.Add(idA, entityA)) // occupies the first slot of the bucket
+	require.True(t, bd.Add(idB, entityB)) // occupies a later slot of the same bucket
+
+	// Fill the cache to capacity and one beyond, so that entityA (the least recently used
+	// entity) is ejected from the pool. Its slot keeps the stale 32-bit prefix, which equals
+	// entityB's prefix. All fillers use a different bucket and a different prefix, so they do
+	// not interfere with the crafted bucket.
+	for i := uint32(0); i < limit-1; i++ {
+		var id flow.Identifier
+		id[0] = byte(4*i + 1) // maps to a different bucket than idA/idB
+		id[8] = 0x11          // different 32-bit prefix than idA/idB
+		require.True(t, bd.Add(id, &unittest.MockEntity{Identifier: id}))
+	}
+
+	// entityA must have been ejected (the pool exceeded its limit by one)
+	_, ok := bd.Get(idA)
+	require.False(t, ok)
+
+	// entityB must still be retrievable, even though the stale slot with the colliding prefix
+	// precedes it in the bucket's scan order
+	actual, ok := bd.Get(idB)
+	require.True(t, ok)
+	require.Equal(t, entityB, actual)
+}
+
 // TestArrayBackData_No_Ejection evaluates correctness of Cache under the writing and retrieving
 // a heavy load of entities beyond its limit. With NoEjection mode, the cache should refuse to add extra entities beyond
 // its limit.

network/alsp/internal/cache.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/onflow/flow-go/model/flow"
 	"github.com/onflow/flow-go/module"
+	"github.com/onflow/flow-go/module/mempool"
 	herocache "github.com/onflow/flow-go/module/mempool/herocache/backdata"
 	"github.com/onflow/flow-go/module/mempool/herocache/backdata/heropool"
 	"github.com/onflow/flow-go/module/mempool/stdmap"
@@ -62,6 +63,10 @@ func NewSpamRecordCache(sizeLimit uint32, logger zerolog.Logger, collector modul
 //   - error any returned error should be considered as an irrecoverable error and indicates a bug.
 func (s *SpamRecordCache) AdjustWithInit(originId flow.Identifier, adjustFunc model.RecordAdjustFunc) (float64, error) {
 	var rErr error
+	// Penalty of the adjusted record, captured while the backend lock is still held: records are
+	// mutated in place by adjust functions, so reading the returned record's fields after
+	// Backend.AdjustWithInit has released the lock would race with concurrent adjustments.
+	var adjustedPenalty float64
 	wrapAdjustFunc := func(record *model.ProtocolSpamRecord) *model.ProtocolSpamRecord {
 		// Adjust the record.
 		adjustedRecord, err := adjustFunc(record)
@@ -70,14 +75,15 @@ func (s *SpamRecordCache) AdjustWithInit(originId flow.Identifier, adjustFunc mo
 			return record // returns the original record (reverse the adjustment).
 		}
 
+		adjustedPenalty = adjustedRecord.Penalty
 		// Return the adjusted record.
 		return adjustedRecord
 	}
 	initFunc := func() *model.ProtocolSpamRecord {
 		return s.recordFactory(originId)
 	}
 
-	adjustedRecord, adjusted := s.Backend.AdjustWithInit(originId, wrapAdjustFunc, initFunc)
+	_, adjusted := s.Backend.AdjustWithInit(originId, wrapAdjustFunc, initFunc)
 	if rErr != nil {
 		return 0, fmt.Errorf("failed to adjust record: %w", rErr)
 	}
@@ -86,7 +92,7 @@ func (s *SpamRecordCache) AdjustWithInit(originId flow.Identifier, adjustFunc mo
 		return 0, fmt.Errorf("adjustment failed for origin id %s", originId)
 	}
 
-	return adjustedRecord.Penalty, nil
+	return adjustedPenalty, nil
 }
 
 // Get returns the spam record of the given origin id.
@@ -97,19 +103,33 @@ func (s *SpamRecordCache) AdjustWithInit(originId flow.Identifier, adjustFunc mo
 // - the record and true if the record exists, nil and false otherwise.
 // Note that the returned record is a copy of the record in the cache (we do not want the caller to modify the record).
 func (s *SpamRecordCache) Get(originId flow.Identifier) (*model.ProtocolSpamRecord, bool) {
-	record, ok := s.Backend.Get(originId)
-	if !ok {
-		return nil, false
+	var copied *model.ProtocolSpamRecord
+	// the copy must be made while holding the backend lock (via Run): records are mutated in
+	// place by adjust functions, so copying the fields outside the lock would race with
+	// concurrent adjustments.
+	err := s.Backend.Run(func(backdata mempool.BackData[flow.Identifier, *model.ProtocolSpamRecord]) error {
+		record, ok := backdata.Get(originId)
+		if !ok {
+			return nil
+		}
+
+		// return a copy of the record (we do not want the caller to modify the record).
+		copied = &model.ProtocolSpamRecord{
+			OriginId:       record.OriginId,
+			Decay:          record.Decay,
+			CutoffCounter:  record.CutoffCounter,
+			Penalty:        record.Penalty,
+			DisallowListed: record.DisallowListed,
+		}
+		return nil
+	})
+	if err != nil {
+		// the Run closure above never returns an error; an error here indicates a bug in the
+		// backend implementation.
+		panic(fmt.Errorf("unexpected error while getting spam record from cache: %w", err))
 	}
 
-	// return a copy of the record (we do not want the caller to modify the record).
-	return &model.ProtocolSpamRecord{
-		OriginId:       record.OriginId,
-		Decay:          record.Decay,
-		CutoffCounter:  record.CutoffCounter,
-		Penalty:        record.Penalty,
-		DisallowListed: record.DisallowListed,
-	}, true
+	return copied, copied != nil
 }
 
 // Identities returns the list of identities of the nodes that have a spam record in the cache.

network/p2p/inspector/validation/control_message_validation_inspector.go

@@ -643,6 +643,11 @@ func (c *ControlMsgValidationInspector) inspectRpcPublishMessages(from peer.ID,
 	if sampleSize > totalMessages {
 		sampleSize = totalMessages
 	}
+	// Clone the slice before sampling: this method runs asynchronously on worker goroutines
+	// while the libp2p pubsub layer continues to process the same RPC object, so shuffling the
+	// RPC's own message slice in place would be a data race. The messages themselves are only
+	// read, hence a shallow clone suffices.
+	messages = slices.Clone(messages)
 	c.performSample(p2pmsg.RpcPublishMessage, uint(totalMessages), uint(sampleSize), func(i, j uint) {
 		messages[i], messages[j] = messages[j], messages[i]
 	})

network/p2p/scoring/internal/appSpecificScoreCache.go

@@ -79,9 +79,15 @@ func (a *AppSpecificScoreCache) AdjustWithInit(peerID peer.ID, score float64, ti
 		}
 	}
 	adjustLogic := func(record *appSpecificScoreRecord) *appSpecificScoreRecord {
-		record.Score = score
-		record.LastUpdated = time
-		return record
+		// copy-on-write: the cache stores records by pointer and Get reads the fields of the
+		// retrieved record after the backend lock is released; mutating the record in place
+		// here would be a data race. Returning a fresh record (which Adjust stores back)
+		// keeps previously returned records immutable.
+		return &appSpecificScoreRecord{
+			PeerID:      record.PeerID,
+			Score:       score,
+			LastUpdated: time,
+		}
 	}
 	_, adjusted := a.c.AdjustWithInit(p2p.MakeId(peerID), adjustLogic, initLogic)
 	if !adjusted {