Use an async Mutex for the Monero distribution

Enables safe async/thread-safe usage.

Cargo.lock

@@ -5027,6 +5027,7 @@ dependencies = [
  "digest_auth",
  "dleq",
  "flexible-transcript",
+ "futures",
  "group",
  "hex",
  "hex-literal",

coins/monero/Cargo.toml

@@ -12,6 +12,8 @@ all-features = true
 rustdoc-args = ["--cfg", "docsrs"]
 
 [dependencies]
+futures = "0.3"
+
 hex-literal = "0.3"
 lazy_static = "1"
 thiserror = "1"

coins/monero/src/wallet/decoys.rs

@@ -1,4 +1,6 @@
-use std::{sync::Mutex, collections::HashSet};
+use std::collections::HashSet;
+
+use futures::lock::{Mutex, MutexGuard};
 
 use lazy_static::lazy_static;
 
@@ -24,16 +26,15 @@ const TIP_APPLICATION: f64 = (LOCK_WINDOW * BLOCK_TIME) as f64;
 lazy_static! {
   static ref GAMMA: Gamma<f64> = Gamma::new(19.28, 1.0 / 1.61).unwrap();
   // TODO: Expose an API to reset this in case a reorg occurs/the RPC fails/returns garbage
-  // TODO: This is not currently thread-safe. This needs to be a tokio Mutex held by select until
+  // TODO: Update this when scanning a block, as possible
-  // it returns
-  // TODO: Update this when scanning a block, as possible.
   static ref DISTRIBUTION: Mutex<Vec<u64>> = Mutex::new(Vec::with_capacity(3000000));
 }
 
 #[allow(clippy::too_many_arguments)]
-async fn select_n<R: RngCore + CryptoRng>(
+async fn select_n<'a, R: RngCore + CryptoRng>(
   rng: &mut R,
   rpc: &Rpc,
+  distribution: &MutexGuard<'a, Vec<u64>>,
   height: usize,
   high: u64,
   per_second: f64,
@@ -65,7 +66,6 @@ async fn select_n<R: RngCore + CryptoRng>(
 
       let o = (age * per_second) as u64;
       if o < high {
-        let distribution = DISTRIBUTION.lock().unwrap();
         let i = distribution.partition_point(|s| *s < (high - 1 - o));
         let prev = i.saturating_sub(1);
         let n = distribution[i] - distribution[prev];
@@ -140,6 +140,8 @@ impl Decoys {
     height: usize,
     inputs: &[SpendableOutput],
   ) -> Result<Vec<Decoys>, RpcError> {
+    let mut distribution = DISTRIBUTION.lock().await;
+
     let decoy_count = ring_len - 1;
 
     // Convert the inputs in question to the raw output data
@@ -150,29 +152,19 @@ impl Decoys {
       outputs.push((real[real.len() - 1], [input.key(), input.commitment().calculate()]));
     }
 
-    let distribution_len = {
+    if distribution.len() <= height {
-      let distribution = DISTRIBUTION.lock().unwrap();
+      let extension = rpc.get_output_distribution(distribution.len(), height).await?;
-      distribution.len()
+      distribution.extend(extension);
-    };
-    if distribution_len <= height {
-      let extension = rpc.get_output_distribution(distribution_len, height).await?;
-      DISTRIBUTION.lock().unwrap().extend(extension);
     }
+    // If asked to use an older height than previously asked, truncate to ensure accuracy
+    // Should never happen, yet risks desyncing if it did
+    distribution.truncate(height + 1); // height is inclusive, and 0 is a valid height
 
-    let high;
+    let high = distribution[distribution.len() - 1];
-    let per_second;
+    let per_second = {
-    {
+      let blocks = distribution.len().min(BLOCKS_PER_YEAR);
-      let mut distribution = DISTRIBUTION.lock().unwrap();
+      let outputs = high - distribution[distribution.len().saturating_sub(blocks + 1)];
-      // If asked to use an older height than previously asked, truncate to ensure accuracy
+      (outputs as f64) / ((blocks * BLOCK_TIME) as f64)
-      // Should never happen, yet risks desyncing if it did
-      distribution.truncate(height + 1); // height is inclusive, and 0 is a valid height
-
-      high = distribution[distribution.len() - 1];
-      per_second = {
-        let blocks = distribution.len().min(BLOCKS_PER_YEAR);
-        let outputs = high - distribution[distribution.len().saturating_sub(blocks + 1)];
-        (outputs as f64) / ((blocks * BLOCK_TIME) as f64)
-      };
     };
 
     let mut used = HashSet::<u64>::new();
@@ -188,9 +180,18 @@ impl Decoys {
     // Select all decoys for this transaction, assuming we generate a sane transaction
     // We should almost never naturally generate an insane transaction, hence why this doesn't
     // bother with an overage
-    let mut decoys =
+    let mut decoys = select_n(
-      select_n(rng, rpc, height, high, per_second, &real, &mut used, inputs.len() * decoy_count)
+      rng,
-        .await?;
+      rpc,
+      &distribution,
+      height,
+      high,
+      per_second,
+      &real,
+      &mut used,
+      inputs.len() * decoy_count,
+    )
+    .await?;
     real.zeroize();
 
     let mut res = Vec::with_capacity(inputs.len());
@@ -228,8 +229,18 @@ impl Decoys {
 
           // Select new outputs until we have a full sized ring again
           ring.extend(
-            select_n(rng, rpc, height, high, per_second, &[], &mut used, ring_len - ring.len())
+            select_n(
-              .await?,
+              rng,
+              rpc,
+              &distribution,
+              height,
+              high,
+              per_second,
+              &[],
+              &mut used,
+              ring_len - ring.len(),
+            )
+            .await?,
           );
           ring.sort_by(|a, b| a.0.cmp(&b.0));
         }