Add functions for multithreaded validation to the Rust side.

Cargo.toml

@@ -135,3 +135,4 @@ zstd = "0.13.2"
 blocking-threadpool = "1.0.1"
 libfuzzer-sys = "0.4"
 wasm-bindgen = "0.2.92"
+parking_lot = "0.12.3"

crates/chia-bls/Cargo.toml

@@ -27,6 +27,7 @@ thiserror = { workspace = true }
 pyo3 = { workspace = true, features = ["multiple-pymethods"], optional = true }
 arbitrary = { workspace = true, optional = true }
 lru = { workspace = true }
+parking_lot = { workspace = true }
 
 [dev-dependencies]
 rand = { workspace = true }

crates/chia-bls/benches/cache.rs

@@ -16,55 +16,50 @@ fn cache_benchmark(c: &mut Criterion) {
 
     let mut agg_sig = Signature::default();
     for i in 0..1000 {
-        let derived = sk.derive_hardened(i as u32);
+        let derived = sk.derive_hardened(i);
         let pk = derived.public_key();
         let sig = sign(&derived, msg);
         agg_sig.aggregate(&sig);
         pks.push(pk);
     }
 
-    let mut bls_cache = BlsCache::default();
+    let bls_cache = BlsCache::default();
 
     c.bench_function("bls_cache.aggregate_verify, 0% cache hits", |b| {
-        let mut cache = bls_cache.clone();
         b.iter(|| {
-            assert!(cache.aggregate_verify(&pks, [&msg].iter().cycle(), &agg_sig));
+            assert!(bls_cache.aggregate_verify(pks.iter().zip([&msg].iter().cycle()), &agg_sig));
         });
     });
 
     // populate 10% of keys
-    bls_cache.aggregate_verify(&pks[0..100], [&msg].iter().cycle(), &agg_sig);
+    bls_cache.aggregate_verify(pks[0..100].iter().zip([&msg].iter().cycle()), &agg_sig);
     c.bench_function("bls_cache.aggregate_verify, 10% cache hits", |b| {
-        let mut cache = bls_cache.clone();
         b.iter(|| {
-            assert!(cache.aggregate_verify(&pks, [&msg].iter().cycle(), &agg_sig));
+            assert!(bls_cache.aggregate_verify(pks.iter().zip([&msg].iter().cycle()), &agg_sig));
         });
     });
 
     // populate another 10% of keys
-    bls_cache.aggregate_verify(&pks[100..200], [&msg].iter().cycle(), &agg_sig);
+    bls_cache.aggregate_verify(pks[100..200].iter().zip([&msg].iter().cycle()), &agg_sig);
     c.bench_function("bls_cache.aggregate_verify, 20% cache hits", |b| {
-        let mut cache = bls_cache.clone();
         b.iter(|| {
-            assert!(cache.aggregate_verify(&pks, [&msg].iter().cycle(), &agg_sig));
+            assert!(bls_cache.aggregate_verify(pks.iter().zip([&msg].iter().cycle()), &agg_sig));
         });
     });
 
     // populate another 30% of keys
-    bls_cache.aggregate_verify(&pks[200..500], [&msg].iter().cycle(), &agg_sig);
+    bls_cache.aggregate_verify(pks[200..500].iter().zip([&msg].iter().cycle()), &agg_sig);
     c.bench_function("bls_cache.aggregate_verify, 50% cache hits", |b| {
-        let mut cache = bls_cache.clone();
         b.iter(|| {
-            assert!(cache.aggregate_verify(&pks, [&msg].iter().cycle(), &agg_sig));
+            assert!(bls_cache.aggregate_verify(pks.iter().zip([&msg].iter().cycle()), &agg_sig));
         });
     });
 
     // populate all other keys
-    bls_cache.aggregate_verify(&pks[500..1000], [&msg].iter().cycle(), &agg_sig);
+    bls_cache.aggregate_verify(pks[500..1000].iter().zip([&msg].iter().cycle()), &agg_sig);
     c.bench_function("bls_cache.aggregate_verify, 100% cache hits", |b| {
-        let mut cache = bls_cache.clone();
         b.iter(|| {
-            assert!(cache.aggregate_verify(&pks, [&msg].iter().cycle(), &agg_sig));
+            assert!(bls_cache.aggregate_verify(pks.iter().zip([&msg].iter().cycle()), &agg_sig));
         });
     });
 

crates/chia-bls/src/bls_cache.rs

@@ -1,11 +1,10 @@
-use std::borrow::Borrow;
-use std::num::NonZeroUsize;
-
-use lru::LruCache;
-use sha2::{Digest, Sha256};
-
 use crate::{aggregate_verify_gt, hash_to_g2};
 use crate::{GTElement, PublicKey, Signature};
+use lru::LruCache;
+use parking_lot::Mutex;
+use sha2::{Digest, Sha256};
+use std::borrow::Borrow;
+use std::num::NonZeroUsize;
 
 /// This is a cache of pairings of public keys and their corresponding message.
 /// It accelerates aggregate verification when some public keys have already
@@ -17,10 +16,10 @@ use crate::{GTElement, PublicKey, Signature};
 /// aggregate_verify() primitive is faster. When long-syncing, that's
 /// preferable.
 #[cfg_attr(feature = "py-bindings", pyo3::pyclass(name = "BLSCache"))]
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub struct BlsCache {
     // sha256(pubkey + message) -> GTElement
-    cache: LruCache<[u8; 32], GTElement>,
+    cache: Mutex<LruCache<[u8; 32], GTElement>>,
 }
 
 impl Default for BlsCache {
@@ -32,47 +31,46 @@ impl Default for BlsCache {
 impl BlsCache {
     pub fn new(cache_size: NonZeroUsize) -> Self {
         Self {
-            cache: LruCache::new(cache_size),
+            cache: Mutex::new(LruCache::new(cache_size)),
         }
     }
 
     pub fn len(&self) -> usize {
-        self.cache.len()
+        self.cache.lock().len()
     }
 
     pub fn is_empty(&self) -> bool {
-        self.cache.is_empty()
+        self.cache.lock().is_empty()
     }
 
-    pub fn aggregate_verify(
-        &mut self,
-        pks: impl IntoIterator<Item = impl Borrow<PublicKey>>,
-        msgs: impl IntoIterator<Item = impl AsRef<[u8]>>,
+    pub fn aggregate_verify<Pk: Borrow<PublicKey>, Msg: AsRef<[u8]>>(
+        &self,
+        pks_msgs: impl IntoIterator<Item = (Pk, Msg)>,
         sig: &Signature,
     ) -> bool {
-        let iter = pks.into_iter().zip(msgs).map(|(pk, msg)| -> GTElement {
+        let iter = pks_msgs.into_iter().map(|(pk, msg)| -> GTElement {
             // Hash pubkey + message
             let mut hasher = Sha256::new();
             hasher.update(pk.borrow().to_bytes());
             hasher.update(msg.as_ref());
             let hash: [u8; 32] = hasher.finalize().into();
 
             // If the pairing is in the cache, we don't need to recalculate it.
-            if let Some(pairing) = self.cache.get(&hash).cloned() {
+            if let Some(pairing) = self.cache.lock().get(&hash).cloned() {
                 return pairing;
             }
 
             // Otherwise, we need to calculate the pairing and add it to the cache.
             let mut aug_msg = pk.borrow().to_bytes().to_vec();
-            aug_msg.extend_from_slice(msg.as_ref());
+            aug_msg.extend(msg.as_ref());
             let aug_hash = hash_to_g2(&aug_msg);
 
             let mut hasher = Sha256::new();
             hasher.update(&aug_msg);
             let hash: [u8; 32] = hasher.finalize().into();
 
             let pairing = aug_hash.pair(pk.borrow());
-            self.cache.put(hash, pairing.clone());
+            self.cache.lock().put(hash, pairing.clone());
             pairing
         });
 
@@ -83,7 +81,6 @@ impl BlsCache {
 #[cfg(feature = "py-bindings")]
 use pyo3::{
     exceptions::PyValueError,
-    pybacked::PyBackedBytes,
     types::{PyAnyMethods, PyList},
     Bound, PyObject, PyResult,
 };
@@ -122,9 +119,9 @@ impl BlsCache {
         let msgs = msgs
             .iter()?
             .map(|item| item?.extract())
-            .collect::<PyResult<Vec<PyBackedBytes>>>()?;
+            .collect::<PyResult<Vec<Vec<u8>>>>()?;
 
-        Ok(self.aggregate_verify(pks, msgs, sig))
+        Ok(self.aggregate_verify(pks.into_iter().zip(msgs), sig))
     }
 
     #[pyo3(name = "len")]
@@ -137,27 +134,22 @@ impl BlsCache {
         use pyo3::prelude::*;
         use pyo3::types::PyBytes;
         let ret = PyList::empty_bound(py);
-        for (key, value) in &self.cache {
-            ret.append((
-                PyBytes::new_bound(py, key),
-                PyBytes::new_bound(py, &value.to_bytes()),
-            ))?;
+        let cache = self.cache.lock();
+        for (key, value) in cache.iter() {
+            ret.append((PyBytes::new_bound(py, key), value.clone().into_py(py)))?;
         }
         Ok(ret.into())
     }
 
     #[pyo3(name = "update")]
     pub fn py_update(&mut self, other: &Bound<'_, PyList>) -> PyResult<()> {
+        let mut cache = self.cache.lock();
         for item in other.borrow().iter()? {
-            let (key, value): (Vec<u8>, Vec<u8>) = item?.extract()?;
-            self.cache.put(
+            let (key, value): (Vec<u8>, GTElement) = item?.extract()?;
+            cache.put(
                 key.try_into()
                     .map_err(|_| PyValueError::new_err("invalid key"))?,
-                GTElement::from_bytes(
-                    (&value[..])
-                        .try_into()
-                        .map_err(|_| PyValueError::new_err("invalid GTElement"))?,
-                ),
+                value,
             );
         }
         Ok(())
@@ -173,7 +165,7 @@ pub mod tests {
 
     #[test]
     fn test_aggregate_verify() {
-        let mut bls_cache = BlsCache::default();
+        let bls_cache = BlsCache::default();
 
         let sk = SecretKey::from_seed(&[0; 32]);
         let pk = sk.public_key();
@@ -187,17 +179,17 @@ pub mod tests {
         assert!(bls_cache.is_empty());
 
         // Verify the signature and add to the cache.
-        assert!(bls_cache.aggregate_verify(pk_list, msg_list, &sig));
+        assert!(bls_cache.aggregate_verify(pk_list.into_iter().zip(msg_list), &sig));
         assert_eq!(bls_cache.len(), 1);
 
         // Now that it's cached, it shouldn't cache it again.
-        assert!(bls_cache.aggregate_verify(pk_list, msg_list, &sig));
+        assert!(bls_cache.aggregate_verify(pk_list.into_iter().zip(msg_list), &sig));
         assert_eq!(bls_cache.len(), 1);
     }
 
     #[test]
     fn test_cache() {
-        let mut bls_cache = BlsCache::default();
+        let bls_cache = BlsCache::default();
 
         let sk1 = SecretKey::from_seed(&[0; 32]);
         let pk1 = sk1.public_key();
@@ -211,7 +203,7 @@ pub mod tests {
         assert!(bls_cache.is_empty());
 
         // Add the first signature to cache.
-        assert!(bls_cache.aggregate_verify(&pk_list, &msg_list, &agg_sig));
+        assert!(bls_cache.aggregate_verify(pk_list.iter().zip(msg_list.iter()), &agg_sig));
         assert_eq!(bls_cache.len(), 1);
 
         // Try with the first key message pair in the cache but not the second.
@@ -223,7 +215,7 @@ pub mod tests {
         pk_list.push(pk2);
         msg_list.push(msg2);
 
-        assert!(bls_cache.aggregate_verify(&pk_list, &msg_list, &agg_sig));
+        assert!(bls_cache.aggregate_verify(pk_list.iter().zip(msg_list.iter()), &agg_sig));
         assert_eq!(bls_cache.len(), 2);
 
         // Try reusing a public key.
@@ -234,14 +226,14 @@ pub mod tests {
         msg_list.push(msg3);
 
         // Verify this signature and add to the cache as well (since it's still a different aggregate).
-        assert!(bls_cache.aggregate_verify(pk_list, msg_list, &agg_sig));
+        assert!(bls_cache.aggregate_verify(pk_list.iter().zip(msg_list), &agg_sig));
         assert_eq!(bls_cache.len(), 3);
     }
 
     #[test]
     fn test_cache_limit() {
         // The cache is limited to only 3 items.
-        let mut bls_cache = BlsCache::new(NonZeroUsize::new(3).unwrap());
+        let bls_cache = BlsCache::new(NonZeroUsize::new(3).unwrap());
 
         // Before we cache anything, it should be empty.
         assert!(bls_cache.is_empty());
@@ -257,11 +249,11 @@ pub mod tests {
             let msg_list = [msg];
 
             // Add to cache by validating them one at a time.
-            assert!(bls_cache.aggregate_verify(pk_list.iter(), msg_list.iter(), &sig));
+            assert!(bls_cache.aggregate_verify(pk_list.into_iter().zip(msg_list), &sig));
         }
 
         // The cache should be full now.
-        assert_eq!(bls_cache.cache.len(), 3);
+        assert_eq!(bls_cache.len(), 3);
 
         // Recreate first key.
         let sk = SecretKey::from_seed(&[1; 32]);
@@ -275,16 +267,16 @@ pub mod tests {
         let hash: [u8; 32] = hasher.finalize().into();
 
         // The first key should have been removed, since it's the oldest that's been accessed.
-        assert!(!bls_cache.cache.contains(&hash));
+        assert!(!bls_cache.cache.lock().contains(&hash));
     }
 
     #[test]
     fn test_empty_sig() {
-        let mut bls_cache = BlsCache::default();
+        let bls_cache = BlsCache::default();
 
         let pks: [&PublicKey; 0] = [];
         let msgs: [&[u8]; 0] = [];
 
-        assert!(bls_cache.aggregate_verify(pks, msgs, &Signature::default()));
+        assert!(bls_cache.aggregate_verify(pks.into_iter().zip(msgs), &Signature::default()));
     }
 }