fix(manager): refactor freezing node using context

block/manager.go

@@ -115,6 +115,8 @@ type Manager struct {
 
 	// validates all non-finalized state updates from settlement, checking there is consistency between DA and P2P blocks, and the information in the state update.
 	SettlementValidator *SettlementValidator
+
+	Cancel context.CancelFunc
 }
 
 // NewManager creates new block Manager.
@@ -197,6 +199,7 @@ func NewManager(
 
 // Start starts the block manager.
 func (m *Manager) Start(ctx context.Context) error {
+	ctx, m.Cancel = context.WithCancel(ctx)
 	// Check if InitChain flow is needed
 	if m.State.IsGenesis() {
 		m.logger.Info("Running InitChain")
@@ -379,8 +382,7 @@ func (m *Manager) setFraudHandler(handler *FreezeHandler) {
 }
 
 func (m *Manager) freezeNode(ctx context.Context, err error) {
+	m.logger.Info("Freezing node", "err", err)
 	uevent.MustPublish(ctx, m.Pubsub, &events.DataHealthStatus{Error: err}, events.HealthStatusList)
-	if m.RunMode == RunModeFullNode {
-		m.unsubscribeFullNodeEvents(ctx)
-	}
+	m.Cancel()
 }

block/manager_test.go

@@ -233,7 +233,6 @@ func TestApplyCachedBlocks_WithFraudCheck(t *testing.T) {
 	t.Log("Taking the manager out of sync by submitting a batch")
 	manager.DAClient = testutil.GetMockDALC(log.TestingLogger())
 	manager.Retriever = manager.DAClient.(da.BatchRetriever)
-
 	mockExecutor := &blockmocks.MockExecutorI{}
 	manager.Executor = mockExecutor
 	mockExecutor.On("GetAppInfo").Return(&abci.ResponseInfo{
@@ -243,12 +242,11 @@ func TestApplyCachedBlocks_WithFraudCheck(t *testing.T) {
 
 	// Check that handle fault is called
 	manager.FraudHandler = block.NewFreezeHandler(manager)
-
 	fraudEventReceived := make(chan *events.DataHealthStatus, 1)
 
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
-
+	_, manager.Cancel = context.WithCancel(context.Background())
 	go event.MustSubscribe(
 		ctx,
 		manager.Pubsub,

block/modes.go

@@ -86,17 +86,3 @@ func (m *Manager) subscribeFullNodeEvents(ctx context.Context) {
 	go uevent.MustSubscribe(ctx, m.Pubsub, p2pGossipLoop, p2p.EventQueryNewGossipedBlock, m.OnReceivedBlock, m.logger)
 	go uevent.MustSubscribe(ctx, m.Pubsub, p2pBlocksyncLoop, p2p.EventQueryNewBlockSyncBlock, m.OnReceivedBlock, m.logger)
 }
-
-func (m *Manager) unsubscribeFullNodeEvents(ctx context.Context) {
-	// unsubscribe for specific event (clientId)
-	unsubscribe := func(clientId string) {
-		err := m.Pubsub.UnsubscribeAll(ctx, clientId)
-		if err != nil {
-			m.logger.Error("Unsubscribe", "clientId", clientId, "error", err)
-		}
-	}
-	unsubscribe(syncLoop)
-	unsubscribe(validateLoop)
-	unsubscribe(p2pGossipLoop)
-	unsubscribe(p2pBlocksyncLoop)
-}

block/produce.go

@@ -38,7 +38,7 @@ func (m *Manager) ProduceBlockLoop(ctx context.Context, bytesProducedC chan int)
 	for {
 		select {
 		case <-ctx.Done():
-			return nil
+			return ctx.Err()
 		case <-ticker.C:
 			// Only produce if I'm the current rollapp proposer.
 			if !m.AmIProposerOnRollapp() {

block/submit.go

@@ -60,6 +60,7 @@ func SubmitLoopInner(
 		// 'trigger': this thread is responsible for waking up the submitter when a new block arrives, and back-pressures the block production loop
 		// if it gets too far ahead.
 		for {
+
 			if maxBatchSkew*maxBatchBytes < pendingBytes.Load() {
 				// too much stuff is pending submission
 				// we block here until we get a progress nudge from the submitter thread

block/sync.go

@@ -50,12 +50,14 @@ func (m *Manager) SettlementSyncLoop(ctx context.Context) error {
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
-
 		case <-m.settlementSyncingC:
-
 			m.logger.Info("syncing to target height", "targetHeight", m.LastSettlementHeight.Load())
 
 			for currH := m.State.NextHeight(); currH <= m.LastSettlementHeight.Load(); currH = m.State.NextHeight() {
+				// if context has been cancelled, stop syncing
+				if ctx.Err() != nil {
+					return ctx.Err()
+				}
 				// if we have the block locally, we don't need to fetch it from the DA.
 				// it will only happen in case of rollback.
 				err := m.applyLocalBlock(currH)

block/validate.go

@@ -29,7 +29,6 @@ func (m *Manager) SettlementValidateLoop(ctx context.Context) error {
 		case <-ctx.Done():
 			return ctx.Err()
 		case <-m.settlementValidationC:
-
 			targetValidationHeight := min(m.LastSettlementHeight.Load(), m.State.Height())
 			m.logger.Info("validating state updates to target height", "targetHeight", targetValidationHeight)