diff --git a/crates/zakura-state/src/lib.rs b/crates/zakura-state/src/lib.rs index e9029f9a3..9b5ab119b 100644 --- a/crates/zakura-state/src/lib.rs +++ b/crates/zakura-state/src/lib.rs @@ -124,7 +124,7 @@ pub use service::read::{ }; pub use service::{ finalized_state::{DiskWriteBatch, FallibleDiskValue, FromDisk, IntoDisk, WriteDisk, ZakuraDb}, - ReadStateService, VctRootRepairState, VctRootRepairStatus, + OwnedBlockRange, ReadStateService, VctRootRepairState, VctRootRepairStatus, }; // Allow use in external tests diff --git a/crates/zakura-state/src/request.rs b/crates/zakura-state/src/request.rs index fe98bb139..e80139b7c 100644 --- a/crates/zakura-state/src/request.rs +++ b/crates/zakura-state/src/request.rs @@ -1900,6 +1900,9 @@ pub enum ReadRequest { /// Returns contiguous committed blocks by height, in ascending order. /// /// The response stops before the first height without a committed body. + /// Callers that charge resources to the database job should instead use + /// [`crate::ReadStateService::read_owned_block_range`] so cancellation of + /// the caller cannot release those resources during a running read. BlocksByHeightRange { /// First height to read. start: block::Height, diff --git a/crates/zakura-state/src/service.rs b/crates/zakura-state/src/service.rs index 34befb2a5..1593d0426 100644 --- a/crates/zakura-state/src/service.rs +++ b/crates/zakura-state/src/service.rs @@ -72,6 +72,7 @@ use crate::{ }; pub mod block_iter; +mod block_range; pub mod chain_tip; pub mod watch_receiver; @@ -91,6 +92,7 @@ pub mod arbitrary; #[cfg(test)] mod tests; +pub use block_range::OwnedBlockRange; pub use finalized_state::{OutputLocation, TransactionIndex, TransactionLocation}; use write::NonFinalizedWriteMessage; pub use write::{VctRootRepairState, VctRootRepairStatus}; diff --git a/crates/zakura-state/src/service/block_range.rs b/crates/zakura-state/src/service/block_range.rs new file mode 100644 index 000000000..198309c50 --- /dev/null +++ b/crates/zakura-state/src/service/block_range.rs @@ -0,0 +1,137 @@ +//! Height-range reads that keep resource reservations with the database work. +//! +//! Cancelling an async caller does not stop a blocking read already in progress. +//! For example, releasing a serving permit on disconnect could free a slot while +//! the database is still reading blocks for that request. These reads move the +//! reservation into the blocking job, then return it alongside the blocks. +//! +//! The generic resource type lets callers supply their own reservations without +//! making the state service depend on network serving policy. + +use std::sync::Arc; + +use futures::future::BoxFuture; +use tower::ServiceExt; +use tracing::Span; +use zakura_chain::{block, diagnostic::CodeTimer}; + +use super::{read, ReadStateService}; +use crate::{request::TimedSpan, BoxError, ReadRequest}; + +/// A bounded block prefix together with the caller's resource reservations. +/// +/// The blocking read transfers its resources into this result. Dropping the +/// async caller cannot release them while that read is running. Dropping an +/// undelivered result releases its blocks before its resources. +#[derive(Debug)] +pub struct OwnedBlockRange { + // Field order keeps resources alive while the retained blocks are dropped. + blocks: Vec<(block::Height, Arc, usize)>, + resources: R, +} + +impl OwnedBlockRange { + /// Transfer the blocks and their resources to the next owner. + /// + /// The caller must retain the resources for as long as its resource policy + /// requires, including while it holds or processes the returned blocks. + pub fn into_parts(self) -> (Vec<(block::Height, Arc, usize)>, R) { + (self.blocks, self.resources) + } +} + +impl ReadStateService { + /// Read a bounded contiguous prefix while the database job owns `resources`. + /// + /// This uses the same readiness checks, chain snapshot, and missing-block + /// behavior as [`ReadRequest::BlocksByHeightRange`], with an additional byte + /// cap. One blocking job transfers its resources into the returned result. + /// + /// `is_cancelled` is checked before the first lookup and between lookups. + /// Cancellation stops further lookups and returns the prefix already read. + /// It cannot interrupt a database lookup already in progress. Dropping or + /// aborting the caller discards delivery, but the job retains its resources + /// until it finishes and drops its undelivered result. + pub async fn read_owned_block_range( + &mut self, + start: block::Height, + count: u32, + max_response_bytes: u32, + resources: R, + is_cancelled: impl FnMut(&R) -> bool + Send + 'static, + ) -> Result, BoxError> { + self.ready().await?; + ReadRequest::BlocksByHeightRange { start, count }.count_metric(); + let state = self.clone(); + let best_chain = state.latest_best_chain(); + spawn_owned_block_range( + start, + count, + max_response_bytes, + resources, + is_cancelled, + move |height| read::block_and_size(best_chain.clone(), &state.db, height.into()), + ) + .await + } +} + +/// Spawn one blocking read that owns `resources` until it returns the blocks. +/// +/// Dropping the returned future stops waiting for the result, but does not stop +/// the blocking job. Capture `resources` inside that job so caller cancellation +/// cannot release reservations still needed by the read. Move them into +/// [`OwnedBlockRange`] on completion so they remain with the returned blocks. +fn spawn_owned_block_range( + start: block::Height, + count: u32, + max_response_bytes: u32, + resources: R, + mut is_cancelled: impl FnMut(&R) -> bool + Send + 'static, + mut get_block: impl FnMut(block::Height) -> Option<(Arc, usize)> + Send + 'static, +) -> BoxFuture<'static, Result, BoxError>> { + let timed_span = TimedSpan::new( + CodeTimer::start_desc("blocks_by_height_range"), + Span::current(), + ); + timed_span.spawn_blocking(move || { + let blocks = collect_bounded_height_range(start, count, max_response_bytes, |height| { + if is_cancelled(&resources) { + None + } else { + get_block(height) + } + }); + Ok(OwnedBlockRange { blocks, resources }) + }) +} + +/// Read a contiguous prefix without retaining more encoded block bytes than +/// the caller permits. +/// +/// The first block that does not fit can be materialized by the lookup, but is +/// dropped immediately and never enters the returned response. +fn collect_bounded_height_range( + start: block::Height, + count: u32, + max_response_bytes: u32, + mut get_block: impl FnMut(block::Height) -> Option<(T, usize)>, +) -> Vec<(block::Height, T, usize)> { + let mut response_bytes = 0u64; + (0..count) + .map_while(|offset| { + let height = start.0.checked_add(offset).map(block::Height)?; + let (block, size) = get_block(height)?; + let size_u64 = u64::try_from(size).ok()?; + let next_response_bytes = response_bytes.checked_add(size_u64)?; + if next_response_bytes > u64::from(max_response_bytes) { + return None; + } + response_bytes = next_response_bytes; + Some((height, block, size)) + }) + .collect() +} + +#[cfg(test)] +mod tests; diff --git a/crates/zakura-state/src/service/block_range/tests.rs b/crates/zakura-state/src/service/block_range/tests.rs new file mode 100644 index 000000000..81e51aac3 --- /dev/null +++ b/crates/zakura-state/src/service/block_range/tests.rs @@ -0,0 +1,375 @@ +use std::{ + panic::AssertUnwindSafe, + sync::{ + atomic::{AtomicBool, AtomicUsize, Ordering}, + mpsc, Arc, + }, + time::Duration, +}; + +use futures::FutureExt; +use tokio::{sync::oneshot, time::timeout}; +use zakura_chain::{ + block, + parameters::Network, + serialization::{ZcashDeserializeInto, ZcashSerialize}, +}; + +use super::spawn_owned_block_range; + +const DEADLINE: Duration = Duration::from_secs(10); + +#[derive(Debug)] +struct DropSignal { + drops: Arc, + finished: Option>, +} + +impl Drop for DropSignal { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::SeqCst); + if let Some(finished) = self.finished.take() { + let _ = finished.send(()); + } + } +} + +fn resources() -> (DropSignal, Arc, oneshot::Receiver<()>) { + let drops = Arc::new(AtomicUsize::new(0)); + let (finished, receiver) = oneshot::channel(); + ( + DropSignal { + drops: drops.clone(), + finished: Some(finished), + }, + drops, + receiver, + ) +} + +fn genesis() -> Arc { + Arc::new( + zakura_test::vectors::BLOCK_MAINNET_GENESIS_BYTES + .zcash_deserialize_into() + .expect("the genesis fixture is a serialized block"), + ) +} + +#[tokio::test] +async fn returned_blocks_retain_resources_and_respect_the_byte_cap() { + let _guard = zakura_test::init(); + let (resources, drops, _finished) = resources(); + let block = genesis(); + let reads = Arc::new(AtomicUsize::new(0)); + let worker_reads = reads.clone(); + let result = timeout( + DEADLINE, + spawn_owned_block_range( + block::Height(1), + 3, + 5, + resources, + |_| false, + move |_| { + worker_reads.fetch_add(1, Ordering::SeqCst); + // Synthetic sizes isolate the response cap from fixture size. + Some((block.clone(), 2)) + }, + ), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(reads.load(Ordering::SeqCst), 3); + assert_eq!( + result + .blocks + .iter() + .map(|(height, _, _)| *height) + .collect::>(), + [block::Height(1), block::Height(2)] + ); + assert_eq!(result.resources.drops.load(Ordering::SeqCst), 0); + let (blocks, resources) = result.into_parts(); + assert_eq!(drops.load(Ordering::SeqCst), 0); + drop(blocks); + drop(resources); + assert_eq!(drops.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn cancellation_before_the_first_lookup_skips_the_range() { + let _guard = zakura_test::init(); + let (resources, drops, _finished) = resources(); + let result = timeout( + DEADLINE, + spawn_owned_block_range( + block::Height(1), + 2, + 10, + resources, + |_| true, + |_| panic!("cancelled work must not reach its first lookup"), + ), + ) + .await + .unwrap() + .unwrap(); + assert!(result.blocks.is_empty()); + assert_eq!(drops.load(Ordering::SeqCst), 0); + drop(result); + assert_eq!(drops.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn cancellation_between_lookups_retains_the_completed_prefix() { + let _guard = zakura_test::init(); + let (resources, drops, _finished) = resources(); + let cancelled = Arc::new(AtomicBool::new(false)); + let worker_cancelled = cancelled.clone(); + let block = genesis(); + let result = timeout( + DEADLINE, + spawn_owned_block_range( + block::Height(1), + 2, + 10, + resources, + move |_| cancelled.load(Ordering::SeqCst), + move |_| { + assert!(!worker_cancelled.swap(true, Ordering::SeqCst)); + Some((block.clone(), 2)) + }, + ), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(result.blocks.len(), 1); + assert_eq!(drops.load(Ordering::SeqCst), 0); + drop(result); + assert_eq!(drops.load(Ordering::SeqCst), 1); +} + +struct BlockedRead { + started: oneshot::Receiver<()>, + resume: mpsc::SyncSender<()>, + drops: Arc, + finished: oneshot::Receiver<()>, +} + +impl BlockedRead { + async fn wait_started(&mut self) { + timeout(DEADLINE, &mut self.started).await.unwrap().unwrap(); + } + + async fn finish(self) { + assert_eq!(self.drops.load(Ordering::SeqCst), 0); + self.resume.send(()).unwrap(); + timeout(DEADLINE, self.finished).await.unwrap().unwrap(); + assert_eq!(self.drops.load(Ordering::SeqCst), 1); + } +} + +fn blocked_read() -> ( + futures::future::BoxFuture< + 'static, + Result, crate::BoxError>, + >, + BlockedRead, +) { + let (resources, drops, finished) = resources(); + let (started_tx, started) = oneshot::channel(); + let mut started_tx = Some(started_tx); + let (resume, blocked) = mpsc::sync_channel(1); + let block = genesis(); + let job = spawn_owned_block_range( + block::Height(1), + 1, + 10, + resources, + |_| false, + move |_| { + started_tx.take().unwrap().send(()).unwrap(); + blocked.recv_timeout(DEADLINE).unwrap(); + Some((block.clone(), 2)) + }, + ); + ( + job, + BlockedRead { + started, + resume, + drops, + finished, + }, + ) +} + +#[tokio::test] +async fn dropping_the_waiter_keeps_a_running_read_charged() { + let _guard = zakura_test::init(); + let (job, mut read) = blocked_read(); + read.wait_started().await; + drop(job); + read.finish().await; +} + +#[tokio::test] +async fn aborting_the_caller_keeps_a_running_read_charged() { + let _guard = zakura_test::init(); + let (job, mut read) = blocked_read(); + let caller = tokio::spawn(job); + read.wait_started().await; + caller.abort(); + assert!(timeout(DEADLINE, caller) + .await + .unwrap() + .unwrap_err() + .is_cancelled()); + read.finish().await; +} + +#[tokio::test] +async fn a_panicking_read_releases_resources_once() { + let _guard = zakura_test::init(); + let (resources, drops, finished) = resources(); + let job = spawn_owned_block_range( + block::Height(1), + 1, + 10, + resources, + |_| false, + |_| panic!("injected database panic"), + ); + let outcome = timeout(DEADLINE, AssertUnwindSafe(job).catch_unwind()) + .await + .unwrap(); + assert!(outcome.is_err()); + timeout(DEADLINE, finished).await.unwrap().unwrap(); + assert_eq!(drops.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn the_state_api_returns_an_owned_empty_range() { + let _guard = zakura_test::init(); + let (resources, drops, _finished) = resources(); + let (_state, mut read_state, _tip, _change) = + timeout(DEADLINE, crate::init_test_services(&Network::Mainnet)) + .await + .unwrap(); + let result = timeout( + DEADLINE, + read_state.read_owned_block_range(block::Height(1), 1, 10, resources, |_| false), + ) + .await + .unwrap() + .unwrap(); + assert!(result.blocks.is_empty()); + assert_eq!(drops.load(Ordering::SeqCst), 0); + drop(result); + assert_eq!(drops.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn the_state_api_returns_committed_blocks_with_their_resources() { + let _guard = zakura_test::init(); + let (resources, drops, _finished) = resources(); + let block = genesis(); + let size = block.zcash_serialized_size(); + let (_state, mut read_state, _tip, _change) = timeout( + DEADLINE, + crate::populated_state([block.clone()], &Network::Mainnet), + ) + .await + .unwrap(); + let result = timeout( + DEADLINE, + read_state.read_owned_block_range( + block::Height(0), + 2, + u32::try_from(size).unwrap(), + resources, + |_| false, + ), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(result.blocks, [(block::Height(0), block, size)]); + assert_eq!(drops.load(Ordering::SeqCst), 0); + drop(result); + assert_eq!(drops.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn readiness_failure_releases_resources_without_dispatching_a_read() { + use crate::service::write::{BlockWriteTaskFailure, HeaderChainAttachmentError}; + + let _guard = zakura_test::init(); + let (resources, drops, finished) = resources(); + let (_state, mut read_state, _tip, _change) = + timeout(DEADLINE, crate::init_test_services(&Network::Mainnet)) + .await + .unwrap(); + let failure = BlockWriteTaskFailure::from(&HeaderChainAttachmentError::MissingGenesis); + read_state.block_write_failure.set(failure.clone()).unwrap(); + let error = timeout( + DEADLINE, + read_state.read_owned_block_range(block::Height(1), 1, 10, resources, |_| { + panic!("a readiness failure must prevent dispatch of the blocking job") + }), + ) + .await + .unwrap() + .unwrap_err(); + assert_eq!(error.to_string(), failure.to_string()); + timeout(DEADLINE, finished).await.unwrap().unwrap(); + assert_eq!(drops.load(Ordering::SeqCst), 1); +} + +#[test] +fn block_range_response_stops_before_crossing_its_byte_limit() { + let sizes = [3usize, 2, 4]; + let exact = super::collect_bounded_height_range(block::Height(10), 3, 5, |height| { + let index = usize::try_from(height.0.checked_sub(10)?).ok()?; + sizes.get(index).copied().map(|size| (index, size)) + }); + assert_eq!( + exact, + vec![(block::Height(10), 0, 3), (block::Height(11), 1, 2)], + "the exact-fit prefix is returned and the first over-limit block is excluded", + ); + + let one_byte_short = super::collect_bounded_height_range(block::Height(10), 3, 4, |height| { + let index = usize::try_from(height.0.checked_sub(10)?).ok()?; + sizes.get(index).copied().map(|size| (index, size)) + }); + assert_eq!( + one_byte_short, + vec![(block::Height(10), 0, 3)], + "a prefix that would cross the cap by one byte stops before that block", + ); + + let first_too_large = super::collect_bounded_height_range(block::Height(10), 3, 2, |height| { + let index = usize::try_from(height.0.checked_sub(10)?).ok()?; + sizes.get(index).copied().map(|size| (index, size)) + }); + assert!( + first_too_large.is_empty(), + "a first block larger than the response budget is not retained", + ); +} + +#[test] +fn block_range_response_includes_a_maximum_size_first_block() { + let maximum = usize::try_from(block::MAX_BLOCK_BYTES).unwrap(); + let cap = u32::try_from(block::MAX_BLOCK_BYTES).unwrap(); + let result = super::collect_bounded_height_range(block::Height(10), 2, cap, |height| { + Some((height, maximum)) + }); + assert_eq!( + result, + vec![(block::Height(10), block::Height(10), maximum)] + ); +} diff --git a/docs/changelog/unreleased/942.md b/docs/changelog/unreleased/942.md new file mode 100644 index 000000000..d65159650 --- /dev/null +++ b/docs/changelog/unreleased/942.md @@ -0,0 +1,4 @@ + + +Add an internal bounded state-read API that retains caller resources through +blocking work and result ownership; production callers are wired separately.