Move CPU-heavy proof preparation into the batch cryptography thread

[teor2345]

Motivation

Zebra's sync can be slow or fail, because proof batch preparation involves some CPU-heavy cryptography.

So we want to move proof preparation into the existing batch rayon CPU thread pool scope.

API Reference

Blocking threads:
https://docs.rs/tokio/latest/tokio/task/fn.spawn_blocking.html

CPU-heavy threads:
https://docs.rs/rayon/latest/rayon/fn.scope_fifo.html

Detailed Analysis

See #4583 (comment)

Designs

When queueing batch items, store them raw, then do preparation inside the rayon scope, for the following verifiers:

• groth16
• halo2

Existing code for rayon in tokio blocking threads:

zebra/zebra-consensus/src/primitives/ed25519.rs

Lines 121 to 134 in 81727d7

	/// Flush the batch using a thread pool, and return the result via the channel.
	/// This function returns a future that becomes ready when the batch is completed.
	fn flush_spawning(batch: BatchVerifier, tx: Sender) -> impl Future<Output = ()> {
	// Correctness: Do CPU-intensive work on a dedicated thread, to avoid blocking other futures.
	tokio::task::spawn_blocking(|| {
	// TODO:
	// - spawn batches so rayon executes them in FIFO order
	// possible implementation: return a closure in a Future,
	// then run it using scope_fifo() in the worker task,
	// limiting the number of concurrent batches to the number of rayon threads
	rayon::scope_fifo(|s| s.spawn_fifo(|_s| Self::verify(batch, tx)))
	})
	.map(|join_result| join_result.expect("panic in ed25519 batch verifier"))
	}

Related Work

[conradoplg]

Note that we wanted to move proof decoding inside tx deserialization: #3179

I don't think this changes anything since the proposed approach in this ticket is probably easier, but I thought it was worth it pointing out (i.e. another solution would be to do #3179, and move the entire deserialization to a thread, which I think we'll do anyway in another ticket).

If we follow the proposed approach in this ticket then we should probably close #3179

[upbqdn]

Good point @conradoplg. I can take this ticket, but if we move the proof preparation to rayon, shouldn't we still keep in mind that we should use bellman::Proof in Groth16Proof instead of a byte array?

[conradoplg]

Yep we need to decide which approach to use:

• Leave the byte array, and decode inside the batch verifier, as suggested in the ticket
• Replace the byte array, such that the decoding happens when deserializing, and then Move network transaction deserialization to a dedicated blocking and CPU-heavy thread #4787 will automatically solve this issue. The downside is that replacing the array requires a lot of changes (not that much, but it seems it will be a bit annoying) as described in Use a bellman::Proof instead of a byte array in Groth16Proof #3179 so we may not want to do this right now.

(Thinking about it, if we do the first, we can keep #3179 open as long as we add a note to it to remember to revert the decoding inside the batch verifier)

[upbqdn]

I think I'd lean towards the second option in terms of future maintenance since it seems a bit cleaner.

[teor2345]

I'd be interested in comparing the performance differences between the two different options.

I think at the moment we're looking for a quick fix, so we can re-test Zebra sync, and find any bugs or excessive CPU usage.

[daira]

I'd be happy to change https://github.com/zcash-hackworks/zcash-test-vectors/ to generate structurally valid (but still unverifiable) Groth16 proofs, which seems to be one of the main blockers for #3179. I won't have time to do this before Zcon, though.

[teor2345]

Right now, we need to do the quick fix of moving proof decoding into the batch rayon threads.

After Zebra is syncing correctly, we can explore cleanup/simplification tasks like #3179.

[teor2345]

We're also seeing stalls during checkpoint verification, which doesn't use the batch verifier code.
So we might want to review the priority of this ticket.

We could check how much time proof preparation is actually taking?
I'll try to open a PR with some easy-to-use functions today.

[teor2345]

Conrado did some testing, and proof preparation takes 1-25ms.

So we don't need to do this ticket right now, but we should do it eventually, to handle large aggregated Orchard proofs.

@upbqdn how much progress have you made on this ticket so far?
It's ok to finish it off and submit a PR if you want.

[upbqdn]

@upbqdn how much progress have you made on this ticket so far?

I have no implementation ready to be pushed yet.

[teor2345]

@upbqdn how much progress have you made on this ticket so far?

I have no implementation ready to be pushed yet.

Feel free to finish it off, or to leave it for now, and move on to something that's more urgent.

[teor2345]

It looks like no-one is working on this right now, and we've finished the epic it was a part of.

[teor2345]

Performance improvements are not a priority for Zebra right now.