perf(crypto): CRP-1827 Add a cache for G2Prepared (#7967)

This improves the benchmark results of the verify_multi_sig_individual
and verify_multi_sig_combined benchmarks by approximately 5%.

This also corrects an error introduced when the G2 point cache was
introduced; it reported its cache statistics into the BLS signature
verification cache instead of the metrics for the point cache.

Author: randombit

rs/crypto/internal/crypto_lib/bls12_381/type/src/cache.rs

@@ -1,4 +1,4 @@
-use crate::G2Affine;
+use crate::{G2Affine, G2Prepared};
 use cached::{Cached, SizedCache};
 use parking_lot::Mutex;
 use std::sync::LazyLock;
@@ -86,3 +86,71 @@ impl G2PublicKeyCache {
         G2PublicKeyCacheStatistics::new(cache_size, hits, misses)
     }
 }
+
+#[derive(Copy, Clone, Eq, PartialEq, Debug)]
+pub struct G2PreparedCacheStatistics {
+    pub size: usize,
+    pub hits: u64,
+    pub misses: u64,
+}
+
+impl G2PreparedCacheStatistics {
+    fn new(size: usize, hits: u64, misses: u64) -> Self {
+        Self { size, hits, misses }
+    }
+}
+
+/// A cache for G2 Public Keys
+pub(crate) struct G2PreparedCache {
+    cache: Mutex<SizedCache<[u8; G2Affine::BYTES], G2Prepared>>,
+}
+
+static GLOBAL_G2PREP_CACHE: LazyLock<G2PreparedCache> =
+    LazyLock::new(|| G2PreparedCache::new(G2PreparedCache::SIZE_OF_GLOBAL_CACHE));
+
+impl G2PreparedCache {
+    /// Specify the size of the global cache used for public keys
+    ///
+    /// This cache is kept small because G2Prepared is quite large,
+    /// just under 19 KiB.
+    pub const SIZE_OF_GLOBAL_CACHE: usize = 50;
+
+    /// Create a new signature cache with the specified maximum size
+    fn new(max_size: usize) -> Self {
+        let cache = Mutex::<SizedCache<[u8; G2Affine::BYTES], G2Prepared>>::new(
+            SizedCache::with_size(max_size),
+        );
+        Self { cache }
+    }
+
+    /// Return a reference to the global signature cache
+    pub(crate) fn global() -> &'static Self {
+        &GLOBAL_G2PREP_CACHE
+    }
+
+    /// Check the cache for an already prepared G2 element
+    pub(crate) fn get(&self, bytes: &[u8; G2Affine::BYTES]) -> Option<G2Prepared> {
+        let mut cache = self.cache.lock();
+        cache.cache_get(bytes).cloned()
+    }
+
+    /// Insert a new G2Prepared into the cache
+    pub(crate) fn insert(&self, bytes: [u8; G2Affine::BYTES], prep: G2Prepared) {
+        let mut cache = self.cache.lock();
+        cache.cache_set(bytes, prep);
+    }
+
+    /// Return statistics about the cache
+    ///
+    /// Returns the size of the cache, the number of cache hits, and
+    /// the number of cache misses
+    pub(crate) fn cache_statistics(&self) -> G2PreparedCacheStatistics {
+        let cache = self.cache.lock();
+
+        let cache_size = cache.cache_size();
+        let hits = cache.cache_hits().unwrap_or(0);
+        let misses = cache.cache_misses().unwrap_or(0);
+
+        G2PreparedCacheStatistics::new(cache_size, hits, misses)
+    }
+}

rs/crypto/internal/crypto_lib/bls12_381/type/src/lib.rs

@@ -2443,11 +2443,23 @@ impl G2Prepared {
     pub fn neg_generator() -> &'static Self {
         &G2PREPARED_NEG_G
     }
+
+    /// Return statistics related to the G2Prepared cache
+    pub fn cache_statistics() -> crate::cache::G2PreparedCacheStatistics {
+        crate::cache::G2PreparedCache::global().cache_statistics()
+    }
 }
 
 impl From<&G2Affine> for G2Prepared {
     fn from(v: &G2Affine) -> Self {
-        Self::new((*v.inner()).into())
+        let bytes = v.serialize();
+        if let Some(prep) = crate::cache::G2PreparedCache::global().get(&bytes) {
+            prep
+        } else {
+            let prep = Self::new((*v.inner()).into());
+            crate::cache::G2PreparedCache::global().insert(bytes, prep.clone());
+            prep
+        }
     }
 }
 

rs/crypto/internal/logmon/src/metrics.rs

@@ -1,5 +1,6 @@
 //! Metrics exported by crypto
 
+mod bls12_381_g2_prep_cache;
 mod bls12_381_point_cache;
 mod bls12_381_sig_cache;
 
@@ -283,7 +284,7 @@ impl CryptoMetrics {
         }
     }
 
