feat: implement reconstruct and broadcast data columns (#15023)

* Implement reconstructAndBroadcastDataColumns

* Fix merge error

* Fix tests

* Minor changes.

---------

Co-authored-by: Manu NALEPA <[email protected]>

Author: 0x00101010

beacon-chain/blockchain/testing/mock.go

@@ -74,6 +74,7 @@ type ChainService struct {
 	BlockSlot                   primitives.Slot
 	SyncingRoot                 [32]byte
 	Blobs                       []blocks.VerifiedROBlob
+	DataColumns                 []blocks.VerifiedRODataColumn
 	TargetRoot                  [32]byte
 }
 
@@ -703,7 +704,8 @@ func (c *ChainService) ReceiveBlob(_ context.Context, b blocks.VerifiedROBlob) e
 }
 
 // ReceiveDataColumn implements the same method in chain service
-func (*ChainService) ReceiveDataColumn(_ blocks.VerifiedRODataColumn) error {
+func (c *ChainService) ReceiveDataColumn(dc blocks.VerifiedRODataColumn) error {
+	c.DataColumns = append(c.DataColumns, dc)
 	return nil
 }
 

beacon-chain/execution/engine_client.go

@@ -642,13 +642,13 @@ func (s *Service) ReconstructBlobSidecars(ctx context.Context, block interfaces.
 }
 
 // ReconstructDataColumnSidecars reconstructs the verified data column sidecars for a given beacon block.
-// It retrieves the KZG commitments from the block body, fetches the associated blobs and cell proofs,
+// It retrieves the KZG commitments from the block body, fetches the associated blobs and cell proofs from the EL,
 // and constructs the corresponding verified read-only data column sidecars.
 func (s *Service) ReconstructDataColumnSidecars(ctx context.Context, block interfaces.ReadOnlySignedBeaconBlock, blockRoot [32]byte) ([]blocks.VerifiedRODataColumn, error) {
 	blockBody := block.Block().Body()
 	kzgCommitments, err := blockBody.BlobKzgCommitments()
 	if err != nil {
-		return nil, wrapWithBlockRoot(err, blockRoot, "could not get blob KZG commitments")
+		return nil, wrapWithBlockRoot(err, blockRoot, "blob KZG commitments")
 	}
 
 	// Collect KZG hashes for all blobs
@@ -657,20 +657,18 @@ func (s *Service) ReconstructDataColumnSidecars(ctx context.Context, block inter
 		kzgHashes = append(kzgHashes, primitives.ConvertKzgCommitmentToVersionedHash(commitment))
 	}
 
-	// Fetch all blobs from EL
-	blobs, err := s.GetBlobsV2(ctx, kzgHashes)
+	// Fetch all blobsAndCellsProofs from EL
+	blobsAndCellsProofs, err := s.GetBlobsV2(ctx, kzgHashes)
 	if err != nil {
-		return nil, wrapWithBlockRoot(err, blockRoot, "could not get blobs")
+		return nil, wrapWithBlockRoot(err, blockRoot, "get blobs V2")
 	}
 
-	for _, blobAndCellProofs := range blobs {
+	var cellsAndProofs []kzg.CellsAndProofs
+	for _, blobAndCellProofs := range blobsAndCellsProofs {
 		if blobAndCellProofs == nil {
 			return nil, wrapWithBlockRoot(errors.New("unable to reconstruct data column sidecars, did not get all blobs from EL"), blockRoot, "")
 		}
-	}
 
-	var cellsAndProofs []kzg.CellsAndProofs
-	for _, blobAndCellProofs := range blobs {
 		var blob kzg.Blob
 		copy(blob[:], blobAndCellProofs.Blob)
 		cells, err := kzg.ComputeCells(&blob)
@@ -713,7 +711,7 @@ func (s *Service) ReconstructDataColumnSidecars(ctx context.Context, block inter
 		verifiedRODataColumns[i] = blocks.NewVerifiedRODataColumn(roDataColumn)
 	}
 
-	log.WithField("root", fmt.Sprintf("%x", blockRoot)).Debug("Data columns reconstructed successfully")
+	log.WithField("root", fmt.Sprintf("%x", blockRoot)).Debug("Data columns reconstructed successfully from EL")
 
 	return verifiedRODataColumns, nil
 }

beacon-chain/execution/engine_client_test.go

@@ -2504,7 +2504,7 @@ func TestReconstructDataColumnSidecars(t *testing.T) {
 	ctx := context.Background()
 	t.Run("GetBlobsV2 is not supported", func(t *testing.T) {
 		_, err := client.ReconstructDataColumnSidecars(ctx, sb, r)
-		require.ErrorContains(t, "could not get blobs", err)
+		require.ErrorContains(t, "get blobs V2 for block", err)
 	})
 
 	t.Run("receiving all blobs", func(t *testing.T) {

beacon-chain/execution/testing/mock_engine_client.go

@@ -38,6 +38,8 @@ type EngineClient struct {
 	ErrGetPayload               error
 	BlobSidecars                []blocks.VerifiedROBlob
 	ErrorBlobSidecars           error
+	DataColumnSidecars          []blocks.VerifiedRODataColumn
+	ErrorDataColumnSidecars     error
 }
 
 // NewPayload --
@@ -114,7 +116,7 @@ func (e *EngineClient) ReconstructBlobSidecars(context.Context, interfaces.ReadO
 }
 
 func (e *EngineClient) ReconstructDataColumnSidecars(ctx context.Context, block interfaces.ReadOnlySignedBeaconBlock, blockRoot [32]byte) ([]blocks.VerifiedRODataColumn, error) {
-	return nil, nil
+	return e.DataColumnSidecars, e.ErrorDataColumnSidecars
 }
 
 // GetTerminalBlockHash --

beacon-chain/sync/initial-sync/blocks_fetcher_utils.go

@@ -8,9 +8,7 @@ import (
 
 	"github.com/libp2p/go-libp2p/core/peer"
 	"github.com/pkg/errors"
-	"github.com/prysmaticlabs/prysm/v5/beacon-chain/core/peerdas"
 	coreTime "github.com/prysmaticlabs/prysm/v5/beacon-chain/core/time"
-	"github.com/prysmaticlabs/prysm/v5/beacon-chain/p2p"
 	p2pTypes "github.com/prysmaticlabs/prysm/v5/beacon-chain/p2p/types"
 	"github.com/prysmaticlabs/prysm/v5/cmd/beacon-chain/flags"
 	"github.com/prysmaticlabs/prysm/v5/config/params"
@@ -382,33 +380,6 @@ func (f *blocksFetcher) calculateHeadAndTargetEpochs() (headEpoch, targetEpoch p
 	return headEpoch, targetEpoch, peers
 }
 
-// custodyGroupsFromPeer compute all the custody groups indexed by peer.
-func (f *blocksFetcher) custodyGroupsFromPeer(peers map[peer.ID]bool) (map[peer.ID]map[uint64]bool, error) {
-	peerCount := len(peers)
-
-	custodyGroupsByPeer := make(map[peer.ID]map[uint64]bool, peerCount)
-	for peer := range peers {
-		// Get the node ID from the peer ID.
-		nodeID, err := p2p.ConvertPeerIDToNodeID(peer)
-		if err != nil {
-			return nil, errors.Wrap(err, "convert peer ID to node ID")
-		}
-
-		// Get the custody group count of the peer.
-		custodyGroupCount := f.p2p.CustodyGroupCountFromPeer(peer)
-
-		// Retrieve the peer info.
-		peerInfo, _, err := peerdas.Info(nodeID, custodyGroupCount)
-		if err != nil {
-			return nil, errors.Wrap(err, "peer info")
-		}
-
-		custodyGroupsByPeer[peer] = peerInfo.CustodyGroups
-	}
-
-	return custodyGroupsByPeer, nil
-}
-
 // uint64MapToSortedSlice produces a sorted uint64 slice from a map.
 func uint64MapToSortedSlice(input map[uint64]bool) []uint64 {
 	output := make([]uint64, 0, len(input))
@@ -420,54 +391,6 @@ func uint64MapToSortedSlice(input map[uint64]bool) []uint64 {
 	return output
 }
 
-// `filterPeerWhichCustodyAtLeastOneDataColumn` filters peers which custody at least one data column
-// specified in `neededDataColumns`. It returns also a list of descriptions for non admissible peers.
-func filterPeerWhichCustodyAtLeastOneDataColumn(
-	neededDataColumns map[uint64]bool,
-	inputDataColumnsByPeer map[peer.ID]map[uint64]bool,
-) (map[peer.ID]map[uint64]bool, []string) {
-	// Get the count of needed data columns.
-	neededDataColumnsCount := uint64(len(neededDataColumns))
-
-	// Create pretty needed data columns for logs.
-	var neededDataColumnsLog interface{} = "all"
-	numberOfColumns := params.BeaconConfig().NumberOfColumns
-
-	if neededDataColumnsCount < numberOfColumns {
-		neededDataColumnsLog = uint64MapToSortedSlice(neededDataColumns)
-	}
-
-	outputDataColumnsByPeer := make(map[peer.ID]map[uint64]bool, len(inputDataColumnsByPeer))
-	descriptions := make([]string, 0)
-
-outerLoop:
-	for peer, peerCustodyDataColumns := range inputDataColumnsByPeer {
-		for neededDataColumn := range neededDataColumns {
-			if peerCustodyDataColumns[neededDataColumn] {
-				outputDataColumnsByPeer[peer] = peerCustodyDataColumns
-
-				continue outerLoop
-			}
-		}
-
-		peerCustodyColumnsCount := uint64(len(peerCustodyDataColumns))
-		var peerCustodyColumnsLog interface{} = "all"
-
-		if peerCustodyColumnsCount < numberOfColumns {
-			peerCustodyColumnsLog = uint64MapToSortedSlice(peerCustodyDataColumns)
-		}
-
-		description := fmt.Sprintf(
-			"peer %s: does not custody any needed column, custody columns: %v, needed columns: %v",
-			peer, peerCustodyColumnsLog, neededDataColumnsLog,
-		)
-
-		descriptions = append(descriptions, description)
-	}
-
-	return outputDataColumnsByPeer, descriptions
-}
-
 // Filter peers with head epoch lower than our target epoch for ByRange requests.
 func (f *blocksFetcher) filterPeersByTargetSlot(peers []peer.ID, targetSlot primitives.Slot) ([]peer.ID, []string, error) {
 	filteredPeers := make([]peer.ID, 0, len(peers))

beacon-chain/sync/subscriber_beacon_blocks.go

@@ -7,6 +7,7 @@ import (
 	"path"
 
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/blockchain"
+	"github.com/prysmaticlabs/prysm/v5/beacon-chain/core/peerdas"
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/core/transition/interop"
 	"github.com/prysmaticlabs/prysm/v5/config/features"
 	"github.com/prysmaticlabs/prysm/v5/consensus-types/blocks"
@@ -35,6 +36,7 @@ func (s *Service) beaconBlockSubscriber(ctx context.Context, msg proto.Message)
 		return err
 	}
 
+	// TODO: do we only do this for super nodes?
 	go s.reconstructAndBroadcastBlobs(ctx, signed)
 
 	if err := s.cfg.chain.ReceiveBlock(ctx, signed, root, nil); err != nil {
@@ -60,18 +62,67 @@ func (s *Service) beaconBlockSubscriber(ctx context.Context, msg proto.Message)
 }
 
 // reconstructAndBroadcastBlobs processes and broadcasts blob sidecars for a given beacon block.
-// This function reconstructs the blob sidecars from the EL using the block's KZG commitments,
-// broadcasts the reconstructed blobs over P2P, and saves them into the blob storage.
 func (s *Service) reconstructAndBroadcastBlobs(ctx context.Context, block interfaces.ReadOnlySignedBeaconBlock) {
-	if block.Version() < version.Deneb {
+	if block.Version() >= version.Fulu {
+		s.reconstructAndBroadcastBlobsInDataColumn(ctx, block)
 		return
 	}
 
-	// TODO: Apply the equivalent strategy for data columns.
-	if block.Version() >= version.Fulu {
+	if block.Version() >= version.Deneb {
+		s.reconstructAndBroadcastFullBlobs(ctx, block)
+		return
+	}
+}
+
+// reconstructAndBroadcastBlobsInDataColumn reconstructs and broadcasts blobs in data column format for a given beacon block, it also saves data column sidecars into the blob storage.
+func (s *Service) reconstructAndBroadcastBlobsInDataColumn(ctx context.Context, block interfaces.ReadOnlySignedBeaconBlock) {
+	blockRoot, err := block.Block().HashTreeRoot()
+	if err != nil {
+		log.WithError(err).Error("Failed to calculate block root")
 		return
 	}
 
+	if s.cfg.blobStorage == nil {
+		log.Warn("Blob storage is not enabled, skip saving data column, but continue to reconstruct and broadcast blobs")
+	}
+
+	// when this function is called, it's from the time when the block is received, so in almost all situations we need to get the data column from EL instead of the blob storage.
+	sidecars, err := s.cfg.executionReconstructor.ReconstructDataColumnSidecars(ctx, block, blockRoot)
+	if err != nil {
+		log.WithError(err).Debug("Cannot reconstruct data column sidecars after receiving the block")
+		return
+	}
+
+	nodeID := s.cfg.p2p.NodeID()
+	s.cfg.custodyInfo.Mut.RLock()
+	defer s.cfg.custodyInfo.Mut.RUnlock()
+	samplingSize := s.cfg.custodyInfo.CustodyGroupSamplingSize(peerdas.Actual)
+	info, _, err := peerdas.Info(nodeID, samplingSize)
+	if err != nil {
+		log.WithError(err).Error("Failed to get peer info")
+		return
+	}
+
+	// Broadcast data column and then save to db (if needs to be in custody)
+	for _, sidecar := range sidecars {
+		if !info.CustodyColumns[sidecar.ColumnIndex] {
+			continue
+		}
+
+		// first broadcast the data column
+		if err := s.cfg.p2p.BroadcastDataColumn(ctx, blockRoot, sidecar.ColumnIndex, sidecar.DataColumnSidecar); err != nil {
+			log.WithFields(dataColumnFields(sidecar.RODataColumn)).WithError(err).Error("Failed to broadcast data column")
+		}
+
+		if err := s.receiveDataColumn(ctx, sidecar); err != nil {
+			log.WithFields(dataColumnFields(sidecar.RODataColumn)).WithError(err).Error("Failed to receive data column")
+		}
+	}
+}
+
+// reconstructAndBroadcastFullBlobs reconstructs the blob sidecars from the EL using the block's KZG commitments,
+// broadcasts the reconstructed blobs over P2P, and saves them into the blob storage.
+func (s *Service) reconstructAndBroadcastFullBlobs(ctx context.Context, block interfaces.ReadOnlySignedBeaconBlock) {
 	startTime, err := slots.ToTime(uint64(s.cfg.chain.GenesisTime().Unix()), block.Block().Slot())
 	if err != nil {
 		log.WithError(err).Error("Failed to convert slot to time")