fix: address slow shutdowns in aggregation loop (#1297)

<!--
Please read and fill out this form before submitting your PR.

Please make sure you have reviewed our contributors guide before
submitting your
first PR.
-->

## Overview

<!-- 
Please provide an explanation of the PR, including the appropriate
context,
background, goal, and rationale. If there is an issue with this
information,
please provide a tl;dr and link the issue. 
-->

I started trying to debug the test failure in #1058 and #1095 by first
cleaning up the test output. I replaced the stdout logger with the file
logger and that highlighted some panic conditions of writing to a closed
file, which meant there were some hanging threads. The test that was
easiest to catch the panic was `TestTx`.

One of the issues was the indexer didn't have a ctx so wasn't shutting
down cleanly.

The second was the edge case of the aggregation loop execution taking
some time, resulting in the node shutting down in the middle. I added
the ctx check at the end so that any updates are ideally saved to the
store before we exit.

I tested this locally by running the failing test in a loop. Prior to
these changes it panicked every time, afterwards, you can run the test
in a loop without failure.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

## Summary by CodeRabbit

Refactor:
- Enhanced the node shutdown process in `block/manager.go` to skip
publishing if the node is already shut down, improving system
efficiency.
- Modified the `createAndStartIndexerService` function in `node/full.go`
to require a context during initialization, allowing for better control
over the service's lifecycle.

Tests:
- Updated `node/full_client_test.go` to use a cancellable context and a
new file logger, improving test reliability and logging capabilities.
- Modified `state/txindex/indexer_service_test.go` to use a cancellable
context, ensuring proper cleanup after tests.

Style:
- Updated `state/txindex/indexer_service.go` to include a context in the
`IndexerService` struct and handle its cancellation, aligning with the
context-oriented design.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Manav Aggarwal <[email protected]>

Author: MSevey

block/manager.go

@@ -699,6 +699,13 @@ func (m *Manager) publishBlock(ctx context.Context) error {
 		return err
 	}
 
+	// Check if the node has shutdown prior to publishing to channels
+	select {
+	case <-ctx.Done():
+		return nil
+	default:
+	}
+
 	// Publish header to channel so that header exchange service can broadcast
 	m.HeaderCh <- &block.SignedHeader
 

node/full.go

@@ -435,7 +435,7 @@ func createAndStartIndexerService(
 	txIndexer = kv.NewTxIndex(ctx, kvStore)
 	blockIndexer = blockidxkv.New(ctx, newPrefixKV(kvStore, "block_events"))
 
-	indexerService := txindex.NewIndexerService(txIndexer, blockIndexer, eventBus)
+	indexerService := txindex.NewIndexerService(ctx, txIndexer, blockIndexer, eventBus)
 	indexerService.SetLogger(logger.With("module", "txindex"))
 
 	if err := indexerService.Start(); err != nil {

node/full_client_test.go

@@ -19,7 +19,6 @@ import (
 	"github.com/cometbft/cometbft/crypto/ed25519"
 	"github.com/cometbft/cometbft/crypto/encoding"
 	"github.com/cometbft/cometbft/libs/bytes"
-	"github.com/cometbft/cometbft/libs/log"
 	"github.com/cometbft/cometbft/p2p"
 	"github.com/cometbft/cometbft/proxy"
 	cmtypes "github.com/cometbft/cometbft/types"
@@ -30,6 +29,7 @@ import (
 	"github.com/rollkit/rollkit/config"
 	mockda "github.com/rollkit/rollkit/da/mock"
 	"github.com/rollkit/rollkit/store"
+	test "github.com/rollkit/rollkit/test/log"
 	"github.com/rollkit/rollkit/test/mocks"
 	"github.com/rollkit/rollkit/types"
 	abciconv "github.com/rollkit/rollkit/types/abci"
@@ -79,7 +79,7 @@ func getRPC(t *testing.T) (*mocks.Application, *FullClient) {
 	key, _, _ := crypto.GenerateEd25519Key(crand.Reader)
 	signingKey, _, _ := crypto.GenerateEd25519Key(crand.Reader)
 	ctx := context.Background()
-	node, err := newFullNode(ctx, config.NodeConfig{DALayer: "mock"}, key, signingKey, proxy.NewLocalClientCreator(app), &cmtypes.GenesisDoc{ChainID: "test"}, log.TestingLogger())
+	node, err := newFullNode(ctx, config.NodeConfig{DALayer: "mock"}, key, signingKey, proxy.NewLocalClientCreator(app), &cmtypes.GenesisDoc{ChainID: "test"}, test.NewFileLogger(t))
 	require.NoError(err)
 	require.NotNil(node)
 
@@ -177,7 +177,7 @@ func TestGenesisChunked(t *testing.T) {
 	mockApp.On(InitChain, mock.Anything).Return(abci.ResponseInitChain{})
 	privKey, _, _ := crypto.GenerateEd25519Key(crand.Reader)
 	signingKey, _, _ := crypto.GenerateEd25519Key(crand.Reader)
-	n, _ := newFullNode(context.Background(), config.NodeConfig{DALayer: "mock"}, privKey, signingKey, proxy.NewLocalClientCreator(mockApp), genDoc, log.TestingLogger())
+	n, _ := newFullNode(context.Background(), config.NodeConfig{DALayer: "mock"}, privKey, signingKey, proxy.NewLocalClientCreator(mockApp), genDoc, test.NewFileLogger(t))
 
 	rpc := NewFullClient(n)
 
@@ -483,15 +483,17 @@ func TestTx(t *testing.T) {
 	mockApp.On(InitChain, mock.Anything).Return(abci.ResponseInitChain{})
 	key, _, _ := crypto.GenerateEd25519Key(crand.Reader)
 	genesisValidators, signingKey := getGenesisValidatorSetWithSigner()
-	node, err := newFullNode(context.Background(), config.NodeConfig{
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	node, err := newFullNode(ctx, config.NodeConfig{
 		DALayer:    "mock",
 		Aggregator: true,
 		BlockManagerConfig: config.BlockManagerConfig{
 			BlockTime: 1 * time.Second, // blocks must be at least 1 sec apart for adjacent headers to get verified correctly
 		}},
 		key, signingKey, proxy.NewLocalClientCreator(mockApp),
 		&cmtypes.GenesisDoc{ChainID: "test", Validators: genesisValidators},
-		log.TestingLogger())
+		test.NewFileLogger(t))
 	require.NoError(err)
 	require.NotNil(node)
 
@@ -509,21 +511,21 @@ func TestTx(t *testing.T) {
 		require.NoError(rpc.node.Stop())
 	}()
 	tx1 := cmtypes.Tx("tx1")
-	res, err := rpc.BroadcastTxSync(context.Background(), tx1)
+	res, err := rpc.BroadcastTxSync(ctx, tx1)
 	assert.NoError(err)
 	assert.NotNil(res)
 
 	time.Sleep(2 * time.Second)
 
-	resTx, errTx := rpc.Tx(context.Background(), res.Hash, true)
+	resTx, errTx := rpc.Tx(ctx, res.Hash, true)
 	assert.NoError(errTx)
 	assert.NotNil(resTx)
 	assert.EqualValues(tx1, resTx.Tx)
 	assert.EqualValues(res.Hash, resTx.Hash)
 
 	tx2 := cmtypes.Tx("tx2")
 	assert.Panics(func() {
-		resTx, errTx := rpc.Tx(context.Background(), tx2.Hash(), true)
+		resTx, errTx := rpc.Tx(ctx, tx2.Hash(), true)
 		assert.Nil(resTx)
 		assert.Error(errTx)
 	})
@@ -727,7 +729,7 @@ func createGenesisValidators(t *testing.T, numNodes int, appCreator func(require
 	dalc := &mockda.DataAvailabilityLayerClient{}
 	ds, err := store.NewDefaultInMemoryKVStore()
 	require.Nil(err)
-	err = dalc.Init([8]byte{}, nil, ds, log.TestingLogger())
+	err = dalc.Init([8]byte{}, nil, ds, test.NewFileLogger(t))
 	require.Nil(err)
 	err = dalc.Start()
 	require.Nil(err)
@@ -755,7 +757,7 @@ func createGenesisValidators(t *testing.T, numNodes int, appCreator func(require
 			signingKey,
 			proxy.NewLocalClientCreator(apps[i]),
 			&cmtypes.GenesisDoc{ChainID: "test", Validators: genesisValidators},
-			log.TestingLogger(),
+			test.NewFileLogger(t),
 		)
 		require.NoError(err)
 		require.NotNil(nodes[i])
@@ -918,7 +920,7 @@ func TestMempool2Nodes(t *testing.T) {
 		BlockManagerConfig: config.BlockManagerConfig{
 			BlockTime: 1 * time.Second,
 		},
-	}, key1, signingKey1, proxy.NewLocalClientCreator(app), &cmtypes.GenesisDoc{ChainID: "test"}, log.TestingLogger())
+	}, key1, signingKey1, proxy.NewLocalClientCreator(app), &cmtypes.GenesisDoc{ChainID: "test"}, test.NewFileLogger(t))
 	require.NoError(err)
 	require.NotNil(node1)
 
@@ -928,7 +930,7 @@ func TestMempool2Nodes(t *testing.T) {
 			ListenAddress: "/ip4/127.0.0.1/tcp/9002",
 			Seeds:         "/ip4/127.0.0.1/tcp/9001/p2p/" + id1.Pretty(),
 		},
-	}, key2, signingKey2, proxy.NewLocalClientCreator(app), &cmtypes.GenesisDoc{ChainID: "test"}, log.TestingLogger())
+	}, key2, signingKey2, proxy.NewLocalClientCreator(app), &cmtypes.GenesisDoc{ChainID: "test"}, test.NewFileLogger(t))
 	require.NoError(err)
 	require.NotNil(node1)
 
@@ -1022,7 +1024,7 @@ func TestStatus(t *testing.T) {
 			ChainID:    "test",
 			Validators: genesisValidators,
 		},
-		log.TestingLogger(),
+		test.NewFileLogger(t),
 	)
 	require.NoError(err)
 	require.NotNil(node)
@@ -1133,7 +1135,7 @@ func TestFutureGenesisTime(t *testing.T) {
 			GenesisTime:   genesisTime,
 			Validators:    genesisValidators,
 		},
-		log.TestingLogger())
+		test.NewFileLogger(t))
 	require.NoError(err)
 	require.NotNil(node)
 

state/txindex/indexer_service.go

@@ -20,19 +20,22 @@ const (
 type IndexerService struct {
 	service.BaseService
 
+	ctx context.Context
+
 	txIdxr    TxIndexer
 	blockIdxr indexer.BlockIndexer
 	eventBus  *types.EventBus
 }
 
 // NewIndexerService returns a new service instance.
 func NewIndexerService(
+	ctx context.Context,
 	txIdxr TxIndexer,
 	blockIdxr indexer.BlockIndexer,
 	eventBus *types.EventBus,
 ) *IndexerService {
 
-	is := &IndexerService{txIdxr: txIdxr, blockIdxr: blockIdxr, eventBus: eventBus}
+	is := &IndexerService{ctx: ctx, txIdxr: txIdxr, blockIdxr: blockIdxr, eventBus: eventBus}
 	is.BaseService = *service.NewBaseService(nil, "IndexerService", is)
 	return is
 }
@@ -59,6 +62,8 @@ func (is *IndexerService) OnStart() error {
 	go func() {
 		for {
 			select {
+			case <-is.ctx.Done():
+				return
 			case <-blockHeadersSub.Cancelled():
 				return
 			case msg := <-blockHeadersSub.Out():
@@ -68,6 +73,8 @@ func (is *IndexerService) OnStart() error {
 
 				for i := int64(0); i < eventDataHeader.NumTxs; i++ {
 					select {
+					case <-is.ctx.Done():
+						return
 					case <-txsSub.Cancelled():
 						return
 					case msg2 := <-txsSub.Out():

state/txindex/indexer_service_test.go

@@ -36,7 +36,9 @@ func TestIndexerServiceIndexesBlocks(t *testing.T) {
 	prefixStore := (ktds.Wrap(kvStore, ktds.PrefixTransform{Prefix: ds.NewKey("block_events")}).Children()[0]).(ds.TxnDatastore)
 	blockIndexer := blockidxkv.New(context.Background(), prefixStore)
 
-	service := txindex.NewIndexerService(txIndexer, blockIndexer, eventBus)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	service := txindex.NewIndexerService(ctx, txIndexer, blockIndexer, eventBus)
 	service.SetLogger(log.TestingLogger())
 	err = service.Start()
 	require.NoError(t, err)