Use an async Mutex for the Monero distribution

Enables safe async/thread-safe usage.
This commit is contained in:
Luke Parker 2023-03-12 04:13:43 -04:00
parent f5fa6f020d
commit 14203bbb46
No known key found for this signature in database
3 changed files with 46 additions and 32 deletions

1
Cargo.lock generated
View file

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

View file

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

View file

@ -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
// it returns
// TODO: Update this when scanning a block, as possible.
// 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 = {
let distribution = DISTRIBUTION.lock().unwrap();
distribution.len()
};
if distribution_len <= height {
let extension = rpc.get_output_distribution(distribution_len, height).await?;
DISTRIBUTION.lock().unwrap().extend(extension);
if distribution.len() <= height {
let extension = rpc.get_output_distribution(distribution.len(), height).await?;
distribution.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 per_second;
{
let mut distribution = DISTRIBUTION.lock().unwrap();
// 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
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 high = distribution[distribution.len() - 1];
let 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 =
select_n(rng, rpc, height, high, per_second, &real, &mut used, inputs.len() * decoy_count)
.await?;
let mut decoys = select_n(
rng,
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())
.await?,
select_n(
rng,
rpc,
&distribution,
height,
high,
per_second,
&[],
&mut used,
ring_len - ring.len(),
)
.await?,
);
ring.sort_by(|a, b| a.0.cmp(&b.0));
}