From bd85c1897005cbb5ace1fc20f9ba506b11ec996f Mon Sep 17 00:00:00 2001
From: David Zhou <david.zhou.0110@gmail.com>
Date: Fri, 9 May 2025 16:56:15 -0400
Subject: [PATCH] TEST: Fixed race issue in existing test files caused by
 integration of goroutine. All tests passing

---
 eth/tracers/firehose.go                       | 39 ++++++++++---------
 .../tracetest/firehose/firehose_test.go       |  4 ++
 2 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/eth/tracers/firehose.go b/eth/tracers/firehose.go
index 6ea929bc47..2fb2e4a009 100644
--- a/eth/tracers/firehose.go
+++ b/eth/tracers/firehose.go
@@ -79,7 +79,6 @@ func NewTracingHooksFromFirehose(tracer *Firehose) *tracing.Hooks {
 		OnBlockEnd:       tracer.OnBlockEnd,
 		OnSkippedBlock:   tracer.OnSkippedBlock,
 		OnClose:          tracer.OnClose,
-		// TODO: OnClose
 
 		OnTxStart: tracer.OnTxStart,
 		OnTxEnd:   tracer.OnTxEnd,
@@ -193,7 +192,8 @@ type Firehose struct {
 
 	// Worker queuefor blockprinting
 	blockPrintQueue chan *blockPrintJob
-	workerWg        sync.WaitGroup
+	flushDone       sync.WaitGroup
+	closeOnce       sync.Once
 }
 
 const FirehoseProtocolVersion = "3.0"
@@ -257,12 +257,8 @@ func NewFirehose(config *FirehoseConfig) *Firehose {
 		}
 	}
 
-	// Worker goroutines
-	numWorkers := 1
-	firehose.workerWg.Add(numWorkers)
-	for i := 0; i < numWorkers; i++ {
-		go firehose.blockPrintWorker()
-	}
+	firehose.flushDone.Add(1)
+	go firehose.blockPrintWorker()
 
 	return firehose
 }
@@ -499,13 +495,13 @@ func (f *Firehose) OnBlockEnd(err error) {
 			f.fixOrdinalsForEndOfBlockChanges()
 		}
 
-		f.printBlockToFirehose(f.block, f.blockFinality)
-		//f.ensureInBlockAndNotInTrx()
-		//job := &blockPrintJob{
-		//	block:    f.block,
-		//	finality: f.blockFinality,
-		//}
-		//f.blockPrintQueue <- job
+		f.ensureInBlockAndNotInTrx()
+		// f.printBlockToFirehose(f.block, f.blockFinality)
+		job := &blockPrintJob{
+			block:    f.block,
+			finality: f.blockFinality,
+		}
+		f.blockPrintQueue <- job
 
 	} else {
 		// An error occurred, could have happen in transaction/call context, we must not check if in trx/call, only check in block
@@ -626,8 +622,7 @@ func (f *Firehose) reorderCallOrdinals(call *pbeth.Call, ordinalBase uint64) (or
 
 func (f *Firehose) OnClose() {
 	log.Info("Firehose closing: waiting for worker goroutines to finish and shutting down channels")
-	close(f.blockPrintQueue)
-	f.workerWg.Wait()
+	f.CloseBlockPrintQueue()
 }
 
 func (f *Firehose) OnSystemCallStart() {
@@ -2827,9 +2822,15 @@ type blockPrintJob struct {
 }
 
 func (f *Firehose) blockPrintWorker() {
-	defer f.workerWg.Done()
-
+	defer f.flushDone.Done()
 	for job := range f.blockPrintQueue {
 		f.printBlockToFirehose(job.block, job.finality)
 	}
 }
+
+func (f *Firehose) CloseBlockPrintQueue() {
+	f.closeOnce.Do(func() {
+		close(f.blockPrintQueue)
+		f.flushDone.Wait()
+	})
+}
diff --git a/eth/tracers/internal/tracetest/firehose/firehose_test.go b/eth/tracers/internal/tracetest/firehose/firehose_test.go
index b226241996..cd959260cc 100644
--- a/eth/tracers/internal/tracetest/firehose/firehose_test.go
+++ b/eth/tracers/internal/tracetest/firehose/firehose_test.go
@@ -48,6 +48,8 @@ func TestFirehosePrestate(t *testing.T) {
 
 				runPrestateBlock(t, filepath.Join(folder, "prestate.json"), tracingHooks)
 
+				tracer.CloseBlockPrintQueue()
+
 				genesisLine, blockLines, unknownLines := readTracerFirehoseLines(t, tracer)
 				require.Len(t, unknownLines, 0, "Lines:\n%s", strings.Join(slicesMap(unknownLines, func(l unknownLine) string { return "- '" + string(l) + "'" }), "\n"))
 				require.NotNil(t, genesisLine)
@@ -203,6 +205,8 @@ func testBlockTracesCorrectly(t *testing.T, genesisSpec *core.Genesis, engine co
 			n, err := chain.InsertChain(blocks)
 			require.NoError(t, err, "failed to insert chain block %d", n)
 
+			tracer.CloseBlockPrintQueue()
+
 			assertBlockEquals(t, tracer, filepath.Join("testdata", goldenDir, string(model)), len(blocks))
 		})
 	}