Refactor

Author: michaelsproul

testing/web3signer_tests/src/lib.rs

@@ -605,12 +605,13 @@ mod tests {
         })
         .await
         .assert_signatures_match("attestation", |pubkey, validator_store| async move {
-            let mut attestation = get_attestation();
+            let attestation = get_attestation();
             validator_store
-                .sign_attestation(pubkey, 0, &mut attestation, Epoch::new(0))
+                .sign_attestations(vec![(pubkey, 0, attestation)])
                 .await
-                .unwrap();
-            attestation
+                .unwrap()
+                .pop()
+                .unwrap()
         })
         .await
         .assert_signatures_match("signed_aggregate", |pubkey, validator_store| async move {
@@ -820,8 +821,6 @@ mod tests {
             block
         };
 
-        let current_epoch = Epoch::new(5);
-
         TestingRig::new(
             network,
             slashing_protection_config,
@@ -830,48 +829,47 @@ mod tests {
         )
         .await
         .assert_signatures_match("first_attestation", |pubkey, validator_store| async move {
-            let mut attestation = first_attestation();
+            let attestation = first_attestation();
             validator_store
-                .sign_attestation(pubkey, 0, &mut attestation, current_epoch)
+                .sign_attestations(vec![(pubkey, 0, attestation)])
                 .await
-                .unwrap();
-            validator_store
-                .check_and_insert_attestations(vec![(attestation, pubkey)])
                 .unwrap()
                 .pop()
                 .unwrap()
-                .0
         })
         .await
         .assert_slashable_message_should_sign(
             "double_vote_attestation",
             move |pubkey, validator_store| async move {
-                let mut attestation = double_vote_attestation();
+                let attestation = double_vote_attestation();
                 validator_store
-                    .sign_attestation(pubkey, 0, &mut attestation, current_epoch)
+                    .sign_attestations(vec![(pubkey, 0, attestation)])
                     .await
+                    .map(|_| ())
             },
             slashable_message_should_sign,
         )
         .await
         .assert_slashable_message_should_sign(
             "surrounding_attestation",
             move |pubkey, validator_store| async move {
-                let mut attestation = surrounding_attestation();
+                let attestation = surrounding_attestation();
                 validator_store
-                    .sign_attestation(pubkey, 0, &mut attestation, current_epoch)
+                    .sign_attestations(vec![(pubkey, 0, attestation)])
                     .await
+                    .map(|_| ())
             },
             slashable_message_should_sign,
         )
         .await
         .assert_slashable_message_should_sign(
             "surrounded_attestation",
             move |pubkey, validator_store| async move {
-                let mut attestation = surrounded_attestation();
+                let attestation = surrounded_attestation();
                 validator_store
-                    .sign_attestation(pubkey, 0, &mut attestation, current_epoch)
+                    .sign_attestations(vec![(pubkey, 0, attestation)])
                     .await
+                    .map(|_| ())
             },
             slashable_message_should_sign,
         )

validator_client/http_api/src/tests/keystores.rs

@@ -1099,20 +1099,15 @@ async fn generic_migration_test(
         check_keystore_import_response(&import_res, all_imported(keystores.len()));
 
         // Sign attestations on VC1.
-        for (validator_index, mut attestation) in first_vc_attestations {
+        for (validator_index, attestation) in first_vc_attestations {
             let public_key = keystore_pubkey(&keystores[validator_index]);
-            let current_epoch = attestation.data().target.epoch;
-            tester1
-                .validator_store
-                .sign_attestation(public_key, 0, &mut attestation, current_epoch)
-                .await
-                .unwrap();
             let safe_attestations = tester1
                 .validator_store
-                .check_and_insert_attestations(vec![(attestation.clone(), public_key)])
+                .sign_attestations(vec![(public_key, 0, attestation.clone())])
+                .await
                 .unwrap();
             assert_eq!(safe_attestations.len(), 1);
-            assert_eq!(safe_attestations, vec![(attestation, public_key)]);
+            assert_eq!(safe_attestations, vec![attestation]);
         }
 
         // Delete the selected keys from VC1.
@@ -1184,27 +1179,24 @@ async fn generic_migration_test(
         check_keystore_import_response(&import_res, all_imported(import_indices.len()));
 
         // Sign attestations on the second VC.
-        for (validator_index, mut attestation, should_succeed) in second_vc_attestations {
+        for (validator_index, attestation, should_succeed) in second_vc_attestations {
             let public_key = keystore_pubkey(&keystores[validator_index]);
-            let current_epoch = attestation.data().target.epoch;
-            if tester2
-                .validator_store
-                .sign_attestation(public_key, 0, &mut attestation, current_epoch)
-                .await
-                .is_err()
-            {
-                // Doppelganger protected.
-                assert!(!should_succeed);
-                continue;
-            }
-            let safe_attestations = tester2
+            let result = tester2
                 .validator_store
-                .check_and_insert_attestations(vec![(attestation.clone(), public_key)])
-                .unwrap();
-            if should_succeed {
-                assert_eq!(safe_attestations[0], (attestation, public_key));
-            } else {
-                assert!(safe_attestations.is_empty());
+                .sign_attestations(vec![(public_key, 0, attestation.clone())])
+                .await;
+            match result {
+                Ok(safe_attestations) => {
+                    if should_succeed {
+                        assert_eq!(safe_attestations, vec![attestation]);
+                    } else {
+                        assert!(safe_attestations.is_empty());
+                    }
+                }
+                Err(_) => {
+                    // Doppelganger protected or other error.
+                    assert!(!should_succeed);
+                }
             }
         }
     })
@@ -1330,10 +1322,10 @@ async fn delete_concurrent_with_signing() {
 
         let handle = handle.spawn(async move {
             for j in 0..num_attestations {
-                let mut att = make_attestation(j, j + 1);
+                let att = make_attestation(j, j + 1);
                 for public_key in thread_pubkeys.iter() {
                     let _ = validator_store
-                        .sign_attestation(*public_key, 0, &mut att, Epoch::new(j + 1))
+                        .sign_attestations(vec![(*public_key, 0, att.clone())])
                         .await;
                 }
             }

validator_client/lighthouse_validator_store/src/lib.rs

@@ -750,8 +750,8 @@ impl<T: SlotClock + 'static, E: EthSpec> ValidatorStore for LighthouseValidatorS
 
     async fn sign_attestations(
         &self,
-        mut attestations: Vec<(PublicKeyBytes, usize, &mut Attestation<Self::E>)>,
-    ) -> Result<(), Error> {
+        mut attestations: Vec<(PublicKeyBytes, usize, Attestation<Self::E>)>,
+    ) -> Result<Vec<Attestation<E>>, Error> {
         // Sign all attestations concurrently.
         let signing_futures = attestations.iter_mut().map(|(pubkey, validator_committee_index, attestation)| {
             let pubkey = *pubkey;
@@ -764,35 +764,43 @@ impl<T: SlotClock + 'static, E: EthSpec> ValidatorStore for LighthouseValidatorS
                         attestation,
                     )
                     .await
-                    .map(|_| (pubkey, validator_committee_index))
+                    .map(|_| pubkey)
             }
         });
 
         // Execute all signing in parallel.
         let results: Vec<_> = join_all(signing_futures).await;
 
-        // Log errors but don't fail the entire batch.
-        for result in results {
-            if let Err(e) = result {
-                match e {
-                    ValidatorStoreError::UnknownPubkey(pubkey) => {
-                        warn!(
-                            info = "a validator may have recently been removed from this VC",
-                            ?pubkey,
-                            "Missing pubkey for attestation"
-                        );
-                    }
-                    e => {
-                        crit!(
-                            error = ?e,
-                            "Failed to sign attestation"
-                        );
-                    }
+        // Collect successfully signed attestations and log errors.
+        let mut signed_attestations = Vec::new();
+        for (result, (pubkey, _, attestation)) in results.into_iter().zip(attestations.into_iter()) {
+            match result {
+                Ok(_) => {
+                    signed_attestations.push((attestation, pubkey));
+                }
+                Err(ValidatorStoreError::UnknownPubkey(pubkey)) => {
+                    warn!(
+                        info = "a validator may have recently been removed from this VC",
+                        ?pubkey,
+                        "Missing pubkey for attestation"
+                    );
+                }
+                Err(e) => {
+                    crit!(
+                        error = ?e,
+                        "Failed to sign attestation"
+                    );
                 }
             }
         }
 
-        Ok(())
+        if signed_attestations.is_empty() {
+            return Ok(Vec::new());
+        }
+
+        // Check slashing protection and insert into database.
+        let safe_attestations = self.check_and_insert_attestations(signed_attestations)?;
+        Ok(safe_attestations.into_iter().map(|(a, _)| a).collect())
     }
 
     #[instrument(skip_all)]

validator_client/validator_services/src/attestation_service.rs

@@ -445,48 +445,21 @@ impl<S: ValidatorStore + 'static, T: SlotClock + 'static> AttestationService<S,
             return Ok(());
         }
 
-        // Sign all attestations concurrently.
-        let attestations_refs: Vec<_> = attestations_to_sign
-            .iter_mut()
-            .map(|(pubkey, index, attestation)| (*pubkey, *index, attestation))
-            .collect();
-
-        self.validator_store
-            .sign_attestations(attestations_refs)
+        // Sign and check all attestations (includes slashing protection).
+        let safe_attestations = self
+            .validator_store
+            .sign_attestations(attestations_to_sign)
             .await
             .map_err(|e| format!("Failed to sign attestations: {e:?}"))?;
 
-        // Collect signed attestations.
-        let signed_attestations: Vec<_> = attestations_to_sign
-            .into_iter()
-            .map(|(pubkey, _, attestation)| (attestation, pubkey))
-            .collect();
-
-        if signed_attestations.is_empty() {
+        if safe_attestations.is_empty() {
             warn!("No attestations were published");
             return Ok(());
         }
         let fork_name = self
             .chain_spec
             .fork_name_at_slot::<S::E>(attestation_data.slot);
 
-        // Check slashing protection in a blocking thread (this is I/O bound).
-        let service = self.clone();
-        let safe_attestations = self
-            .inner
-            .executor
-            .spawn_blocking_handle(
-                move || {
-                    service
-                        .validator_store
-                        .check_and_insert_attestations(signed_attestations)
-                },
-                "check_and_insert_attestations",
-            )
-            .ok_or("shutting down")?
-            .await
-            .map_err(|e| format!("thread error checking slashability: {e:?}"))?
-            .map_err(|e| format!("error checking slashability: {e:?}"))?;
         let safe_attestations = &safe_attestations;
 
         // Post the attestations to the BN.
@@ -502,7 +475,7 @@ impl<S: ValidatorStore + 'static, T: SlotClock + 'static> AttestationService<S,
                 let single_attestations = safe_attestations
                     .iter()
                     .zip(validator_indices_ref)
-                    .filter_map(|((a, _), i)| {
+                    .filter_map(|(a, i)| {
                         match a.to_single_attestation_with_attester_index(*i) {
                             Ok(a) => Some(a),
                             Err(e) => {

validator_client/validator_store/src/lib.rs

@@ -103,10 +103,11 @@ pub trait ValidatorStore: Send + Sync {
         current_slot: Slot,
     ) -> impl Future<Output = Result<SignedBlock<Self::E>, Error<Self::Error>>> + Send;
 
+    #[allow(clippy::type_complexity)]
     fn sign_attestations(
         &self,
-        attestations: Vec<(PublicKeyBytes, usize, &mut Attestation<Self::E>)>,
-    ) -> impl Future<Output = Result<(), Error<Self::Error>>> + Send;
+        attestations: Vec<(PublicKeyBytes, usize, Attestation<Self::E>)>,
+    ) -> impl Future<Output = Result<Vec<Attestation<Self::E>>, Error<Self::Error>>> + Send;
 
     fn sign_attestation_no_checks(
         &self,