feat: Lock-free apply_block refactor - #2345
Conversation
…ckfree-store-state
…ckfree-store-state
…ckfree-store-state
Mirko-von-Leipzig
left a comment
There was a problem hiding this comment.
LGTM, thank you; this is already much better.
I left some comments regarding splitting the read and write portion sooner somehow; but thinking on it some more I'm not sure its tractable to implement at the moment.
| pub fn reader(&self) -> AccountTreeWithHistory<S::Reader> { | ||
| let latest = self.latest.reader().expect("snapshot creation should not fail"); | ||
| AccountTreeWithHistory { | ||
| block_number: self.block_number, | ||
| latest, | ||
| overlays: self.overlays.clone(), | ||
| } | ||
| } |
There was a problem hiding this comment.
I think the overall API shape here is incorrect, but perhaps that can be part of the followup.
What I would like is a guarantee that readers can never access writing. But with this API everybody has access to AccountTreeWithHistory and can therefore invoke apply_mutations if they so wish.
Something similar to what I suggested here #2353 (comment).
i.e. on Smt::load you get a single SmtWriter (non-clone), and a SmtReader. This would go together with a single WriteConnection and a ReadOnlyConnectionPool for the database.
The write task would get StateWriter (1 db connection, 1 smt writer per smt), and everything else gets ReadOnlyState for queries. This way we guarantee very early on that nothing can escalate their privilege.
So ito this code, reader would become SmtReader::snapshot(BlockNumber) -> AccountTreeSnapshot, or something similar.
There was a problem hiding this comment.
What I would like is a guarantee that readers can never access writing. But with this API everybody has access to AccountTreeWithHistory and can therefore invoke apply_mutations if they so wish.
The apply_mutations fn lives in impl<S: SmtStorage> AccountTreeWithHistory<S> { so it is not callable from the AccountTreeWithHistory<SmtStorageReader> which is returned by the reader() here.
There was a problem hiding this comment.
Yes but readers already have access to super struct i.e. they could elect to write.
I don't want to audit this by reading every read query and checking they do the correct thing.
Ideally, right at the very start, before anything even gets launched, we already split things up. But perhaps that isn't possible because apply_block method lives together with the read methods so we cannot do that.
As an example, if we had the write methods in a separate gRPC service, we could give the write service state the write struct. And the readers would only ever get the read struct. But this isn't possible unless we split the services, which seems a bit silly (? perhaps ?).
There was a problem hiding this comment.
But lets ignore this for.
There was a problem hiding this comment.
Have unresolved to keep track of this. Might add to a stacked PR.
There was a problem hiding this comment.
I would create a separate file per function instead.
There was a problem hiding this comment.
I also didn't review the contents under the assumption that this was a copy-paste move.
| #[must_use = "call `start` to spawn the block writer and obtain the state"] | ||
| pub struct LoadedState { | ||
| state: State, | ||
| writer: BlockWriter, | ||
| } | ||
|
|
||
| impl LoadedState { | ||
| /// Spawns the block writer onto the current runtime and returns the state together with the | ||
| /// writer task's handle. | ||
| /// | ||
| /// The writer exits once the shutdown token passed to [`State::load`] is cancelled or the last | ||
| /// state reference (holding the only write handle) is dropped — an in-flight block write | ||
| /// always completes first. Awaiting the returned handle after either event guarantees the | ||
| /// writer has released the tree storage it owns; a join error carries a writer panic. | ||
| /// | ||
| /// Callers without a token to cancel should stop the store via [`State::stop`] rather than | ||
| /// dropping and joining by hand. | ||
| pub fn start(self) -> (Arc<State>, WriterTask) { | ||
| let writer_task = tokio::spawn(self.writer.run()); | ||
| (Arc::new(self.state), WriterTask(writer_task)) | ||
| } | ||
| } | ||
|
|
||
| // WRITER TASK | ||
| // ================================================================================================ | ||
|
|
||
| /// Handle of the store's block writer task, returned by [`LoadedState::start`]. | ||
| /// | ||
| /// Awaiting it resolves once the writer has exited and released the tree storage it owns; a join | ||
| /// error carries a writer panic. The newtype ensures [`State::stop`] can only be given the store's | ||
| /// own writer task, and deliberately does not expose [`tokio::task::JoinHandle::abort`]: aborting | ||
| /// the writer mid-write could leave the trees lagging the committed database state, voiding the | ||
| /// guarantee that an in-flight block write always completes. | ||
| #[must_use = "await the writer task to observe its exit, or pass it to `State::stop`"] | ||
| pub struct WriterTask(tokio::task::JoinHandle<()>); | ||
|
|
||
| impl Future for WriterTask { | ||
| type Output = Result<(), tokio::task::JoinError>; | ||
|
|
||
| fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { | ||
| Pin::new(&mut self.0).poll(cx) | ||
| } | ||
| } |
There was a problem hiding this comment.
As mentioned; I think this is the wrong abstraction. I haven't checked how WriterTask works yet though.
I feel like State::load could return ReadOnlyState(Arc<State>) + StateWriter. The task spawning etc and what that means in various contexts should be more obvious at the call site, and not hidden imo.
e.g. I would expect to see, at the command level:
- full node's to pass
WriteStateto an sync-via-rpc thingy, and - sequencer to pass it.. to
apply_blocksomehow.
There was a problem hiding this comment.
Have updated to use a readonly State with BlockWriter and ProofWriter
| /// Readers are expected to be request-scoped, so snapshot lifetimes should be well under a block | ||
| /// interval. A lifetime exceeding [`SNAPSHOT_LIFETIME_WARN_THRESHOLD`] is logged at warn level. | ||
| pub(crate) struct SnapshotGuard { | ||
| live: Arc<AtomicUsize>, |
There was a problem hiding this comment.
This could probably be
| live: Arc<AtomicUsize>, | |
| live: Arc<()>, |
because count is tracked by arc already via Arc::strong_count.
There was a problem hiding this comment.
I think I prefer AtomicUsize for clarity. The strong_count approach feels a bit too implicit. And potentially our warn/debug logs could be misleading / erroneous regarding the live count if multiple drops occur concurrently?
There was a problem hiding this comment.
One would have to prematurely "drop" the arc e.g.:
let count = live.downgrade().strong_count();I figured you'd want to actually store something in the Arc eventually. But if you don't like relying on arc semantics, then lets not block on this. It just stood out to me that we were replicating arc internals.
There was a problem hiding this comment.
Thank you, I think this is really cool!
| if snapshots_live > SNAPSHOTS_LIVE_WARN_THRESHOLD { | ||
| tracing::warn!( | ||
| target: COMPONENT, | ||
| block_num = block_num.as_u32(), | ||
| snapshots.live = snapshots_live, | ||
| "too many live state snapshots; slow readers are pinning old generations", | ||
| ); | ||
| } |
There was a problem hiding this comment.
I like this.
I would like to explore some additional stuff we could diagnose here. For example if we could figure out a way to have each snapshot given the read request name as a static string.
And that here we could lookup the oldest snapshot and log that.
An alternative to this would be including a timeout task to each snapshot, which automatically does this logging instead. And the timeout task gets cancelled/aborted when the snapshot is dropped.
But lets leave that for a follow-up PR to explore this space.
There was a problem hiding this comment.
I didn't inspect this code too deeply; @kkovaacs if you could perhaps do a more thorough job here.
Mirko-von-Leipzig
left a comment
There was a problem hiding this comment.
Blocking merge until we are done with the demo
Summary
Closes #1539.
Closes #1853.
Closes #2414.
Makes the store's block-write path lock-free for readers, and makes the resulting read-consistency rules type-enforced. Reads previously contended on
RwLocks over the in-memory trees (and were blocked during the DB-commit window ofapply_block); they now load an immutable snapshot viaArcSwapand are never blocked by writes. All reads flow through a request-scopedStateView, so a query that combines tree and DB data at different chain heights is no longer expressible.Why:
apply_block(oneshot handshakes between the DB task and the in-memory update), which was fragile and hard to reason about.StateViewmakes the scoping structural.How:
Lock-free write path (
state/writer/):WriteWorkertask owns the mutable nullifier tree, account tree, blockchain MMR, and account-state forest, processing blocks serially from an mpsc channel — no locks anywhere. In-flight writes always complete: shutdown is only observed between requests.StateSnapshot(trees backed by read-only RocksDB snapshot views) and publishes it atomically viaArcSwap. Readers keep a consistent frozen view even while the next block commits.Db::apply_blockis now a plain transaction — the oneshotallow_acquire/acquire_donesynchronization is removed.Write capabilities (
state/lifecycle.rs):LoadedState::startspawns the worker and returns the read-onlyArc<State>plus non-cloneableBlockWriterandProofWritercapabilities and aWriterTaskjoin handle — statically limiting each write path to one task. The capabilities expose no read access; tasks that read and write receive theArc<State>alongside their capability.BlockWriter::stopdrains and joins the worker so tree storage is released deterministically before the data directory is re-opened or deleted (used byrecoverand stress-test seeding).Type-enforced reads (
state/view/):StateView, pinned to one snapshot per request (State::view()). DB queries are scoped by the view's tip internally; callers cannot supply their own. RocksDB-backed trees are only reachable throughblock_in_placehelpers; snapshot fields are only visible inside the view module.Dbqueries require view-issued proof types (ScopedBlockNum/ScopedBlockRange), constructible only by aStateViewafter validating the bound against its tip — extending the enforcement to the DB boundary itself.range.end() <= tipthemselves via a newRangeBeyondTiperror (sameInvalidArgumentresponse as before); the RPC layer'srange_bounds_checkis deleted and pagination'schain_tipis now the tip the query actually ran against.get_account, the block producer'sget_tx_inputs);sync_chain_mmrclamps the proven tip to the view's tip.Live tips (
state/tip.rs):State::committed_tip()/proven_tip()read the watch channels their writers publish to (mirroringsubscribe_committed_tip/subscribe_proven_tip); theFinalityenum is removed. The committed tip is published after the snapshot, so it never reports a block a fresh view cannot serve.Observability:
SnapshotGuardtracks live snapshot generations and lifetimes; warns when a snapshot outlives 10s or more than 4 generations are pinned (a leaked/slow reader pins a RocksDB snapshot).Supporting changes: read-only
reader()views forAccountStateForest/AccountTreeWithHistory(relaxed toBackendReader/SmtStorageReaderbounds);statemodule restructured intoview/(read endpoints) andwriter/(worker + capabilities); new tracing field names allowlisted.Changelog