-    /// Observes the cache statistics for the verification of threshold BLS12-381 signatures.
+    /// Observes the cache statistics for parsing of BLS12-381 points
     pub fn observe_bls12_381_point_cache_stats(&self, size: usize, hits: u64, misses: u64) {
         if let Some(metrics) = &self.metrics {
             let m = &metrics.crypto_bls12_381_point_cache_metrics;
@@ -299,6 +300,22 @@ impl CryptoMetrics {
         }
     }
 
+    /// Observes the cache statistics for parsing of BLS12-381 G2Prepared
+    pub fn observe_bls12_381_g2_prep_cache_stats(&self, size: usize, hits: u64, misses: u64) {
+        if let Some(metrics) = &self.metrics {
+            let m = &metrics.crypto_bls12_381_g2_prep_cache_metrics;
+            m.cache_size.set(size as i64);
+
+            let prev_hits = m.cache_hits.get();
+            debug_assert!(prev_hits <= hits);
+            m.cache_hits.inc_by(hits - prev_hits);
+
+            let prev_misses = m.cache_misses.get();
+            debug_assert!(prev_misses <= misses);
+            m.cache_misses.inc_by(misses - prev_misses);
+        }
+    }
+
     /// Observes the minimum epoch in active NI-DKG transcripts
     pub fn observe_minimum_epoch_in_active_nidkg_transcripts(&self, epoch: u32) {
         if let Some(metrics) = &self.metrics {
@@ -609,6 +626,9 @@ struct Metrics {
     /// Metrics for the cache of successfully decoded BLS12-381 points
     pub crypto_bls12_381_point_cache_metrics: bls12_381_point_cache::Metrics,
 
+    /// Metrics for the cache of successfully created BLS12-381 G2Prepared
+    pub crypto_bls12_381_g2_prep_cache_metrics: bls12_381_g2_prep_cache::Metrics,
+
     /// Gauge for the minimum epoch in active NI-DKG transcripts.
     observe_minimum_epoch_in_active_nidkg_transcripts: Gauge,
 
@@ -774,6 +794,18 @@ impl Metrics {
                     "crypto_bls12_381_point_cache_misses",
                 "Number of cache misses for successfully decoded BLS12-381 points"),
             },
+            crypto_bls12_381_g2_prep_cache_metrics: bls12_381_g2_prep_cache::Metrics {
+                cache_size: r.int_gauge(
+                    "crypto_bls12_381_g2_prep_cache_size",
+                    "Size of cache of BLS12-381 G2Prepared",
+                ),
+                cache_hits: r.int_counter(
+                    "crypto_bls12_381_g2_prep_cache_hits",
+                "Number of cache hits of BLS12-381 G2Prepared cache"),
+                cache_misses: r.int_counter(
+                    "crypto_bls12_381_g2_prep_cache_misses",
+                "Number of cache misses of BLS12-381 G2Prepared cache"),
+            },
             observe_minimum_epoch_in_active_nidkg_transcripts: r.gauge(
                 "crypto_minimum_epoch_in_active_nidkg_transcripts",
                 "Minimum epoch in active NI-DKG transcripts"

rs/crypto/internal/logmon/src/metrics/bls12_381_g2_prep_cache.rs

@@ -0,0 +1,14 @@
+use prometheus::{IntCounter, IntGauge};
+
+/// Metrics for the cache of successfully created BLS12-381 G2Prepared
+pub struct Metrics {
+    /// [`IntGauge`] for tracking the size of the cache. The size is expected
+    /// to increase with the time and remain at its max value until the process ends.
+    /// But if it happens that the cache size decreases, we want to know about that,
+    /// therefore we use an [`IntGauge`] and not a [`Counter`].
+    pub cache_size: IntGauge,
+    /// [`Counter`] for tracking the cache hits.
+    pub cache_hits: IntCounter,
+    /// [`Counter`] for tracking the cache misses.
+    pub cache_misses: IntCounter,
+}

rs/crypto/src/sign/threshold_sig/ni_dkg.rs

@@ -175,7 +175,11 @@ impl<C: CryptoServiceProvider> NiDkgAlgorithm for CryptoComponentImpl<C> {
         // locking congestion or similar, we should be able to notice that.
         let stats = ic_crypto_internal_bls12_381_type::G2Affine::deserialize_cached_statistics();
         self.metrics
-            .observe_bls12_381_sig_cache_stats(stats.size, stats.hits, stats.misses);
+            .observe_bls12_381_point_cache_stats(stats.size, stats.hits, stats.misses);
+
+        let stats = ic_crypto_internal_bls12_381_type::G2Prepared::cache_statistics();
+        self.metrics
+            .observe_bls12_381_g2_prep_cache_stats(stats.size, stats.hits, stats.misses);
 
         self.metrics.observe_parameter_size(
             MetricsDomain::NiDkgAlgorithm,