Use rayon to speed up batch KZG verification (#7921)

Addresses #7866.


  Use Rayon to speed up batch KZG verification during range / backfill sync.

While I was analysing the traces, I also discovered a bug that resulted in only the first 128 columns in a chain segment batch being verified. This PR fixes it, so we might actually observe slower range sync due to more cells being KZG verified.

I've also updated the handling of batch KZG failure to only find the first invalid KZG column when verification fails as this gets very expensive during range/backfill sync.

beacon_node/beacon_chain/src/kzg_utils.rs

@@ -1,6 +1,6 @@
 use kzg::{
-    Blob as KzgBlob, Bytes48, CELLS_PER_EXT_BLOB, Cell as KzgCell, CellRef as KzgCellRef,
-    CellsAndKzgProofs, Error as KzgError, Kzg,
+    Blob as KzgBlob, Bytes48, Cell as KzgCell, CellRef as KzgCellRef, CellsAndKzgProofs,
+    Error as KzgError, Kzg,
 };
 use rayon::prelude::*;
 use ssz_types::{FixedVector, VariableList};
@@ -45,38 +45,11 @@ pub fn validate_blob<E: EthSpec>(
     kzg.verify_blob_kzg_proof(&kzg_blob, kzg_commitment, kzg_proof)
 }
 
-/// Validates a list of blobs along with their corresponding KZG commitments and
-/// cell proofs for the extended blobs.
-pub fn validate_blobs_and_cell_proofs<E: EthSpec>(
-    kzg: &Kzg,
-    blobs: Vec<&Blob<E>>,
-    cell_proofs: &[KzgProof],
-    kzg_commitments: &KzgCommitments<E>,
-) -> Result<(), KzgError> {
-    let cells = compute_cells::<E>(&blobs, kzg)?;
-    let cell_refs = cells.iter().map(|cell| cell.as_ref()).collect::<Vec<_>>();
-    let cell_indices = (0..blobs.len())
-        .flat_map(|_| 0..CELLS_PER_EXT_BLOB as u64)
-        .collect::<Vec<_>>();
-
-    let proofs = cell_proofs
-        .iter()
-        .map(|&proof| Bytes48::from(proof))
-        .collect::<Vec<_>>();
-
-    let commitments = kzg_commitments
-        .iter()
-        .flat_map(|&commitment| std::iter::repeat_n(Bytes48::from(commitment), CELLS_PER_EXT_BLOB))
-        .collect::<Vec<_>>();
-
-    kzg.verify_cell_proof_batch(&cell_refs, &proofs, cell_indices, &commitments)
-}
-
 /// Validate a batch of `DataColumnSidecar`.
 pub fn validate_data_columns<'a, E: EthSpec, I>(
     kzg: &Kzg,
     data_column_iter: I,
-) -> Result<(), KzgError>
+) -> Result<(), (Option<u64>, KzgError)>
 where
     I: Iterator<Item = &'a Arc<DataColumnSidecar<E>>> + Clone,
 {
@@ -88,8 +61,12 @@ where
     for data_column in data_column_iter {
         let col_index = data_column.index;
 
+        if data_column.column.is_empty() {
+            return Err((Some(col_index), KzgError::KzgVerificationFailed));
+        }
+
         for cell in &data_column.column {
-            cells.push(ssz_cell_to_crypto_cell::<E>(cell)?);
+            cells.push(ssz_cell_to_crypto_cell::<E>(cell).map_err(|e| (Some(col_index), e))?);
             column_indices.push(col_index);
         }
 
@@ -100,6 +77,19 @@ where
         for &commitment in &data_column.kzg_commitments {
             commitments.push(Bytes48::from(commitment));
         }
+
+        let expected_len = column_indices.len();
+
+        // We make this check at each iteration so that the error is attributable to a specific column
+        if cells.len() != expected_len
+            || proofs.len() != expected_len
+            || commitments.len() != expected_len
+        {
+            return Err((
+                Some(col_index),
+                KzgError::InconsistentArrayLength("Invalid data column".to_string()),
+            ));
+        }
     }
 
     kzg.verify_cell_proof_batch(&cells, &proofs, column_indices, &commitments)
@@ -418,7 +408,7 @@ pub fn reconstruct_data_columns<E: EthSpec>(
 mod test {
     use crate::kzg_utils::{
         blobs_to_data_column_sidecars, reconstruct_blobs, reconstruct_data_columns,
-        validate_blobs_and_cell_proofs,
+        validate_data_columns,
     };
     use bls::Signature;
     use eth2::types::BlobsBundle;
@@ -442,21 +432,20 @@ mod test {
         test_build_data_columns(&kzg, &spec);
         test_reconstruct_data_columns(&kzg, &spec);
         test_reconstruct_blobs_from_data_columns(&kzg, &spec);
-        test_verify_blob_and_cell_proofs(&kzg);
+        test_validate_data_columns(&kzg, &spec);
     }
 
     #[track_caller]
-    fn test_verify_blob_and_cell_proofs(kzg: &Kzg) {
-        let (blobs_bundle, _) = generate_blobs::<E>(3, ForkName::Fulu).unwrap();
-        let BlobsBundle {
-            blobs,
-            commitments,
-            proofs,
-        } = blobs_bundle;
-
-        let result =
-            validate_blobs_and_cell_proofs::<E>(kzg, blobs.iter().collect(), &proofs, &commitments);
+    fn test_validate_data_columns(kzg: &Kzg, spec: &ChainSpec) {
+        let num_of_blobs = 6;
+        let (signed_block, blobs, proofs) =
+            create_test_fulu_block_and_blobs::<E>(num_of_blobs, spec);
+        let blob_refs = blobs.iter().collect::<Vec<_>>();
+        let column_sidecars =
+            blobs_to_data_column_sidecars(&blob_refs, proofs.to_vec(), &signed_block, kzg, spec)
+                .unwrap();
 
+        let result = validate_data_columns::<E, _>(kzg, column_sidecars.iter());
         assert!(result.is_ok());
     }