diff --git a/crates/blockchain/src/key_manager.rs b/crates/blockchain/src/key_manager.rs index 98adbb3e8..6cf38c24e 100644 --- a/crates/blockchain/src/key_manager.rs +++ b/crates/blockchain/src/key_manager.rs @@ -169,6 +169,35 @@ impl KeyManager { self.sign_with_attestation_key(validator_id, slot, &message_hash) } + /// Signs `attestation_data` with each of `validator_ids`' attestation keys, + /// spread over at most one thread per core, and returns one result per id + /// in `validator_ids` order. + /// + /// Every validator a node runs votes at the same instant, and one XMSS + /// signature takes tens of milliseconds, so signing them one after another + /// held a node's last vote back by most of a second, and the aggregators + /// waited on that last vote before their first proof covered the subnet. + /// The keys are independent and `sign` needs only `&self`, so the + /// signatures run concurrently instead. + pub fn sign_attestations( + &self, + validator_ids: &[u64], + attestation_data: &AttestationData, + ) -> Vec> { + let message_hash = attestation_data.hash_tree_root(); + let slot = attestation_data.slot as u32; + let max_threads = thread::available_parallelism().map_or(1, NonZeroUsize::get); + // One key per batch at minimum: a signature costs far more than the + // thread spawn it would amortize. + map_batched( + validator_ids, + max_threads, + 1, + "xmss-sign", + |&validator_id| self.sign_with_attestation_key(validator_id, slot, &message_hash), + ) + } + /// Signs a block root using the validator's proposal key. pub fn sign_block_root( &mut self, @@ -180,14 +209,14 @@ impl KeyManager { } fn sign_with_attestation_key( - &mut self, + &self, validator_id: u64, slot: u32, message: &H256, ) -> Result { let key_pair = self .keys - .get_mut(&validator_id) + .get(&validator_id) .ok_or(KeyManagerError::ValidatorKeyNotFound(validator_id))?; // A slot outside the key's range can never be signed, however long the @@ -285,34 +314,62 @@ fn prepare_key(validator_id: u64, role: KeyRole, key: &ValidatorSecretKey, slot: /// Runs `f` on every item, in batches of at least [`MIN_KEYS_PER_WARM_THREAD`] /// items spread over at most `max_threads` threads, and returns once all are -/// done. +/// done. See [`map_batched`]. +fn for_each_batched(items: &[T], max_threads: usize, f: impl Fn(&T) + Sync) { + map_batched(items, max_threads, MIN_KEYS_PER_WARM_THREAD, "xmss-warm", f); +} + +/// Maps `f` over `items`, in batches of at least `min_batch` items spread over +/// at most `max_threads` threads named `thread_name`, and returns the results +/// in `items` order once all are done. /// /// The calling thread takes the first batch, so a single batch spawns nothing. /// A batch whose thread cannot be spawned also runs on the calling thread, /// which costs latency instead of panicking the caller. -fn for_each_batched(items: &[T], max_threads: usize, f: impl Fn(&T) + Sync) { +fn map_batched( + items: &[T], + max_threads: usize, + min_batch: usize, + thread_name: &str, + f: impl Fn(&T) -> R + Sync, +) -> Vec { let batch_len = items .len() .div_ceil(max_threads.max(1)) - .max(MIN_KEYS_PER_WARM_THREAD); + .max(min_batch.max(1)); let mut batches = items.chunks(batch_len); let Some(first) = batches.next() else { - return; + return Vec::new(); }; let f = &f; + let run = move |batch: &[T]| batch.iter().map(f).collect::>(); thread::scope(|scope| { - let mut inline = vec![first]; - for batch in batches { - let spawned = thread::Builder::new() - .name("xmss-warm".into()) - .spawn_scoped(scope, move || batch.iter().for_each(f)); - if let Err(err) = spawned { - warn!(%err, "Failed to spawn an XMSS warm thread, warming its keys inline"); - inline.push(batch); + // Spawned before the first batch runs, so every batch is in flight + // while the calling thread works through its own. + let rest: Vec<_> = batches + .map(|batch| { + thread::Builder::new() + .name(thread_name.into()) + .spawn_scoped(scope, move || run(batch)) + .map_err(|err| { + warn!(%err, thread_name, "Failed to spawn a key thread, running its batch inline"); + batch + }) + }) + .collect(); + let mut results = run(first); + for batch in rest { + match batch { + Ok(handle) => results.extend( + handle + .join() + .unwrap_or_else(|panic| std::panic::resume_unwind(panic)), + ), + Err(batch) => results.extend(run(batch)), } } - inline.into_iter().flatten().for_each(f); - }); + results + }) } #[cfg(test)] @@ -320,6 +377,8 @@ mod tests { use std::collections::HashSet; use std::sync::Mutex; + use ethlambda_types::checkpoint::Checkpoint; + use super::*; #[test] @@ -332,7 +391,7 @@ mod tests { #[test] fn test_sign_attestation_validator_not_found() { let keys = HashMap::new(); - let mut key_manager = KeyManager::new(keys); + let key_manager = KeyManager::new(keys); let message = H256::default(); let result = key_manager.sign_with_attestation_key(123, 0, &message); @@ -443,6 +502,63 @@ mod tests { assert_eq!(threads_used(3 * batch, 64).len(), 3); } + #[test] + fn map_batched_returns_results_in_item_order() { + for len in [0, 1, 2, 7, 40] { + for max_threads in [0, 1, 3, 64] { + for min_batch in [0, 1, MIN_KEYS_PER_WARM_THREAD] { + let items: Vec = (0..len).collect(); + let doubled = map_batched(&items, max_threads, min_batch, "test", |&i| 2 * i); + let expected: Vec = items.iter().map(|i| 2 * i).collect(); + assert_eq!( + doubled, expected, + "len {len}, max_threads {max_threads}, min_batch {min_batch}" + ); + } + } + } + } + + #[test] + fn map_batched_with_unit_batches_uses_one_thread_per_item() { + let items: Vec = (0..5).collect(); + let threads: HashSet<_> = map_batched(&items, 64, 1, "test", |_| thread::current().id()) + .into_iter() + .collect(); + assert_eq!(threads.len(), items.len()); + } + + #[test] + fn sign_attestations_signs_each_id_in_order() { + let count = 5; + let key_manager = tiny_key_manager(count); + let data = AttestationData { + slot: 1, + head: Checkpoint::default(), + target: Checkpoint::default(), + source: Checkpoint::default(), + }; + let message = data.hash_tree_root(); + // An unknown id in the middle keeps its place and reports itself. + let ids = [3, 0, 99, 4, 1, 2]; + + let results = key_manager.sign_attestations(&ids, &data); + + assert_eq!(results.len(), ids.len()); + for (&id, result) in ids.iter().zip(&results) { + let Some(pair) = key_manager.keys.get(&id) else { + assert!( + matches!(result, Err(KeyManagerError::ValidatorKeyNotFound(missing)) if *missing == id), + "id {id}: {result:?}" + ); + continue; + }; + let signature = ValidatorSignature::from_bytes(result.as_ref().unwrap()).unwrap(); + let pubkey = pair.attestation_key.public_key(); + assert!(signature.is_valid(&pubkey, 1, &message), "id {id}"); + } + } + #[test] fn prepare_keys_for_leaves_every_key_signable() { // More keys than one batch holds, so the warm spans several threads. diff --git a/crates/blockchain/src/lib.rs b/crates/blockchain/src/lib.rs index 7692bc2bd..7c8cb5af1 100644 --- a/crates/blockchain/src/lib.rs +++ b/crates/blockchain/src/lib.rs @@ -1503,17 +1503,17 @@ impl BlockChainServer { // Produce attestation data once for all validators let attestation_data = store::produce_attestation_data(&self.store, slot); - // For each registered validator, produce and publish attestation - for validator_id in validator_ids { - // Sign the attestation - let Ok(signature) = self - .lean_mut() - .key_manager - .sign_attestation(validator_id, &attestation_data) - .inspect_err( - |err| error!(%slot, %validator_id, %err, "Failed to sign attestation"), - ) - else { + // Sign every validator's attestation at once, then deliver and publish + // each. Signing them in parallel is what keeps a node's last vote from + // trailing its first by the sum of every signature before it. + let signatures = self + .lean() + .key_manager + .sign_attestations(&validator_ids, &attestation_data); + for (validator_id, signature) in validator_ids.into_iter().zip(signatures) { + let Ok(signature) = signature.inspect_err( + |err| error!(%slot, %validator_id, %err, "Failed to sign attestation"), + ) else { continue; };