Refactor retriever and fix more tests

Manav-Aggarwal · Apr 29, 2025 · Apr 30, 2025 · Apr 30, 2025 · Apr 30, 2025 · May 2, 2025
commit 6cb6b1b1060b3e5063bce2323bc3f351165d8771
diff --git a/block/retriever.go b/block/retriever.go
@@ -35,7 +35,7 @@ func (m *Manager) RetrieveLoop(ctx context.Context) {
 		case <-blobsFoundCh:
 		}
 		daHeight := m.daHeight.Load()
-		err := m.processNextDAHeaderAndBlock(ctx)
+		err := m.processNextDAHeaderAndData(ctx)
 		if err != nil && ctx.Err() == nil {
 			// if the requested da height is not yet available, wait silently, otherwise log the error and wait
 			if !m.areAllErrorsHeightFromFuture(err) {
@@ -52,7 +52,7 @@ func (m *Manager) RetrieveLoop(ctx context.Context) {
 	}
 }
 
-func (m *Manager) processNextDAHeaderAndBlock(ctx context.Context) error {
+func (m *Manager) processNextDAHeaderAndData(ctx context.Context) error {
 	select {
 	case <-ctx.Done():
 		return ctx.Err()
@@ -72,66 +72,15 @@ func (m *Manager) processNextDAHeaderAndBlock(ctx context.Context) error {
 		blobsResp, fetchErr := m.fetchBlobs(ctx, daHeight)
 		if fetchErr == nil {
 			if blobsResp.Code == coreda.StatusNotFound {
-				m.logger.Debug("no data found", "daHeight", daHeight, "reason", blobsResp.Message)
+				m.logger.Debug("no blob data found", "daHeight", daHeight, "reason", blobsResp.Message)
 				return nil
 			}
 			m.logger.Debug("retrieved potential data", "n", len(blobsResp.Data), "daHeight", daHeight)
 			for _, bz := range blobsResp.Data {
-				// Try to decode as header first
-				header := new(types.SignedHeader)
-				var headerPb pb.SignedHeader
-				err := proto.Unmarshal(bz, &headerPb)
-				if err == nil {
-					err = header.FromProto(&headerPb)
-					if err == nil {
-						// early validation to reject junk headers
-						if !m.isUsingExpectedSingleSequencer(header) {
-							m.logger.Debug("skipping header from unexpected sequencer",
-								"headerHeight", header.Height(),
-								"headerHash", header.Hash().String())
-							continue
-						}
-						headerHash := header.Hash().String()
-						m.headerCache.SetDAIncluded(headerHash)
-						m.sendNonBlockingSignalToDAIncluderCh()
-						m.logger.Info("header marked as DA included", "headerHeight", header.Height(), "headerHash", headerHash)
-						if !m.headerCache.IsSeen(headerHash) {
-							select {
-							case <-ctx.Done():
-								return fmt.Errorf("unable to send header to headerInCh, context done: %w", ctx.Err())
-							default:
-								m.logger.Warn("headerInCh backlog full, dropping header", "daHeight", daHeight)
-							}
-							m.headerInCh <- NewHeaderEvent{header, daHeight}
-						}
-						continue
-					}
-				}
-
-				// If not a header, try to decode as batch
-				var batchPb pb.Batch
-				err = proto.Unmarshal(bz, &batchPb)
-				if err == nil {
-					data := &types.Data{
-						Txs: make(types.Txs, len(batchPb.Txs)),
-					}
-					for i, tx := range batchPb.Txs {
-						data.Txs[i] = types.Tx(tx)
-					}
-					dataHashStr := data.DACommitment().String()
-					m.dataCache.SetDAIncluded(dataHashStr)
-					m.sendNonBlockingSignalToDAIncluderCh()
-					m.logger.Info("batch marked as DA included", "batchHash", dataHashStr, "daHeight", daHeight)
-					if !m.dataCache.IsSeen(dataHashStr) {
-						select {
-						case <-ctx.Done():
-							return fmt.Errorf("unable to send batch to dataInCh, context done: %w", ctx.Err())
-						default:
-							m.logger.Warn("dataInCh backlog full, dropping batch", "daHeight", daHeight)
-						}
-						m.dataInCh <- NewDataEvent{data, daHeight}
-					}
+				if m.handlePotentialHeader(ctx, bz, daHeight) {
+					continue
 				}
+				m.handlePotentialBatch(ctx, bz, daHeight)
 			}
 			return nil
 		}
@@ -148,6 +97,73 @@ func (m *Manager) processNextDAHeaderAndBlock(ctx context.Context) error {
 	return err
 }
 
+// handlePotentialHeader tries to decode and process a header. Returns true if successful or skipped, false if not a header.
+func (m *Manager) handlePotentialHeader(ctx context.Context, bz []byte, daHeight uint64) bool {
+	header := new(types.SignedHeader)
+	var headerPb pb.SignedHeader
+	err := proto.Unmarshal(bz, &headerPb)
+	if err != nil {
+		m.logger.Debug("failed to unmarshal header", "error", err)
+		return false
+	}
+	err = header.FromProto(&headerPb)
+	if err != nil {
+		// treat as handled, but not valid
+		m.logger.Debug("failed to decode unmarshalled header", "error", err)
+		return true
+	}
+	// early validation to reject junk headers
+	if !m.isUsingExpectedSingleSequencer(header) {
+		m.logger.Debug("skipping header from unexpected sequencer",
+			"headerHeight", header.Height(),
+			"headerHash", header.Hash().String())
+		return true
+	}
+	headerHash := header.Hash().String()
+	m.headerCache.SetDAIncluded(headerHash)
+	m.sendNonBlockingSignalToDAIncluderCh()
+	m.logger.Info("header marked as DA included", "headerHeight", header.Height(), "headerHash", headerHash)
+	if !m.headerCache.IsSeen(headerHash) {
+		select {
+		case <-ctx.Done():
+			return true
+		default:
+			m.logger.Warn("headerInCh backlog full, dropping header", "daHeight", daHeight)
+		}
+		m.headerInCh <- NewHeaderEvent{header, daHeight}
+	}
+	return true
+}
+
+// handlePotentialBatch tries to decode and process a batch. No return value.
+func (m *Manager) handlePotentialBatch(ctx context.Context, bz []byte, daHeight uint64) {
+	var batchPb pb.Batch
+	err := proto.Unmarshal(bz, &batchPb)
+	if err != nil {
+		m.logger.Debug("failed to unmarshal batch", "error", err)
+		return
+	}
+	data := &types.Data{
+		Txs: make(types.Txs, len(batchPb.Txs)),
+	}
+	for i, tx := range batchPb.Txs {
+		data.Txs[i] = types.Tx(tx)
+	}
+	dataHashStr := data.DACommitment().String()
+	m.dataCache.SetDAIncluded(dataHashStr)
+	m.sendNonBlockingSignalToDAIncluderCh()
+	m.logger.Info("batch marked as DA included", "batchHash", dataHashStr, "daHeight", daHeight)
+	if !m.dataCache.IsSeen(dataHashStr) {
+		select {
+		case <-ctx.Done():
+			return
+		default:
+			m.logger.Warn("dataInCh backlog full, dropping batch", "daHeight", daHeight)
+		}
+		m.dataInCh <- NewDataEvent{data, daHeight}
+	}
+}
+
 // areAllErrorsHeightFromFuture checks if all errors in a joined error are ErrHeightFromFutureStr
 func (m *Manager) areAllErrorsHeightFromFuture(err error) bool {
 	if err == nil {

diff --git a/block/retriever_test.go b/block/retriever_test.go
@@ -102,8 +102,8 @@ func setupManagerForRetrieverTest(t *testing.T, initialDAHeight uint64) (*Manage
 	return manager, mockDAClient, mockStore, mockLogger, manager.headerCache, manager.dataCache, cancel
 }
 
-// TestProcessNextDAHeader_Success_SingleHeaderAndBatch tests the processNextDAHeaderAndBlock function for a single header and block.
-func TestProcessNextDAHeader_Success_SingleHeaderAndBatch(t *testing.T) {
+// TestProcessNextDAHeader_Success_SingleHeaderAndData tests the processNextDAHeaderAndData function for a single header and data.
+func TestProcessNextDAHeader_Success_SingleHeaderAndData(t *testing.T) {
 	daHeight := uint64(20)
 	blockHeight := uint64(100)
 	manager, mockDAClient, mockStore, _, headerCache, dataCache, cancel := setupManagerForRetrieverTest(t, daHeight)
@@ -146,7 +146,7 @@ func TestProcessNextDAHeader_Success_SingleHeaderAndBatch(t *testing.T) {
 	}, nil).Once()
 
 	ctx := context.Background()
-	err = manager.processNextDAHeaderAndBlock(ctx)
+	err = manager.processNextDAHeaderAndData(ctx)
 	require.NoError(t, err)
 
 	// Validate header event
@@ -170,13 +170,14 @@ func TestProcessNextDAHeader_Success_SingleHeaderAndBatch(t *testing.T) {
 	case <-time.After(100 * time.Millisecond):
 		t.Fatal("Expected block data event not received")
 	}
-	assert.True(t, dataCache.IsDAIncluded(blockData.DACommitment().String()), "Block data hash should be marked as DA included in cache")
+	assert.True(t, dataCache.IsDAIncluded(blockData.DACommitment().String()), "Block data commitment should be marked as DA included in cache")
 
 	mockDAClient.AssertExpectations(t)
 	mockStore.AssertExpectations(t)
 }
 
-func TestProcessNextDAHeader_NotFound(t *testing.T) {
+// TestProcessNextDAHeaderAndData_NotFound tests the processNextDAHeaderAndData function for a NotFound error.
+func TestProcessNextDAHeaderAndData_NotFound(t *testing.T) {
 	daHeight := uint64(25)
 	manager, mockDAClient, _, _, _, _, cancel := setupManagerForRetrieverTest(t, daHeight)
 	defer cancel()
@@ -185,7 +186,7 @@ func TestProcessNextDAHeader_NotFound(t *testing.T) {
 	}, nil).Once()
 
 	ctx := context.Background()
-	err := manager.processNextDAHeaderAndBlock(ctx)
+	err := manager.processNextDAHeaderAndData(ctx)
 	require.NoError(t, err)
 
 	select {
@@ -194,11 +195,17 @@ func TestProcessNextDAHeader_NotFound(t *testing.T) {
 	default:
 	}
 
+	select {
+	case <-manager.dataInCh:
+		t.Fatal("No data event should be received for NotFound")
+	default:
+	}
+
 	mockDAClient.AssertExpectations(t)
 }
 
-// TestProcessNextDAHeader_UnmarshalError tests the processNextDAHeaderAndBlock function for an unmarshal error.
-func TestProcessNextDAHeader_UnmarshalError(t *testing.T) {
+// TestProcessNextDAHeaderAndData_UnmarshalHeaderError tests the processNextDAHeaderAndData function for an unmarshal error.
+func TestProcessNextDAHeaderAndData_UnmarshalHeaderError(t *testing.T) {
 	daHeight := uint64(30)
 	manager, mockDAClient, _, mockLogger, _, _, cancel := setupManagerForRetrieverTest(t, daHeight)
 	defer cancel()
@@ -212,23 +219,29 @@ func TestProcessNextDAHeader_UnmarshalError(t *testing.T) {
 
 	mockLogger.ExpectedCalls = nil
 	mockLogger.On("Debug", "failed to unmarshal header", mock.Anything).Return().Once()
+	mockLogger.On("Debug", "failed to unmarshal batch", mock.Anything).Return().Once()
 	mockLogger.On("Debug", mock.Anything, mock.Anything).Maybe() // Allow other debug logs
 
 	ctx := context.Background()
-	err := manager.processNextDAHeaderAndBlock(ctx)
+	err := manager.processNextDAHeaderAndData(ctx)
 	require.NoError(t, err)
 
 	select {
 	case <-manager.headerInCh:
 		t.Fatal("No header event should be received for unmarshal error")
 	default:
 	}
+	select {
+	case <-manager.dataInCh:
+		t.Fatal("No data event should be received for unmarshal error")
+	default:
+	}
 
 	mockDAClient.AssertExpectations(t)
 	mockLogger.AssertExpectations(t)
 }
 
-// TestProcessNextDAHeader_UnexpectedSequencer tests the processNextDAHeaderAndBlock function for an unexpected sequencer.
+// TestProcessNextDAHeaderAndData_UnexpectedSequencer tests the processNextDAHeaderAndData function for an unexpected sequencer.
 func TestProcessNextDAHeader_UnexpectedSequencer(t *testing.T) {
 	daHeight := uint64(35)
 	blockHeight := uint64(110)
@@ -261,7 +274,7 @@ func TestProcessNextDAHeader_UnexpectedSequencer(t *testing.T) {
 	mockLogger.On("Debug", mock.Anything, mock.Anything).Maybe() // Allow other debug logs
 
 	ctx := context.Background()
-	err = manager.processNextDAHeaderAndBlock(ctx)
+	err = manager.processNextDAHeaderAndData(ctx)
 	require.NoError(t, err)
 
 	select {
@@ -275,6 +288,7 @@ func TestProcessNextDAHeader_UnexpectedSequencer(t *testing.T) {
 	mockLogger.AssertExpectations(t)
 }
 
+// TestProcessNextDAHeader_FetchError_RetryFailure tests the processNextDAHeaderAndBlock function for a fetch error.
 func TestProcessNextDAHeader_FetchError_RetryFailure(t *testing.T) {
 	daHeight := uint64(40)
 	manager, mockDAClient, _, _, _, _, cancel := setupManagerForRetrieverTest(t, daHeight)
@@ -287,7 +301,7 @@ func TestProcessNextDAHeader_FetchError_RetryFailure(t *testing.T) {
 	}, fetchErr).Times(dAFetcherRetries)
 
 	ctx := context.Background()
-	err := manager.processNextDAHeaderAndBlock(ctx)
+	err := manager.processNextDAHeaderAndData(ctx)
 	require.Error(t, err)
 	assert.ErrorContains(t, err, fetchErr.Error(), "Expected the final error after retries")
 
@@ -300,7 +314,7 @@ func TestProcessNextDAHeader_FetchError_RetryFailure(t *testing.T) {
 	mockDAClient.AssertExpectations(t)
 }
 
-// TestProcessNextDAHeader_HeaderAlreadySeen tests the processNextDAHeaderAndBlock function for a header that has already been seen.
+// TestProcessNextDAHeader_HeaderAlreadySeen tests the processNextDAHeaderAndData function for a header that has already been seen.
 func TestProcessNextDAHeader_HeaderAlreadySeen(t *testing.T) {
 	daHeight := uint64(45)
 	blockHeight := uint64(120)
@@ -331,7 +345,7 @@ func TestProcessNextDAHeader_HeaderAlreadySeen(t *testing.T) {
 	binary.LittleEndian.PutUint64(heightBytes, blockHeight)
 
 	ctx := context.Background()
-	err = manager.processNextDAHeaderAndBlock(ctx)
+	err = manager.processNextDAHeaderAndData(ctx)
 	require.NoError(t, err)
 
 	select {