zcash_client_backend/
scan.rs

1use flume as channel;
2use std::collections::HashMap;
3use std::fmt;
4use std::mem;
5use std::sync::{
6    Arc,
7    atomic::{AtomicUsize, Ordering},
8};
9
10use memuse::DynamicUsage;
11use zcash_note_encryption::{BatchDomain, Domain, ShieldedOutput, batch};
12use zcash_primitives::{block::BlockHash, transaction::TxId};
13
14/// A decrypted transaction output.
15pub(crate) struct DecryptedOutput<IvkTag, D: Domain, M> {
16    /// The tag corresponding to the incoming viewing key used to decrypt the note.
17    pub(crate) ivk_tag: IvkTag,
18    /// The recipient of the note.
19    pub(crate) recipient: D::Recipient,
20    /// The note!
21    pub(crate) note: D::Note,
22    /// The memo field, or `()` if this is a decrypted compact output.
23    pub(crate) memo: M,
24}
25
26impl<IvkTag, D: Domain, M> fmt::Debug for DecryptedOutput<IvkTag, D, M>
27where
28    IvkTag: fmt::Debug,
29    D::IncomingViewingKey: fmt::Debug,
30    D::Recipient: fmt::Debug,
31    D::Note: fmt::Debug,
32    M: fmt::Debug,
33{
34    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
35        f.debug_struct("DecryptedOutput")
36            .field("ivk_tag", &self.ivk_tag)
37            .field("recipient", &self.recipient)
38            .field("note", &self.note)
39            .field("memo", &self.memo)
40            .finish()
41    }
42}
43
44/// A decryptor of transaction outputs.
45pub(crate) trait Decryptor<D: BatchDomain, Output> {
46    type Memo;
47
48    fn batch_decrypt<IvkTag: Clone>(
49        tags: &[IvkTag],
50        ivks: &[D::IncomingViewingKey],
51        outputs: &[(D, Output)],
52    ) -> impl Iterator<Item = Option<DecryptedOutput<IvkTag, D, Self::Memo>>>;
53}
54
55/// A decryptor of outputs as encoded in transactions.
56#[allow(dead_code)]
57pub(crate) struct FullDecryptor;
58
59impl<D: BatchDomain, Output: ShieldedOutput<D>> Decryptor<D, Output> for FullDecryptor {
60    type Memo = D::Memo;
61
62    fn batch_decrypt<IvkTag: Clone>(
63        tags: &[IvkTag],
64        ivks: &[D::IncomingViewingKey],
65        outputs: &[(D, Output)],
66    ) -> impl Iterator<Item = Option<DecryptedOutput<IvkTag, D, Self::Memo>>> {
67        batch::try_note_decryption(ivks, outputs)
68            .into_iter()
69            .map(|res| {
70                res.map(|((note, recipient, memo), ivk_idx)| DecryptedOutput {
71                    ivk_tag: tags[ivk_idx].clone(),
72                    recipient,
73                    note,
74                    memo,
75                })
76            })
77    }
78}
79
80/// A decryptor of outputs as encoded in compact blocks.
81pub(crate) struct CompactDecryptor;
82
83impl<D: BatchDomain, Output: ShieldedOutput<D>> Decryptor<D, Output> for CompactDecryptor {
84    type Memo = ();
85
86    fn batch_decrypt<IvkTag: Clone>(
87        tags: &[IvkTag],
88        ivks: &[D::IncomingViewingKey],
89        outputs: &[(D, Output)],
90    ) -> impl Iterator<Item = Option<DecryptedOutput<IvkTag, D, Self::Memo>>> {
91        batch::try_compact_note_decryption(ivks, outputs)
92            .into_iter()
93            .map(|res| {
94                res.map(|((note, recipient), ivk_idx)| DecryptedOutput {
95                    ivk_tag: tags[ivk_idx].clone(),
96                    recipient,
97                    note,
98                    memo: (),
99                })
100            })
101    }
102}
103
104/// A value correlated with an output index.
105struct OutputIndex<V> {
106    /// The index of the output within the corresponding shielded bundle.
107    output_index: usize,
108    /// The value for the output index.
109    value: V,
110}
111
112type OutputItem<IvkTag, D, M> = OutputIndex<DecryptedOutput<IvkTag, D, M>>;
113
114/// The sender for the result of batch scanning a specific transaction output.
115struct OutputReplier<IvkTag, D: Domain, M>(OutputIndex<channel::Sender<OutputItem<IvkTag, D, M>>>);
116
117impl<IvkTag, D: Domain, M> DynamicUsage for OutputReplier<IvkTag, D, M> {
118    #[inline(always)]
119    fn dynamic_usage(&self) -> usize {
120        // We count the memory usage of items in the channel on the receiver side.
121        0
122    }
123
124    #[inline(always)]
125    fn dynamic_usage_bounds(&self) -> (usize, Option<usize>) {
126        (0, Some(0))
127    }
128}
129
130/// The receiver for the result of batch scanning a specific transaction.
131pub(crate) struct BatchReceiver<IvkTag, D: Domain, M>(channel::Receiver<OutputItem<IvkTag, D, M>>);
132
133impl<IvkTag, D: Domain, M> DynamicUsage for BatchReceiver<IvkTag, D, M> {
134    fn dynamic_usage(&self) -> usize {
135        // We count the memory usage of items still buffered in the channel on the
136        // receiver side. This is a loose lower bound rather than an exact figure:
137        // `flume` stores queued items in an internal buffer that may be over-allocated
138        // relative to its length and carries per-item bookkeeping, so the true heap use
139        // can exceed this. It is only used as a soft memory-pressure heuristic.
140        self.0.len() * std::mem::size_of::<OutputItem<IvkTag, D, M>>()
141    }
142
143    fn dynamic_usage_bounds(&self) -> (usize, Option<usize>) {
144        let usage = self.dynamic_usage();
145        (usage, Some(usage))
146    }
147}
148
149impl<IvkTag, D: Domain, M> BatchReceiver<IvkTag, D, M> {
150    /// Blocks until the results of the batch are ready.
151    pub(crate) fn into_results(self) -> HashMap<usize, DecryptedOutput<IvkTag, D, M>> {
152        // This iterator will end once the channel becomes empty and disconnected. We
153        // created one sender per output, and each sender is dropped after the batch it is
154        // in completes (and in the case of successful decryptions, after the decrypted
155        // note has been sent to the channel). Completion of the iterator therefore
156        // corresponds to complete knowledge of the outputs of this transaction that could
157        // be decrypted.
158        self.0
159            .into_iter()
160            .map(
161                |OutputIndex {
162                     output_index,
163                     value,
164                 }| { (output_index, value) },
165            )
166            .collect()
167    }
168
169    /// Waits, without blocking the current thread, until the results of the batch are
170    /// ready.
171    ///
172    /// This is the asynchronous counterpart of [`Self::into_results`]; the channel
173    /// completion semantics are identical.
174    #[cfg(feature = "sync-decryptor")]
175    pub(crate) async fn into_results_async(self) -> HashMap<usize, DecryptedOutput<IvkTag, D, M>> {
176        let mut results = HashMap::new();
177        while let Ok(OutputIndex {
178            output_index,
179            value,
180        }) = self.0.recv_async().await
181        {
182            results.insert(output_index, value);
183        }
184        results
185    }
186}
187
188/// A tracker for the batch scanning tasks that are currently running.
189///
190/// This enables a [`BatchRunner`] to be optionally configured to track heap memory usage.
191pub(crate) trait Tasks<Item> {
192    type Task: Task;
193    fn new() -> Self;
194    fn add_task(&self, item: Item) -> Self::Task;
195    fn run_task(&self, item: Item) {
196        let task = self.add_task(item);
197        rayon::spawn_fifo(|| task.run());
198    }
199}
200
201/// A batch scanning task.
202pub(crate) trait Task: Send + 'static {
203    fn run(self);
204}
205
206impl<Item: Task> Tasks<Item> for () {
207    type Task = Item;
208    fn new() -> Self {}
209    fn add_task(&self, item: Item) -> Self::Task {
210        // Return the item itself as the task; we aren't tracking anything about it, so
211        // there is no need to wrap it in a newtype.
212        item
213    }
214}
215
216/// A task tracker that measures heap usage.
217///
218/// This struct implements `DynamicUsage` without any item bounds, but that works because
219/// it only implements `Tasks` for items that implement `DynamicUsage`.
220#[allow(dead_code)]
221pub(crate) struct WithUsage {
222    // The current heap usage for all running tasks.
223    running_usage: Arc<AtomicUsize>,
224}
225
226impl DynamicUsage for WithUsage {
227    fn dynamic_usage(&self) -> usize {
228        self.running_usage.load(Ordering::Relaxed)
229    }
230
231    fn dynamic_usage_bounds(&self) -> (usize, Option<usize>) {
232        // Tasks are relatively short-lived, so we accept the inaccuracy of treating the
233        // tasks's approximate usage as its bounds.
234        let usage = self.dynamic_usage();
235        (usage, Some(usage))
236    }
237}
238
239impl<Item: Task + DynamicUsage> Tasks<Item> for WithUsage {
240    type Task = WithUsageTask<Item>;
241
242    fn new() -> Self {
243        Self {
244            running_usage: Arc::new(AtomicUsize::new(0)),
245        }
246    }
247
248    fn add_task(&self, item: Item) -> Self::Task {
249        // Create the task that will move onto the heap with the batch item.
250        let mut task = WithUsageTask {
251            item,
252            own_usage: 0,
253            running_usage: self.running_usage.clone(),
254        };
255
256        // `rayon::spawn_fifo` creates a `HeapJob` holding a closure. The size of a
257        // closure is (to good approximation) the size of the captured environment, which
258        // in this case is two moved variables:
259        // - An `Arc<Registry>`, which is a pointer to data that is amortized over the
260        //   entire `rayon` thread pool, so we only count the pointer size here.
261        // - The spawned closure, which in our case moves `task` into it.
262        task.own_usage =
263            mem::size_of::<Arc<()>>() + mem::size_of_val(&task) + task.item.dynamic_usage();
264
265        // Approximate now as when the heap cost of this running batch begins. In practice
266        // this is fine, because `Self::add_task` is called from `Self::run_task` which
267        // immediately moves the task to the heap.
268        self.running_usage
269            .fetch_add(task.own_usage, Ordering::SeqCst);
270
271        task
272    }
273}
274
275/// A task that will clean up its own heap usage from the overall running usage once it is
276/// complete.
277#[allow(dead_code)]
278pub(crate) struct WithUsageTask<Item> {
279    /// The item being run.
280    item: Item,
281    /// Size of this task on the heap. We assume that the size of the task does not change
282    /// once it has been created, to avoid needing to maintain bidirectional channels
283    /// between [`WithUsage`] and its tasks.
284    own_usage: usize,
285    /// Pointer to the parent [`WithUsage`]'s heap usage tracker for running tasks.
286    running_usage: Arc<AtomicUsize>,
287}
288
289impl<Item: Task> Task for WithUsageTask<Item> {
290    fn run(self) {
291        // Run the item.
292        self.item.run();
293
294        // Signal that the heap memory for this task has been freed.
295        self.running_usage
296            .fetch_sub(self.own_usage, Ordering::SeqCst);
297    }
298}
299
300/// A batch of outputs to trial decrypt.
301pub(crate) struct Batch<IvkTag, D: BatchDomain, Output, Dec: Decryptor<D, Output>> {
302    tags: Vec<IvkTag>,
303    ivks: Vec<D::IncomingViewingKey>,
304    /// We currently store outputs and repliers as parallel vectors, because
305    /// [`batch::try_note_decryption`] accepts a slice of domain/output pairs
306    /// rather than a value that implements `IntoIterator`, and therefore we
307    /// can't just use `map` to select the parts we need in order to perform
308    /// batch decryption. Ideally the domain, output, and output replier would
309    /// all be part of the same struct, which would also track the output index
310    /// (that is captured in the outer `OutputIndex` of each `OutputReplier`).
311    outputs: Vec<(D, Output)>,
312    repliers: Vec<OutputReplier<IvkTag, D, Dec::Memo>>,
313}
314
315impl<IvkTag, D, Output, Dec> DynamicUsage for Batch<IvkTag, D, Output, Dec>
316where
317    IvkTag: DynamicUsage,
318    D: BatchDomain + DynamicUsage,
319    D::IncomingViewingKey: DynamicUsage,
320    Output: DynamicUsage,
321    Dec: Decryptor<D, Output>,
322{
323    fn dynamic_usage(&self) -> usize {
324        self.tags.dynamic_usage()
325            + self.ivks.dynamic_usage()
326            + self.outputs.dynamic_usage()
327            + self.repliers.dynamic_usage()
328    }
329
330    fn dynamic_usage_bounds(&self) -> (usize, Option<usize>) {
331        let (tags_lower, tags_upper) = self.tags.dynamic_usage_bounds();
332        let (ivks_lower, ivks_upper) = self.ivks.dynamic_usage_bounds();
333        let (outputs_lower, outputs_upper) = self.outputs.dynamic_usage_bounds();
334        let (repliers_lower, repliers_upper) = self.repliers.dynamic_usage_bounds();
335
336        (
337            tags_lower + ivks_lower + outputs_lower + repliers_lower,
338            tags_upper
339                .zip(ivks_upper)
340                .zip(outputs_upper)
341                .zip(repliers_upper)
342                .map(|(((a, b), c), d)| a + b + c + d),
343        )
344    }
345}
346
347impl<IvkTag, D, Output, Dec> Batch<IvkTag, D, Output, Dec>
348where
349    IvkTag: Clone,
350    D: BatchDomain,
351    Dec: Decryptor<D, Output>,
352{
353    /// Constructs a new batch.
354    fn new(tags: Vec<IvkTag>, ivks: Vec<D::IncomingViewingKey>) -> Self {
355        assert_eq!(tags.len(), ivks.len());
356        Self {
357            tags,
358            ivks,
359            outputs: vec![],
360            repliers: vec![],
361        }
362    }
363
364    /// Returns `true` if the batch is currently empty.
365    fn is_empty(&self) -> bool {
366        self.outputs.is_empty()
367    }
368}
369
370impl<IvkTag, D, Output, Dec> Task for Batch<IvkTag, D, Output, Dec>
371where
372    IvkTag: Clone + Send + 'static,
373    D: BatchDomain + Send + 'static,
374    D::IncomingViewingKey: Send,
375    D::Memo: Send,
376    D::Note: Send,
377    D::Recipient: Send,
378    Output: Send + 'static,
379    Dec: Decryptor<D, Output> + 'static,
380    Dec::Memo: Send,
381{
382    /// Runs the batch of trial decryptions, and reports the results.
383    fn run(self) {
384        // Deconstruct self so we can consume the pieces individually.
385        let Self {
386            tags,
387            ivks,
388            outputs,
389            repliers,
390        } = self;
391
392        assert_eq!(outputs.len(), repliers.len());
393
394        let decryption_results = Dec::batch_decrypt(&tags, &ivks, &outputs);
395        for (decryption_result, OutputReplier(replier)) in decryption_results.zip(repliers) {
396            // If `decryption_result` is `None` then we will just drop `replier`,
397            // indicating to the parent `BatchRunner` that this output was not for us.
398            if let Some(value) = decryption_result {
399                let result = OutputIndex {
400                    output_index: replier.output_index,
401                    value,
402                };
403
404                if replier.value.send(result).is_err() {
405                    tracing::debug!("BatchRunner was dropped before batch finished");
406                    break;
407                }
408            }
409        }
410    }
411}
412
413impl<IvkTag, D, Output, Dec> Batch<IvkTag, D, Output, Dec>
414where
415    D: BatchDomain,
416    Output: Clone,
417    Dec: Decryptor<D, Output>,
418{
419    /// Adds the given outputs to this batch.
420    ///
421    /// `replier` will be called with the result of every output.
422    fn add_outputs(
423        &mut self,
424        domain: impl Fn(&Output) -> D,
425        outputs: impl IntoIterator<Item = Output>,
426        replier: channel::Sender<OutputItem<IvkTag, D, Dec::Memo>>,
427    ) {
428        for (output_index, output) in outputs.into_iter().enumerate() {
429            self.outputs.push((domain(&output), output));
430            self.repliers.push(OutputReplier(OutputIndex {
431                output_index,
432                value: replier.clone(),
433            }));
434        }
435    }
436}
437
438/// A `HashMap` key for looking up the result of a batch scanning a specific transaction.
439#[derive(PartialEq, Eq, Hash)]
440struct ResultKey(BlockHash, TxId);
441
442impl DynamicUsage for ResultKey {
443    #[inline(always)]
444    fn dynamic_usage(&self) -> usize {
445        0
446    }
447
448    #[inline(always)]
449    fn dynamic_usage_bounds(&self) -> (usize, Option<usize>) {
450        (0, Some(0))
451    }
452}
453
454/// Logic to run batches of trial decryptions on the global threadpool.
455pub(crate) struct BatchRunner<IvkTag, D, Output, Dec, T>
456where
457    D: BatchDomain,
458    Dec: Decryptor<D, Output>,
459    T: Tasks<Batch<IvkTag, D, Output, Dec>>,
460{
461    batch_size_threshold: usize,
462    // The batch currently being accumulated.
463    acc: Batch<IvkTag, D, Output, Dec>,
464    // The running batches.
465    running_tasks: T,
466    // Receivers for the results of the running batches.
467    pending_results: HashMap<ResultKey, BatchReceiver<IvkTag, D, Dec::Memo>>,
468}
469
470impl<IvkTag, D, Output, Dec, T> DynamicUsage for BatchRunner<IvkTag, D, Output, Dec, T>
471where
472    IvkTag: DynamicUsage,
473    D: BatchDomain + DynamicUsage,
474    D::IncomingViewingKey: DynamicUsage,
475    Output: DynamicUsage,
476    Dec: Decryptor<D, Output>,
477    T: Tasks<Batch<IvkTag, D, Output, Dec>> + DynamicUsage,
478{
479    fn dynamic_usage(&self) -> usize {
480        self.acc.dynamic_usage()
481            + self.running_tasks.dynamic_usage()
482            + self.pending_results.dynamic_usage()
483    }
484
485    fn dynamic_usage_bounds(&self) -> (usize, Option<usize>) {
486        let running_usage = self.running_tasks.dynamic_usage();
487
488        let bounds = (
489            self.acc.dynamic_usage_bounds(),
490            self.pending_results.dynamic_usage_bounds(),
491        );
492        (
493            bounds.0.0 + running_usage + bounds.1.0,
494            bounds
495                .0
496                .1
497                .zip(bounds.1.1)
498                .map(|(a, b)| a + running_usage + b),
499        )
500    }
501}
502
503impl<IvkTag, D, Output, Dec, T> BatchRunner<IvkTag, D, Output, Dec, T>
504where
505    IvkTag: Clone,
506    D: BatchDomain,
507    Dec: Decryptor<D, Output>,
508    T: Tasks<Batch<IvkTag, D, Output, Dec>>,
509{
510    /// Constructs a new batch runner for the given incoming viewing keys.
511    pub(crate) fn new(
512        batch_size_threshold: usize,
513        ivks: impl Iterator<Item = (IvkTag, D::IncomingViewingKey)>,
514    ) -> Self {
515        let (tags, ivks) = ivks.unzip();
516        Self {
517            batch_size_threshold,
518            acc: Batch::new(tags, ivks),
519            running_tasks: T::new(),
520            pending_results: HashMap::default(),
521        }
522    }
523}
524
525impl<IvkTag, D, Output, Dec, T> BatchRunner<IvkTag, D, Output, Dec, T>
526where
527    IvkTag: Clone + Send + 'static,
528    D: BatchDomain + Send + 'static,
529    D::IncomingViewingKey: Clone + Send,
530    D::Memo: Send,
531    D::Note: Send,
532    D::Recipient: Send,
533    Output: Clone + Send + 'static,
534    Dec: Decryptor<D, Output>,
535    T: Tasks<Batch<IvkTag, D, Output, Dec>>,
536{
537    /// Batches the given outputs for trial decryption.
538    ///
539    /// `block_tag` is the hash of the block that triggered this txid being added to the
540    /// batch, or the all-zeros hash to indicate that no block triggered it (i.e. it was a
541    /// mempool change).
542    ///
543    /// The decryption results can be obtained via [`Self::collect_results`]. To manage
544    /// result collection manually, use [`Self::process_outputs`] instead of this method.
545    ///
546    /// If after adding the given outputs, the accumulated batch size is at least the size
547    /// threshold that was set via `Self::new`, `Self::flush` is called. Subsequent calls
548    /// to `Self::add_outputs` will be accumulated into a new batch.
549    pub(crate) fn add_outputs(
550        &mut self,
551        block_tag: BlockHash,
552        txid: TxId,
553        domain: impl Fn(&Output) -> D,
554        outputs: impl IntoIterator<Item = Output>,
555    ) {
556        let batch_receiver = self.process_outputs(domain, outputs);
557        self.pending_results
558            .insert(ResultKey(block_tag, txid), batch_receiver);
559    }
560
561    /// Batches the given outputs for trial decryption.
562    ///
563    /// Returns a handle for receiving the results of the batch. To have the batch runner
564    /// manage this for you, use [`Self::add_outputs`] instead of this method.
565    ///
566    /// If after adding the given outputs, the accumulated batch size is at least the size
567    /// threshold that was set via `Self::new`, `Self::flush` is called. Subsequent calls
568    /// to either `Self::process_outputs` or `Self::add_outputs` will be accumulated into a
569    /// new batch.
570    pub(crate) fn process_outputs(
571        &mut self,
572        domain: impl Fn(&Output) -> D,
573        outputs: impl IntoIterator<Item = Output>,
574    ) -> BatchReceiver<IvkTag, D, Dec::Memo> {
575        // Each output is given its own clone of the sending half of the channel, and the
576        // returned `BatchReceiver` holds the only receiving half. Successful decryptions
577        // are sent to the channel, and every sender is dropped once the batch it belongs
578        // to has run; the receiver therefore yields exactly the decryptable outputs and
579        // then completes once all the senders are gone.
580        let (tx, rx) = channel::unbounded();
581        self.acc.add_outputs(domain, outputs, tx);
582
583        // Run the batch eagerly once it reaches the configured size, rather than letting
584        // it grow unbounded; any remaining outputs accumulate into the next batch.
585        if self.acc.outputs.len() >= self.batch_size_threshold {
586            self.flush();
587        }
588
589        BatchReceiver(rx)
590    }
591
592    /// Runs the currently accumulated batch on the global threadpool.
593    ///
594    /// Subsequent calls to `Self::add_outputs` will be accumulated into a new batch.
595    pub(crate) fn flush(&mut self) {
596        if !self.acc.is_empty() {
597            let mut batch = Batch::new(self.acc.tags.clone(), self.acc.ivks.clone());
598            mem::swap(&mut batch, &mut self.acc);
599            self.running_tasks.run_task(batch);
600        }
601    }
602
603    /// Collects the pending decryption results for the given transaction.
604    ///
605    /// `block_tag` is the hash of the block that triggered this txid being added to the
606    /// batch, or the all-zeros hash to indicate that no block triggered it (i.e. it was a
607    /// mempool change).
608    pub(crate) fn collect_results(
609        &mut self,
610        block_tag: BlockHash,
611        txid: TxId,
612    ) -> HashMap<usize, DecryptedOutput<IvkTag, D, Dec::Memo>> {
613        self.pending_results
614            .remove(&ResultKey(block_tag, txid))
615            // We won't have a pending result if the transaction didn't have outputs of
616            // this runner's kind.
617            .map(BatchReceiver::into_results)
618            .unwrap_or_default()
619    }
620}
zcash_client_backend/scan.rs

zcash_client_backend/
scan.rs
