fix: recover on restart (#2313)

Resolves #2220

tlDr; Payloads require sequential height when send to the pubsub topics.
Failed on gaps after restart.

* Buffers in `HeaderCh, DataCh` were not persisted
* The `syncService.Close` method never completed successfully due to
timeout and non empty topic subscribers. This lead to data loss
* Any error during `WriteToStoreAndBroadcast` was closing the consumer
loop only, not the app
* On graceful shutdown, not waited for the Go routines
(producers/consumers) which caused data loss


The `TestSlowConsumers` replicates the behaviour that I have seen on my
box. When the heights have gaps after restart, the channel consumers in
`node/full.go` stop their work and everything freezes.

With `TestHeaderSyncServiceRestart`, I confirm that the go-headers store
is persisting the data when submitted.


<!--
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.

NOTE: PR titles should follow semantic commits:
https://www.conventionalcommits.org/en/v1.0.0/
-->

## 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. 

Ex: Closes #<issue number>
-->


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

- **New Features**
- Improved block broadcasting with enhanced modularity and error
handling.
- Added robust tests simulating slow consumers and verifying persistence
across restarts.
- Enhanced synchronization service with explicit subscription management
and restart tests.
- Improved node process management in end-to-end tests, including
command-based tracking and graceful shutdown.

- **Bug Fixes**
- Ensured all background tasks complete before node shutdown for safer
operation.
- Fixed potential issues with process cleanup and synchronization during
tests.

- **Refactor**
- Replaced direct channel usage with broadcaster interfaces for block
publishing.
- Streamlined goroutine management and error handling in node
operations.
  - Simplified and improved test setup and teardown procedures.

- **Chores**
  - Updated dependencies to explicitly require key libraries.
  - Reorganized import statements for clarity.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: alpe

block/manager.go

@@ -17,6 +17,7 @@ import (
 	goheader "github.com/celestiaorg/go-header"
 	ds "github.com/ipfs/go-datastore"
 	"github.com/libp2p/go-libp2p/core/crypto"
+	"golang.org/x/sync/errgroup"
 
 	coreda "github.com/rollkit/rollkit/core/da"
 	coreexecutor "github.com/rollkit/rollkit/core/execution"
@@ -46,9 +47,6 @@ const (
 	// This is temporary solution. It will be removed in future versions.
 	maxSubmitAttempts = 30
 
-	// Applies to most channels, 100 is a large enough buffer to avoid blocking
-	channelLength = 100
-
 	// Applies to the headerInCh and dataInCh, 10000 is a large enough number for headers per DA block.
 	eventInChLength = 10000
 
@@ -90,6 +88,10 @@ type BatchData struct {
 	Data [][]byte
 }
 
+type broadcaster[T any] interface {
+	WriteToStoreAndBroadcast(ctx context.Context, payload T) error
+}
+
 // Manager is responsible for aggregating transactions into blocks.
 type Manager struct {
 	lastState types.State
@@ -104,8 +106,8 @@ type Manager struct {
 
 	daHeight *atomic.Uint64
 
-	HeaderCh chan *types.SignedHeader
-	DataCh   chan *types.Data
+	headerBroadcaster broadcaster[*types.SignedHeader]
+	dataBroadcaster   broadcaster[*types.Data]
 
 	headerInCh  chan NewHeaderEvent
 	headerStore goheader.Store[*types.SignedHeader]
@@ -268,6 +270,8 @@ func NewManager(
 	logger log.Logger,
 	headerStore goheader.Store[*types.SignedHeader],
 	dataStore goheader.Store[*types.Data],
+	headerBroadcaster broadcaster[*types.SignedHeader],
+	dataBroadcaster broadcaster[*types.Data],
 	seqMetrics *Metrics,
 	gasPrice float64,
 	gasMultiplier float64,
@@ -326,15 +330,15 @@ func NewManager(
 	daH.Store(s.DAHeight)
 
 	m := &Manager{
-		signer:    signer,
-		config:    config,
-		genesis:   genesis,
-		lastState: s,
-		store:     store,
-		daHeight:  &daH,
+		signer:            signer,
+		config:            config,
+		genesis:           genesis,
+		lastState:         s,
+		store:             store,
+		daHeight:          &daH,
+		headerBroadcaster: headerBroadcaster,
+		dataBroadcaster:   dataBroadcaster,
 		// channels are buffered to avoid blocking on input/output operations, buffer sizes are arbitrary
-		HeaderCh:            make(chan *types.SignedHeader, channelLength),
-		DataCh:              make(chan *types.Data, channelLength),
 		headerInCh:          make(chan NewHeaderEvent, eventInChLength),
 		dataInCh:            make(chan NewDataEvent, eventInChLength),
 		headerStoreCh:       make(chan struct{}, 1),
@@ -654,24 +658,14 @@ func (m *Manager) publishBlockInternal(ctx context.Context) error {
 
 	m.recordMetrics(data)
 
-	// Check for shut down event prior to sending the header and block to
-	// their respective channels. The reason for checking for the shutdown
-	// event separately is due to the inconsistent nature of the select
-	// statement when multiple cases are satisfied.
-	select {
-	case <-ctx.Done():
-		return fmt.Errorf("unable to send header and block, context done: %w", ctx.Err())
-	default:
+	g, ctx := errgroup.WithContext(ctx)
+	g.Go(func() error { return m.headerBroadcaster.WriteToStoreAndBroadcast(ctx, header) })
+	g.Go(func() error { return m.dataBroadcaster.WriteToStoreAndBroadcast(ctx, data) })
+	if err := g.Wait(); err != nil {
+		return err
 	}
 
-	// Publish header to channel so that header exchange service can broadcast
-	m.HeaderCh <- header
-
-	// Publish block to channel so that block exchange service can broadcast
-	m.DataCh <- data
-
 	m.logger.Debug("successfully proposed header", "proposer", hex.EncodeToString(header.ProposerAddress), "height", headerHeight)
-
 	return nil
 }
 

block/publish_block2_test.go

@@ -0,0 +1,245 @@
+package block
+
+import (
+	"context"
+	cryptoRand "crypto/rand"
+	"errors"
+	"fmt"
+	"math/rand"
+	"path/filepath"
+	"sync"
+	"testing"
+	"time"
+
+	"cosmossdk.io/log"
+	ds "github.com/ipfs/go-datastore"
+	ktds "github.com/ipfs/go-datastore/keytransform"
+	syncdb "github.com/ipfs/go-datastore/sync"
+	logging "github.com/ipfs/go-log/v2"
+	"github.com/libp2p/go-libp2p/core/crypto"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	coresequencer "github.com/rollkit/rollkit/core/sequencer"
+	"github.com/rollkit/rollkit/pkg/config"
+	genesispkg "github.com/rollkit/rollkit/pkg/genesis"
+	"github.com/rollkit/rollkit/pkg/p2p"
+	"github.com/rollkit/rollkit/pkg/p2p/key"
+	"github.com/rollkit/rollkit/pkg/signer"
+	"github.com/rollkit/rollkit/pkg/signer/noop"
+	"github.com/rollkit/rollkit/pkg/store"
+	rollkitSync "github.com/rollkit/rollkit/pkg/sync"
+	"github.com/rollkit/rollkit/types"
+)
+
+func TestSlowConsumers(t *testing.T) {
+	logging.SetDebugLogging()
+	blockTime := 100 * time.Millisecond
+	specs := map[string]struct {
+		headerConsumerDelay time.Duration
+		dataConsumerDelay   time.Duration
+	}{
+		"slow header consumer": {
+			headerConsumerDelay: blockTime * 2,
+			dataConsumerDelay:   0,
+		},
+		"slow data consumer": {
+			headerConsumerDelay: 0,
+			dataConsumerDelay:   blockTime * 2,
+		},
+		"both slow": {
+			headerConsumerDelay: blockTime,
+			dataConsumerDelay:   blockTime,
+		},
+		"both fast": {
+			headerConsumerDelay: 0,
+			dataConsumerDelay:   0,
+		},
+	}
+	for name, spec := range specs {
+		t.Run(name, func(t *testing.T) {
+			workDir := t.TempDir()
+			dbm := syncdb.MutexWrap(ds.NewMapDatastore())
+			ctx, cancel := context.WithCancel(t.Context())
+
+			pk, _, err := crypto.GenerateEd25519Key(cryptoRand.Reader)
+			require.NoError(t, err)
+			noopSigner, err := noop.NewNoopSigner(pk)
+			require.NoError(t, err)
+
+			manager, headerSync, dataSync := setupBlockManager(t, ctx, workDir, dbm, blockTime, noopSigner)
+			var lastCapturedDataPayload *types.Data
+			var lastCapturedHeaderPayload *types.SignedHeader
+			manager.dataBroadcaster = capturingTailBroadcaster[*types.Data](spec.dataConsumerDelay, &lastCapturedDataPayload, dataSync)
+			manager.headerBroadcaster = capturingTailBroadcaster[*types.SignedHeader](spec.headerConsumerDelay, &lastCapturedHeaderPayload, headerSync)
+
+			blockTime := manager.config.Node.BlockTime.Duration
+			aggCtx, aggCancel := context.WithCancel(ctx)
+			errChan := make(chan error, 1)
+			var wg sync.WaitGroup
+			wg.Add(1)
+			go func() {
+				manager.AggregationLoop(aggCtx, errChan)
+				wg.Done()
+			}()
+
+			// wait for messages to pile up
+			select {
+			case err := <-errChan:
+				require.NoError(t, err)
+			case <-time.After(spec.dataConsumerDelay + spec.headerConsumerDelay + 3*blockTime):
+			}
+			aggCancel()
+			wg.Wait() // await aggregation loop to finish
+			t.Log("shutting down block manager")
+			require.NoError(t, dataSync.Stop(ctx))
+			require.NoError(t, headerSync.Stop(ctx))
+			cancel()
+			require.NotNil(t, lastCapturedHeaderPayload)
+			require.NotNil(t, lastCapturedDataPayload)
+
+			t.Log("restart with new block manager")
+			ctx, cancel = context.WithCancel(t.Context())
+			manager, headerSync, dataSync = setupBlockManager(t, ctx, workDir, dbm, blockTime, noopSigner)
+
+			var firstCapturedDataPayload *types.Data
+			var firstCapturedHeaderPayload *types.SignedHeader
+			manager.dataBroadcaster = capturingHeadBroadcaster[*types.Data](0, &firstCapturedDataPayload, dataSync)
+			manager.headerBroadcaster = capturingHeadBroadcaster[*types.SignedHeader](0, &firstCapturedHeaderPayload, headerSync)
+			go manager.AggregationLoop(ctx, errChan)
+			select {
+			case err := <-errChan:
+				require.NoError(t, err)
+			case <-time.After(spec.dataConsumerDelay + spec.headerConsumerDelay + 3*blockTime):
+			}
+			cancel()
+			require.NotNil(t, firstCapturedHeaderPayload)
+			assert.InDelta(t, lastCapturedDataPayload.Height(), firstCapturedDataPayload.Height(), 1)
+			require.NotNil(t, firstCapturedDataPayload)
+			assert.InDelta(t, lastCapturedHeaderPayload.Height(), firstCapturedHeaderPayload.Height(), 1)
+		})
+	}
+}
+
+func capturingTailBroadcaster[T interface{ Height() uint64 }](waitDuration time.Duration, target *T, next ...broadcaster[T]) broadcaster[T] {
+	var lastHeight uint64
+	return broadcasterFn[T](func(ctx context.Context, payload T) error {
+		if payload.Height() <= lastHeight {
+			panic(fmt.Sprintf("got height %d, want %d", payload.Height(), lastHeight+1))
+		}
+
+		time.Sleep(waitDuration)
+		lastHeight = payload.Height()
+		*target = payload
+		var err error
+		for _, n := range next {
+			err = errors.Join(n.WriteToStoreAndBroadcast(ctx, payload))
+		}
+
+		return err
+	})
+}
+
+func capturingHeadBroadcaster[T interface{ Height() uint64 }](waitDuration time.Duration, target *T, next ...broadcaster[T]) broadcaster[T] {
+	var once sync.Once
+	return broadcasterFn[T](func(ctx context.Context, payload T) error {
+		once.Do(func() {
+			*target = payload
+		})
+		var err error
+		for _, n := range next {
+			err = errors.Join(n.WriteToStoreAndBroadcast(ctx, payload))
+		}
+		time.Sleep(waitDuration)
+		return err
+	})
+}
+
+type broadcasterFn[T any] func(ctx context.Context, payload T) error
+
+func (b broadcasterFn[T]) WriteToStoreAndBroadcast(ctx context.Context, payload T) error {
+	return b(ctx, payload)
+}
+
+func setupBlockManager(t *testing.T, ctx context.Context, workDir string, mainKV ds.Batching, blockTime time.Duration, signer signer.Signer) (*Manager, *rollkitSync.HeaderSyncService, *rollkitSync.DataSyncService) {
+	t.Helper()
+	nodeConfig := config.DefaultConfig
+	nodeConfig.Node.BlockTime = config.DurationWrapper{Duration: blockTime}
+	nodeConfig.RootDir = workDir
+	nodeKey, err := key.LoadOrGenNodeKey(filepath.Dir(nodeConfig.ConfigPath()))
+	require.NoError(t, err)
+
+	proposerAddr, err := signer.GetAddress()
+	require.NoError(t, err)
+	genesisDoc := genesispkg.Genesis{
+		ChainID:            "test-chain-id",
+		GenesisDAStartTime: time.Now(),
+		InitialHeight:      1,
+		ProposerAddress:    proposerAddr,
+	}
+
+	logger := log.NewTestLogger(t)
+	p2pClient, err := p2p.NewClient(nodeConfig, nodeKey, mainKV, logger, p2p.NopMetrics())
+	require.NoError(t, err)
+
+	// Start p2p client before creating sync service
+	err = p2pClient.Start(ctx)
+	require.NoError(t, err)
+
+	const RollkitPrefix = "0"
+	ktds.Wrap(mainKV, ktds.PrefixTransform{Prefix: ds.NewKey(RollkitPrefix)})
+	headerSyncService, err := rollkitSync.NewHeaderSyncService(mainKV, nodeConfig, genesisDoc, p2pClient, logger.With("module", "HeaderSyncService"))
+	require.NoError(t, err)
+	require.NoError(t, headerSyncService.Start(ctx))
+	dataSyncService, err := rollkitSync.NewDataSyncService(mainKV, nodeConfig, genesisDoc, p2pClient, logger.With("module", "DataSyncService"))
+	require.NoError(t, err)
+	require.NoError(t, dataSyncService.Start(ctx))
+
+	result, err := NewManager(
+		ctx,
+		signer,
+		nodeConfig,
+		genesisDoc,
+		store.New(mainKV),
+		&mockExecutor{},
+		coresequencer.NewDummySequencer(),
+		nil,
+		logger.With("module", "BlockManager"),
+		headerSyncService.Store(),
+		dataSyncService.Store(),
+		nil,
+		nil,
+		NopMetrics(),
+		1.,
+		1.,
+	)
+	require.NoError(t, err)
+	return result, headerSyncService, dataSyncService
+}
+
+type mockExecutor struct{}
+
+func (m mockExecutor) InitChain(ctx context.Context, genesisTime time.Time, initialHeight uint64, chainID string) (stateRoot []byte, maxBytes uint64, err error) {
+	return bytesN(32), 10_000, nil
+}
+
+func (m mockExecutor) GetTxs(ctx context.Context) ([][]byte, error) {
+	panic("implement me")
+}
+
+func (m mockExecutor) ExecuteTxs(ctx context.Context, txs [][]byte, blockHeight uint64, timestamp time.Time, prevStateRoot []byte) (updatedStateRoot []byte, maxBytes uint64, err error) {
+	return bytesN(32), 10_000, nil
+
+}
+
+func (m mockExecutor) SetFinal(ctx context.Context, blockHeight uint64) error {
+	return nil
+}
+
+var rnd = rand.New(rand.NewSource(1)) //nolint:gosec // test code only
+
+func bytesN(n int) []byte {
+	data := make([]byte, n)
+	_, _ = rnd.Read(data)
+	return data
+}