DEV: Implemented onClose()

This commit is contained in:
David Zhou 2025-05-09 16:20:57 -04:00
parent e9e894d090
commit 7520990fa6
2 changed files with 18 additions and 15 deletions

View file

@ -78,6 +78,7 @@ func NewTracingHooksFromFirehose(tracer *Firehose) *tracing.Hooks {
OnBlockStart: tracer.OnBlockStart, OnBlockStart: tracer.OnBlockStart,
OnBlockEnd: tracer.OnBlockEnd, OnBlockEnd: tracer.OnBlockEnd,
OnSkippedBlock: tracer.OnSkippedBlock, OnSkippedBlock: tracer.OnSkippedBlock,
OnClose: tracer.OnClose,
// TODO: OnClose // TODO: OnClose
OnTxStart: tracer.OnTxStart, OnTxStart: tracer.OnTxStart,
@ -498,13 +499,13 @@ func (f *Firehose) OnBlockEnd(err error) {
f.fixOrdinalsForEndOfBlockChanges() f.fixOrdinalsForEndOfBlockChanges()
} }
// f.printBlockToFirehose(f.block, f.blockFinality) f.printBlockToFirehose(f.block, f.blockFinality)
f.ensureInBlockAndNotInTrx() //f.ensureInBlockAndNotInTrx()
job := &blockPrintJob{ //job := &blockPrintJob{
block: f.block, // block: f.block,
finality: f.blockFinality, // finality: f.blockFinality,
} //}
f.blockPrintQueue <- job //f.blockPrintQueue <- job
} else { } else {
// An error occurred, could have happen in transaction/call context, we must not check if in trx/call, only check in block // An error occurred, could have happen in transaction/call context, we must not check if in trx/call, only check in block
@ -623,6 +624,12 @@ func (f *Firehose) reorderCallOrdinals(call *pbeth.Call, ordinalBase uint64) (or
return call.EndOrdinal return call.EndOrdinal
} }
func (f *Firehose) OnClose() {
log.Info("Firehose closing: waiting for worker goroutines to finish and shutting down channels")
close(f.blockPrintQueue)
f.workerWg.Wait()
}
func (f *Firehose) OnSystemCallStart() { func (f *Firehose) OnSystemCallStart() {
firehoseInfo("system call start") firehoseInfo("system call start")
f.ensureInBlockAndNotInTrx() f.ensureInBlockAndNotInTrx()
@ -2826,8 +2833,3 @@ func (f *Firehose) blockPrintWorker() {
f.printBlockToFirehose(job.block, job.finality) f.printBlockToFirehose(job.block, job.finality)
} }
} }
func (f *Firehose) Shutdown() {
close(f.blockPrintQueue)
f.workerWg.Wait()
}

View file

@ -37,7 +37,7 @@ func TestFirehose_BlockPrintsToFirehose_SingleBlock(t *testing.T) {
f.OnBlockEnd(nil) f.OnBlockEnd(nil)
} }
f.Shutdown() f.OnClose()
output := f.InternalTestingBuffer().String() output := f.InternalTestingBuffer().String()
@ -50,11 +50,12 @@ func TestFirehose_BlockPrintsToFirehose_SingleBlock(t *testing.T) {
fields := strings.SplitN(line, " ", 4) fields := strings.SplitN(line, " ", 4)
if len(fields) >= 3 { if len(fields) >= 3 {
require.Equal(t, "FIRE", fields[0])
require.Equal(t, "BLOCK", fields[1])
outNumber = append(outNumber, fields[2]) outNumber = append(outNumber, fields[2])
} }
} }
require.Contains(t, output, "FIRE BLOCK", "expected FIRE BLOCK output not found")
require.Equal(t, []string{"123", "124", "125"}, outNumber) require.Equal(t, []string{"123", "124", "125"}, outNumber)
} }
@ -89,7 +90,7 @@ func TestFirehose_BlocksPrintToFirehose_MultipleBlocksInOrder(t *testing.T) {
f.OnBlockEnd(nil) f.OnBlockEnd(nil)
} }
f.Shutdown() f.OnClose()
output := f.InternalTestingBuffer().String() output := f.InternalTestingBuffer().String()
extractedBlocks := extractBlocksFromOutput(t, output) extractedBlocks := extractBlocksFromOutput(t, output)