zec
/
librustzcash
mirror of https://github.com/zcash/librustzcash.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Merge pull request #1173 from zcash/1171-generalise-batch-runner 

zcash_client_backend: Extend `BatchRunner` to support full outputs
This commit is contained in:
str4d 2024-02-08 23:14:02 +00:00 committed by GitHub
parent 74c036e870 f6224e7e80
commit 121438eaef
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 135 additions and 57 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
zcash_client_backend/src/data_api/chain.rs
View File

@ -303,7 +303,7 @@ where
.get_sapling_nullifiers(NullifierQuery::Unspent)
.map_err(Error::Wallet)?;
let mut batch_runner = BatchRunner::<_, _, _, ()>::new(
let mut batch_runner = BatchRunner::<_, _, _, _, ()>::new(
100,
dfvks
.iter()

171
zcash_client_backend/src/scan.rs
View File

@ -8,36 +8,107 @@ use std::sync::{
};
use memuse::DynamicUsage;
use zcash_note_encryption::{batch, BatchDomain, Domain, ShieldedOutput, COMPACT_NOTE_SIZE};
use zcash_note_encryption::{
batch, BatchDomain, Domain, ShieldedOutput, COMPACT_NOTE_SIZE, ENC_CIPHERTEXT_SIZE,
};
use zcash_primitives::{block::BlockHash, transaction::TxId};
/// A decrypted note.
pub(crate) struct DecryptedNote<A, D: Domain> {
/// A decrypted transaction output.
pub(crate) struct DecryptedOutput<A, D: Domain, M> {
/// The tag corresponding to the incoming viewing key used to decrypt the note.
pub(crate) ivk_tag: A,
/// The recipient of the note.
pub(crate) recipient: D::Recipient,
/// The note!
pub(crate) note: D::Note,
/// The memo field, or `()` if this is a decrypted compact output.
pub(crate) memo: M,
}
impl<A, D: Domain> fmt::Debug for DecryptedNote<A, D>
impl<A, D: Domain, M> fmt::Debug for DecryptedOutput<A, D, M>
where
A: fmt::Debug,
D::IncomingViewingKey: fmt::Debug,
D::Recipient: fmt::Debug,
D::Note: fmt::Debug,
D::Memo: fmt::Debug,
M: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("DecryptedNote")
f.debug_struct("DecryptedOutput")
.field("ivk_tag", &self.ivk_tag)
.field("recipient", &self.recipient)
.field("note", &self.note)
.field("memo", &self.memo)
.finish()
}
}
/// A decryptor of transaction outputs.
pub(crate) trait Decryptor<D: BatchDomain, Output> {
type Memo;
// Once we reach MSRV 1.75.0, this can return `impl Iterator`.
fn batch_decrypt<A: Clone>(
tags: &[A],
ivks: &[D::IncomingViewingKey],
outputs: &[(D, Output)],
) -> Vec<Option<DecryptedOutput<A, D, Self::Memo>>>;
}
/// A decryptor of outputs as encoded in transactions.
pub(crate) struct FullDecryptor;
impl<D: BatchDomain, Output: ShieldedOutput<D, ENC_CIPHERTEXT_SIZE>> Decryptor<D, Output>
for FullDecryptor
{
type Memo = D::Memo;
fn batch_decrypt<A: Clone>(
tags: &[A],
ivks: &[D::IncomingViewingKey],
outputs: &[(D, Output)],
) -> Vec<Option<DecryptedOutput<A, D, Self::Memo>>> {
batch::try_note_decryption(ivks, outputs)
.into_iter()
.map(|res| {
res.map(|((note, recipient, memo), ivk_idx)| DecryptedOutput {
ivk_tag: tags[ivk_idx].clone(),
recipient,
note,
memo,
})
})
.collect()
}
}
/// A decryptor of outputs as encoded in compact blocks.
pub(crate) struct CompactDecryptor;
impl<D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>> Decryptor<D, Output>
for CompactDecryptor
{
type Memo = ();
fn batch_decrypt<A: Clone>(
tags: &[A],
ivks: &[D::IncomingViewingKey],
outputs: &[(D, Output)],
) -> Vec<Option<DecryptedOutput<A, D, Self::Memo>>> {
batch::try_compact_note_decryption(ivks, outputs)
.into_iter()
.map(|res| {
res.map(|((note, recipient), ivk_idx)| DecryptedOutput {
ivk_tag: tags[ivk_idx].clone(),
recipient,
note,
memo: (),
})
})
.collect()
}
}
/// A value correlated with an output index.
struct OutputIndex<V> {
/// The index of the output within the corresponding shielded bundle.
@ -46,12 +117,12 @@ struct OutputIndex<V> {
value: V,
}
type OutputItem<A, D> = OutputIndex<DecryptedNote<A, D>>;
type OutputItem<A, D, M> = OutputIndex<DecryptedOutput<A, D, M>>;
/// The sender for the result of batch scanning a specific transaction output.
struct OutputReplier<A, D: Domain>(OutputIndex<channel::Sender<OutputItem<A, D>>>);
struct OutputReplier<A, D: Domain, M>(OutputIndex<channel::Sender<OutputItem<A, D, M>>>);
impl<A, D: Domain> DynamicUsage for OutputReplier<A, D> {
impl<A, D: Domain, M> DynamicUsage for OutputReplier<A, D, M> {
#[inline(always)]
fn dynamic_usage(&self) -> usize {
// We count the memory usage of items in the channel on the receiver side.
@ -65,9 +136,9 @@ impl<A, D: Domain> DynamicUsage for OutputReplier<A, D> {
}
/// The receiver for the result of batch scanning a specific transaction.
struct BatchReceiver<A, D: Domain>(channel::Receiver<OutputItem<A, D>>);
struct BatchReceiver<A, D: Domain, M>(channel::Receiver<OutputItem<A, D, M>>);
impl<A, D: Domain> DynamicUsage for BatchReceiver<A, D> {
impl<A, D: Domain, M> DynamicUsage for BatchReceiver<A, D, M> {
fn dynamic_usage(&self) -> usize {
// We count the memory usage of items in the channel on the receiver side.
let num_items = self.0.len();
@ -84,7 +155,7 @@ impl<A, D: Domain> DynamicUsage for BatchReceiver<A, D> {
// - Space for an item.
// - The state of the slot, stored as an AtomicUsize.
const PTR_SIZE: usize = std::mem::size_of::<usize>();
let item_size = std::mem::size_of::<OutputItem<A, D>>();
let item_size = std::mem::size_of::<OutputItem<A, D, M>>();
const ATOMIC_USIZE_SIZE: usize = std::mem::size_of::<AtomicUsize>();
let block_size = PTR_SIZE + ITEMS_PER_BLOCK * (item_size + ATOMIC_USIZE_SIZE);
@ -208,7 +279,7 @@ impl<Item: Task> Task for WithUsageTask<Item> {
}
/// A batch of outputs to trial decrypt.
pub(crate) struct Batch<A, D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>> {
pub(crate) struct Batch<A, D: BatchDomain, Output, Dec: Decryptor<D, Output>> {
tags: Vec<A>,
ivks: Vec<D::IncomingViewingKey>,
/// We currently store outputs and repliers as parallel vectors, because
@ -219,15 +290,16 @@ pub(crate) struct Batch<A, D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOT
/// all be part of the same struct, which would also track the output index
/// (that is captured in the outer `OutputIndex` of each `OutputReplier`).
outputs: Vec<(D, Output)>,
repliers: Vec<OutputReplier<A, D>>,
repliers: Vec<OutputReplier<A, D, Dec::Memo>>,
}
impl<A, D, Output> DynamicUsage for Batch<A, D, Output>
impl<A, D, Output, Dec> DynamicUsage for Batch<A, D, Output, Dec>
where
A: DynamicUsage,
D: BatchDomain + DynamicUsage,
D::IncomingViewingKey: DynamicUsage,
Output: ShieldedOutput<D, COMPACT_NOTE_SIZE> + DynamicUsage,
Output: DynamicUsage,
Dec: Decryptor<D, Output>,
{
fn dynamic_usage(&self) -> usize {
self.tags.dynamic_usage()
@ -253,11 +325,11 @@ where
}
}
impl<A, D, Output> Batch<A, D, Output>
impl<A, D, Output, Dec> Batch<A, D, Output, Dec>
where
A: Clone,
D: BatchDomain,
Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>,
Dec: Decryptor<D, Output>,
{
/// Constructs a new batch.
fn new(tags: Vec<A>, ivks: Vec<D::IncomingViewingKey>) -> Self {
@ -276,7 +348,7 @@ where
}
}
impl<A, D, Output> Task for Batch<A, D, Output>
impl<A, D, Output, Dec> Task for Batch<A, D, Output, Dec>
where
A: Clone + Send + 'static,
D: BatchDomain + Send + 'static,
@ -284,7 +356,9 @@ where
D::Memo: Send,
D::Note: Send,
D::Recipient: Send,
Output: ShieldedOutput<D, COMPACT_NOTE_SIZE> + Send + 'static,
Output: Send + 'static,
Dec: Decryptor<D, Output> + 'static,
Dec::Memo: Send,
{
/// Runs the batch of trial decryptions, and reports the results.
fn run(self) {
@ -298,20 +372,16 @@ where
assert_eq!(outputs.len(), repliers.len());
let decryption_results = batch::try_compact_note_decryption(&ivks, &outputs);
let decryption_results = Dec::batch_decrypt(&tags, &ivks, &outputs);
for (decryption_result, OutputReplier(replier)) in
decryption_results.into_iter().zip(repliers.into_iter())
{
// If `decryption_result` is `None` then we will just drop `replier`,
// indicating to the parent `BatchRunner` that this output was not for us.
if let Some(((note, recipient), ivk_idx)) = decryption_result {
if let Some(value) = decryption_result {
let result = OutputIndex {
output_index: replier.output_index,
value: DecryptedNote {
ivk_tag: tags[ivk_idx].clone(),
recipient,
note,
},
value,
};
if replier.value.send(result).is_err() {
@ -323,7 +393,12 @@ where
}
}
impl<A, D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE> + Clone> Batch<A, D, Output> {
impl<A, D, Output, Dec> Batch<A, D, Output, Dec>
where
D: BatchDomain,
Output: Clone,
Dec: Decryptor<D, Output>,
{
/// Adds the given outputs to this batch.
///
/// `replier` will be called with the result of every output.
@ -331,7 +406,7 @@ impl<A, D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE> + Clone> Ba
&mut self,
domain: impl Fn() -> D,
outputs: &[Output],
replier: channel::Sender<OutputItem<A, D>>,
replier: channel::Sender<OutputItem<A, D, Dec::Memo>>,
) {
self.outputs
.extend(outputs.iter().cloned().map(|output| (domain(), output)));
@ -361,28 +436,29 @@ impl DynamicUsage for ResultKey {
}
/// Logic to run batches of trial decryptions on the global threadpool.
pub(crate) struct BatchRunner<A, D, Output, T>
pub(crate) struct BatchRunner<A, D, Output, Dec, T>
where
D: BatchDomain,
Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>,
T: Tasks<Batch<A, D, Output>>,
Dec: Decryptor<D, Output>,
T: Tasks<Batch<A, D, Output, Dec>>,
{
batch_size_threshold: usize,
// The batch currently being accumulated.
acc: Batch<A, D, Output>,
acc: Batch<A, D, Output, Dec>,
// The running batches.
running_tasks: T,
// Receivers for the results of the running batches.
pending_results: HashMap<ResultKey, BatchReceiver<A, D>>,
pending_results: HashMap<ResultKey, BatchReceiver<A, D, Dec::Memo>>,
}
impl<A, D, Output, T> DynamicUsage for BatchRunner<A, D, Output, T>
impl<A, D, Output, Dec, T> DynamicUsage for BatchRunner<A, D, Output, Dec, T>
where
A: DynamicUsage,
D: BatchDomain + DynamicUsage,
D::IncomingViewingKey: DynamicUsage,
Output: ShieldedOutput<D, COMPACT_NOTE_SIZE> + DynamicUsage,
T: Tasks<Batch<A, D, Output>> + DynamicUsage,
Output: DynamicUsage,
Dec: Decryptor<D, Output>,
T: Tasks<Batch<A, D, Output, Dec>> + DynamicUsage,
{
fn dynamic_usage(&self) -> usize {
self.acc.dynamic_usage()
@ -408,12 +484,12 @@ where
}
}
impl<A, D, Output, T> BatchRunner<A, D, Output, T>
impl<A, D, Output, Dec, T> BatchRunner<A, D, Output, Dec, T>
where
A: Clone,
D: BatchDomain,
Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>,
T: Tasks<Batch<A, D, Output>>,
Dec: Decryptor<D, Output>,
T: Tasks<Batch<A, D, Output, Dec>>,
{
/// Constructs a new batch runner for the given incoming viewing keys.
pub(crate) fn new(
@ -430,7 +506,7 @@ where
}
}
impl<A, D, Output, T> BatchRunner<A, D, Output, T>
impl<A, D, Output, Dec, T> BatchRunner<A, D, Output, Dec, T>
where
A: Clone + Send + 'static,
D: BatchDomain + Send + 'static,
@ -438,8 +514,9 @@ where
D::Memo: Send,
D::Note: Send,
D::Recipient: Send,
Output: ShieldedOutput<D, COMPACT_NOTE_SIZE> + Clone + Send + 'static,
T: Tasks<Batch<A, D, Output>>,
Output: Clone + Send + 'static,
Dec: Decryptor<D, Output>,
T: Tasks<Batch<A, D, Output, Dec>>,
{
/// Batches the given outputs for trial decryption.
///
@ -447,9 +524,9 @@ where
/// batch, or the all-zeros hash to indicate that no block triggered it (i.e. it was a
/// mempool change).
///
/// If after adding the given outputs, the accumulated batch size is at least
/// `BATCH_SIZE_THRESHOLD`, `Self::flush` is called. Subsequent calls to
/// `Self::add_outputs` will be accumulated into a new batch.
/// If after adding the given outputs, the accumulated batch size is at least the size
/// threshold that was set via `Self::new`, `Self::flush` is called. Subsequent calls
/// to `Self::add_outputs` will be accumulated into a new batch.
pub(crate) fn add_outputs(
&mut self,
block_tag: BlockHash,
@ -487,7 +564,7 @@ where
&mut self,
block_tag: BlockHash,
txid: TxId,
) -> HashMap<(TxId, usize), DecryptedNote<A, D>> {
) -> HashMap<(TxId, usize), DecryptedOutput<A, D, Dec::Memo>> {
self.pending_results
.remove(&ResultKey(block_tag, txid))
// We won't have a pending result if the transaction didn't have outputs of

19
zcash_client_backend/src/scanning.rs
View File

@ -21,7 +21,7 @@ use zcash_primitives::{
use crate::data_api::{BlockMetadata, ScannedBlock, ScannedBundles};
use crate::{
proto::compact_formats::CompactBlock,
scan::{Batch, BatchRunner, Tasks},
scan::{Batch, BatchRunner, CompactDecryptor, Tasks},
wallet::{WalletSaplingOutput, WalletSaplingSpend, WalletTx},
ShieldedProtocol,
};
@ -268,9 +268,10 @@ pub fn scan_block<P: consensus::Parameters + Send + 'static, K: ScanningKey>(
)
}
type TaggedBatch<S> = Batch<(AccountId, S), SaplingDomain, CompactOutputDescription>;
type TaggedBatch<S> =
Batch<(AccountId, S), SaplingDomain, CompactOutputDescription, CompactDecryptor>;
type TaggedBatchRunner<S, T> =
BatchRunner<(AccountId, S), SaplingDomain, CompactOutputDescription, T>;
BatchRunner<(AccountId, S), SaplingDomain, CompactOutputDescription, CompactDecryptor, T>;
#[tracing::instrument(skip_all, fields(height = block.height))]
pub(crate) fn add_block_to_runner<P, S, T>(
@ -524,13 +525,13 @@ pub(crate) fn scan_block_with_runner<
let mut decrypted = runner.collect_results(cur_hash, txid);
(0..decoded.len())
.map(|i| {
decrypted.remove(&(txid, i)).map(|d_note| {
let a = d_note.ivk_tag.0;
let nk = vks.get(&d_note.ivk_tag).expect(
decrypted.remove(&(txid, i)).map(|d_out| {
let a = d_out.ivk_tag.0;
let nk = vks.get(&d_out.ivk_tag).expect(
"The batch runner and scan_block must use the same set of IVKs.",
);
(d_note.note, a, d_note.ivk_tag.1, (*nk).clone())
(d_out.note, a, d_out.ivk_tag.1, (*nk).clone())
})
})
.collect()
@ -834,7 +835,7 @@ mod tests {
assert_eq!(cb.vtx.len(), 2);
let mut batch_runner = if scan_multithreaded {
let mut runner = BatchRunner::<_, _, _, ()>::new(
let mut runner = BatchRunner::<_, _, _, _, ()>::new(
10,
dfvk.to_sapling_keys()
.iter()
@ -921,7 +922,7 @@ mod tests {
assert_eq!(cb.vtx.len(), 3);
let mut batch_runner = if scan_multithreaded {
let mut runner = BatchRunner::<_, _, _, ()>::new(
let mut runner = BatchRunner::<_, _, _, _, ()>::new(
10,
dfvk.to_sapling_keys()
.iter()