diff --git a/crates/zakura-network/src/zakura/block_sync/peer_routine.rs b/crates/zakura-network/src/zakura/block_sync/peer_routine.rs index 838479bc42..2cd6b14c79 100644 --- a/crates/zakura-network/src/zakura/block_sync/peer_routine.rs +++ b/crates/zakura-network/src/zakura/block_sync/peer_routine.rs @@ -45,6 +45,7 @@ use super::{ BlockSyncMessage, BlockSyncMisbehavior, BlockSyncPeerSession, BlockSyncStatus, ZakuraBlockSyncConfig, ZakuraPeerId, ZakuraTrace, MSG_BS_BLOCK, }; +use crate::zakura::transport::OrderedStreamFailure; use crate::zakura::{ trace::BlockBodySource, Admit, FramedRecv, OrderedSendError, SinkReject, ZakuraConnId, }; @@ -413,6 +414,27 @@ impl PeerRoutine { /// reject. A reject returns `Err(SinkReject::protocol(..))` so the supervised /// pipe tears the whole connection down. pub(super) async fn run(mut self) -> Result<(), SinkReject> { + let mut guard = block_sync_guard(); + let result = self.run_inner(&mut guard).await; + // A transport failure can cancel the session before its queued responses + // reach us. Validate them under the existing decode-capacity bound before + // scoring unanswered work. Connection shutdown can still stop this drain. + if result.is_ok() { + if let Some(failure) = self.recv.failure() { + self.recv.close(); + while let Ok(frame) = self.recv.try_recv() { + self.handle_frame(&mut guard, frame).await?; + } + return self.handle_stream_failure(Instant::now(), failure); + } + } + result + } + + async fn run_inner( + &mut self, + guard: &mut crate::zakura::SessionGuard, + ) -> Result<(), SinkReject> { // Local clones so the `Notified` futures below borrow these handles, not // `self` — `self.try_fill()` needs `&mut self` while the notifications are // pinned. The clones share the same underlying `Arc`, so the wakes still @@ -420,13 +442,14 @@ impl PeerRoutine { // `self.work`. let budget = self.budget.clone(); let work = self.work.clone(); - // The per-connection oversize guard applied to inbound frames at ingress. - let mut guard = block_sync_guard(); // Per-peer BBR heartbeat cadence. `Skip` so a routine busy past a tick emits one // fresh sample rather than a catch-up burst. Observability only. let mut bbr_trace_ticks = time::interval(BBR_TRACE_INTERVAL); bbr_trace_ticks.set_missed_tick_behavior(time::MissedTickBehavior::Skip); loop { + if self.cancel.is_cancelled() { + return Ok(()); + } // missed-wake safety: register both `Notify`s via // `Notified::enable()` BEFORE the fill attempt. The budget/work // `Notify`s use `notify_waiters` (no stored permit), so a @@ -470,12 +493,15 @@ impl PeerRoutine { // in this same task. A protocol reject propagates out so // the supervised pipe cancels the connection; the `Drop` // guard returns unreceived work on the way out. - Some(frame) => self.handle_frame(&mut guard, frame).await?, + Some(frame) => self.handle_frame(guard, frame).await?, // Stream closed by the peer. With no outstanding work this // is a clean exit; with unanswered requests it is a // no-progress stall (park or disconnect). `Drop` returns // unreceived outstanding heights and releases their budget. - None => return self.handle_remote_stream_closed(Instant::now()), + None => return self.handle_stream_failure( + Instant::now(), + self.recv.failure().unwrap_or(OrderedStreamFailure::RemoteClose), + ), } } changed = self.sequencer_view.changed() => { @@ -525,7 +551,20 @@ impl PeerRoutine { let frame_payload_bytes = frame.payload.len(); let body_permit = if is_block_frame(&frame) { - Some(self.reserve_body_decode_permit().await?) + let permit = self.reserve_body_decode_permit(); + tokio::pin!(permit); + Some(tokio::select! { + biased; + () = self.cancel.cancelled() => { + if self.recv.failure().is_none() { + return Ok(()); + } + // A remote close cannot excuse an unvalidated body. Keep + // this bounded raw frame until its decode slot is available. + permit.await? + } + permit = &mut permit => permit?, + }) } else { None }; @@ -1340,29 +1379,29 @@ impl PeerRoutine { Err(SinkReject::local(error)) } - /// Handle the peer closing its send side of the stream. Honest peers close - /// *connections*, not lone streams; a stream-only EOF while block-progress - /// liveness is still armed is the same signal as the liveness stall and must - /// not reset the park/second-stall state machine — otherwise a peer could - /// take work, deliver nothing, let its requests time out (which drains - /// `outstanding` without disarming liveness), EOF before the liveness - /// deadline, and be readmitted fresh forever. An armed deadline with empty - /// `outstanding` can only mean charged requests were consumed without - /// accepted progress — every answered-everything path disarms it — and - /// `DownloadWindow::check_liveness` parks at that deadline - /// regardless of `outstanding`, so parking here only moves the already - /// scheduled outcome earlier. Frames are processed in-order in this task, so - /// at EOF everything the peer sent has already been counted. The liveness - /// grace does not apply: it waits for in-flight frames stuck behind our full - /// outbound queue, and a closed stream has none. - fn handle_remote_stream_closed(&mut self, now: Instant) -> Result<(), SinkReject> { + /// Retiring a failed session must not erase unanswered download work. + /// Request expiry can empty `outstanding` without disarming liveness, so + /// that pending no-progress deadline must also survive stream replacement. + fn handle_stream_failure( + &mut self, + now: Instant, + failure: OrderedStreamFailure, + ) -> Result<(), SinkReject> { + // Local reset may have withdrawn the requests while decoding waited + // for capacity. It must remain neutral when the failed stream retires. + self.on_view_changed(); if self.window.outstanding.is_empty() && self.window.block_liveness_deadline.is_none() { return Ok(()); } - self.no_progress_stall( - now, - "block-sync peer closed the stream with its requests unanswered", - ) + let error = match failure { + OrderedStreamFailure::RemoteClose => { + "block-sync peer closed the stream with its requests unanswered" + } + OrderedStreamFailure::WriteTimeout => { + "block-sync stream write stalled with download requests unanswered" + } + }; + self.no_progress_stall(now, error) } /// Drop this routine's outstanding requests whose whole range is at or below @@ -2672,6 +2711,423 @@ mod tests { )); } + #[tokio::test] + async fn cancelled_session_settles_peer_failures_before_reopening() { + for failure in [ + super::OrderedStreamFailure::RemoteClose, + super::OrderedStreamFailure::WriteTimeout, + ] { + check_cancelled_session_policy(Some(failure), true, false, false).await; + check_cancelled_session_policy(Some(failure), true, false, true).await; + check_cancelled_session_policy(Some(failure), true, true, false).await; + } + } + + #[tokio::test] + async fn local_cancel_and_idle_write_timeout_do_not_park_peers() { + check_cancelled_session_policy(None, true, false, false).await; + for buffered in [ + BufferedResponse::Backpressured, + BufferedResponse::AlreadyProcessing, + ] { + check_cancelled_session_with_buffered_response(None, true, false, false, buffered) + .await; + } + check_cancelled_session_policy( + Some(super::OrderedStreamFailure::WriteTimeout), + false, + false, + false, + ) + .await; + } + + async fn check_cancelled_session_policy( + failure: Option, + unanswered: bool, + readmitted: bool, + expired: bool, + ) { + check_cancelled_session_with_buffered_response( + failure, + unanswered, + readmitted, + expired, + BufferedResponse::None, + ) + .await; + } + + #[derive(Clone, Copy, PartialEq)] + enum BufferedResponse { + None, + Complete, + CompleteWhileOutboundFull, + StatusOnly, + Malformed, + MalformedBlock, + MalformedBlockBackpressured, + MalformedBlockAlreadyProcessing, + Backpressured, + AlreadyProcessing, + ResetDuringBackpressure, + } + + #[tokio::test] + async fn malformed_blocks_are_validated_before_stream_failure_settlement() { + for failure in [ + super::OrderedStreamFailure::RemoteClose, + super::OrderedStreamFailure::WriteTimeout, + ] { + for readmitted in [false, true] { + for buffered in [ + BufferedResponse::MalformedBlock, + BufferedResponse::MalformedBlockBackpressured, + BufferedResponse::MalformedBlockAlreadyProcessing, + ] { + check_cancelled_session_with_buffered_response( + Some(failure), + true, + readmitted, + false, + buffered, + ) + .await; + } + } + } + } + + #[tokio::test] + async fn buffered_responses_are_considered_before_stream_failure() { + for failure in [ + super::OrderedStreamFailure::RemoteClose, + super::OrderedStreamFailure::WriteTimeout, + ] { + for readmitted in [false, true] { + for buffered in [ + BufferedResponse::Complete, + BufferedResponse::CompleteWhileOutboundFull, + BufferedResponse::StatusOnly, + BufferedResponse::Malformed, + BufferedResponse::Backpressured, + BufferedResponse::ResetDuringBackpressure, + ] { + check_cancelled_session_with_buffered_response( + Some(failure), + true, + readmitted, + false, + buffered, + ) + .await; + } + } + } + } + + #[tokio::test] + async fn cancellation_during_local_body_backpressure_does_not_charge_a_stall() { + for readmitted in [false, true] { + check_cancelled_session_with_buffered_response( + Some(super::OrderedStreamFailure::RemoteClose), + true, + readmitted, + false, + BufferedResponse::AlreadyProcessing, + ) + .await; + } + } + + async fn check_cancelled_session_with_buffered_response( + failure: Option, + unanswered: bool, + readmitted: bool, + expired: bool, + buffered: BufferedResponse, + ) { + use super::super::peer_registry::SessionAdmission; + use crate::zakura::transport::OrderedStreamFailureCause; + use crate::zakura::{ServicePeerDirection, SinkReject}; + + use zakura_chain::serialization::ZcashDeserializeInto; + let body: Arc = Arc::new( + zakura_test::vectors::BLOCK_MAINNET_1_BYTES + .zcash_deserialize_into() + .unwrap(), + ); + let config = ZakuraBlockSyncConfig::default(); + let budget = ByteBudget::new(1_000_000); + let work = Arc::new(WorkQueue::new(block::Height(0))); + work.set_estimate_floor_for_tests(1); + if unanswered { + work.extend( + super::super::test_work_scope(), + [( + block::Height(1), + body.hash(), + BlockSizeEstimate::Confirmed( + u32::try_from(zakura_test::vectors::BLOCK_MAINNET_1_BYTES.len()).unwrap(), + ), + )], + ); + } + let peer = ZakuraPeerId::new(vec![19; 32]).unwrap(); + let registry = Arc::new(PeerRegistry::new()); + let now = Instant::now(); + let mut generation = registry + .admit_session(&peer, ServicePeerDirection::Outbound, &config, 0, now) + .generation(); + if readmitted { + assert!(registry.park_session(&peer, 0, generation, now)); + let admission = + registry.admit_session(&peer, ServicePeerDirection::Outbound, &config, 0, now); + assert!(matches!(admission, SessionAdmission::Readmitted { .. })); + generation = admission.generation(); + } + let cancel = CancellationToken::new(); + let (out_send, mut out_recv) = crate::zakura::transport::worker_framed_channel(4); + let fill_outbound = out_send.clone(); + let (in_send, in_recv) = framed_channel(4); + let cause = OrderedStreamFailureCause::default(); + let session = BlockSyncPeerSession::for_test(peer.clone(), out_send, cancel.clone()); + let (sequencer_input, mut sequencer_recv) = mpsc::channel(4); + let mut held_capacity = Vec::new(); + if matches!( + buffered, + BufferedResponse::Backpressured + | BufferedResponse::AlreadyProcessing + | BufferedResponse::ResetDuringBackpressure + | BufferedResponse::MalformedBlockBackpressured + | BufferedResponse::MalformedBlockAlreadyProcessing + ) { + for _ in 0..4 { + held_capacity.push(sequencer_input.clone().try_reserve_owned().unwrap()); + } + } + let (reactor_input, _reactor_recv) = mpsc::channel(4); + let (view_tx, view_rx) = watch::channel(initial_view(BlockSyncFrontiers { + finalized_height: block::Height(0), + verified_block_tip: block::Height(0), + verified_block_hash: block::Hash([0; 32]), + })); + let mut routine = PeerRoutine::new( + peer.clone(), + 0, + session, + in_recv.with_failure_cause(cause.clone()), + config.clone(), + !readmitted, + generation, + budget.clone(), + work.clone(), + registry.clone(), + Arc::new(Mutex::new(ThroughputMeter::new(now))), + sequencer_input, + Arc::new(AtomicU64::new(0)), + Arc::new(AtomicU64::new(0)), + reactor_input, + view_rx, + cancel.clone(), + ZakuraTrace::noop(), + ); + if unanswered { + routine.handle_status(super::BlockSyncStatus { + servable_low: block::Height(1), + servable_high: block::Height(1), + max_blocks_per_response: 1, + ..super::BlockSyncStatus::default() + }); + routine.try_fill().await; + assert_eq!(routine.window.outstanding.len(), 1); + timeout(Duration::from_secs(1), out_recv.recv()) + .await + .unwrap() + .unwrap() + .write_with(|_| async { Ok::<_, std::convert::Infallible>(()) }) + .await + .unwrap(); + assert!(work.in_flight_contains(block::Height(1))); + if expired { + let after_deadline = + routine.window.outstanding[0].deadline + Duration::from_millis(1); + routine.expire_due_timeouts(after_deadline); + assert!(routine.window.outstanding.is_empty()); + assert!(routine.window.block_liveness_deadline.is_some()); + } + } + match buffered { + BufferedResponse::Complete + | BufferedResponse::CompleteWhileOutboundFull + | BufferedResponse::Backpressured + | BufferedResponse::AlreadyProcessing + | BufferedResponse::ResetDuringBackpressure => { + in_send + .send( + super::BlockSyncMessage::Block(body.clone()) + .encode_frame() + .unwrap(), + ) + .await + .unwrap(); + in_send + .send( + super::BlockSyncMessage::BlocksDone { + start_height: block::Height(1), + returned: 1, + } + .encode_frame() + .unwrap(), + ) + .await + .unwrap(); + } + BufferedResponse::StatusOnly => { + in_send + .send( + super::BlockSyncMessage::Status(super::BlockSyncStatus::default()) + .encode_frame() + .unwrap(), + ) + .await + .unwrap(); + } + BufferedResponse::Malformed => { + in_send + .send(crate::zakura::Frame { + message_type: u16::MAX, + flags: 0, + payload: vec![], + }) + .await + .unwrap(); + } + BufferedResponse::MalformedBlock + | BufferedResponse::MalformedBlockBackpressured + | BufferedResponse::MalformedBlockAlreadyProcessing => { + let frame = crate::zakura::Frame { + message_type: u16::from(super::MSG_BS_BLOCK), + flags: 0, + payload: vec![super::MSG_BS_BLOCK], + }; + assert!(super::BlockSyncMessage::decode_frame(frame.clone()).is_err()); + in_send.send(frame).await.unwrap(); + } + BufferedResponse::None => {} + } + if buffered == BufferedResponse::CompleteWhileOutboundFull { + for _ in 0..4 { + fill_outbound + .try_send( + super::BlockSyncMessage::Status(super::BlockSyncStatus::default()) + .encode_frame() + .unwrap(), + ) + .unwrap(); + } + assert_eq!(routine.session.outbound_capacity(), 0); + } + let mut running = Box::pin(routine.run()); + if matches!( + buffered, + BufferedResponse::AlreadyProcessing | BufferedResponse::MalformedBlockAlreadyProcessing + ) { + // Poll until the body leaves the frame queue and waits for local + // capacity. Cancellation must not turn that wait into a peer fault. + assert!(futures::poll!(running.as_mut()).is_pending()); + assert_eq!( + in_send.capacity(), + if buffered == BufferedResponse::MalformedBlockAlreadyProcessing { + 4 + } else { + 3 + } + ); + } + if let Some(failure) = failure { + cause.record(failure); + } + cancel.cancel(); + if failure.is_some() && !held_capacity.is_empty() { + assert!( + futures::poll!(running.as_mut()).is_pending(), + "a failed stream must retain its unvalidated response until decode capacity returns" + ); + assert!(work.in_flight_contains(block::Height(1))); + assert!(budget.reserved() > 0); + assert!(sequencer_recv.try_recv().is_err()); + assert!(!registry.is_peer_parked(&peer, Instant::now())); + if buffered == BufferedResponse::ResetDuringBackpressure { + budget.clone().release(work.reset_above(block::Height(0))); + view_tx.send_modify(|view| view.reset_epoch += 1); + } + drop(held_capacity); + } + let result = timeout(Duration::from_secs(1), running).await.unwrap(); + if matches!( + buffered, + BufferedResponse::Malformed + | BufferedResponse::MalformedBlock + | BufferedResponse::MalformedBlockBackpressured + | BufferedResponse::MalformedBlockAlreadyProcessing + ) { + assert!(matches!(result, Err(SinkReject::Protocol(_))), "{result:?}"); + assert!(!registry.is_peer_parked(&peer, Instant::now())); + } else if failure.is_some() + && unanswered + && matches!( + buffered, + BufferedResponse::None | BufferedResponse::StatusOnly + ) + { + if readmitted { + assert!(matches!(result, Err(SinkReject::Protocol(_))), "{result:?}"); + } else { + assert!(matches!(result, Err(SinkReject::Local(_))), "{result:?}"); + assert!(matches!( + registry.admit_session( + &peer, + ServicePeerDirection::Outbound, + &config, + 0, + Instant::now() + ), + SessionAdmission::Parked, + )); + assert!(matches!( + registry.admit_session( + &peer, + ServicePeerDirection::Outbound, + &config, + 0, + Instant::now() + config.effective_no_progress_peer_cooldown() + ), + SessionAdmission::Readmitted { .. }, + )); + } + } else { + assert!(result.is_ok(), "{result:?}"); + assert!(!registry.is_peer_parked(&peer, Instant::now())); + } + assert_eq!(budget.reserved(), 0); + let received = failure.is_some() + && matches!( + buffered, + BufferedResponse::Complete + | BufferedResponse::CompleteWhileOutboundFull + | BufferedResponse::Backpressured + | BufferedResponse::AlreadyProcessing + ); + assert_eq!( + work.pending_contains(block::Height(1)), + unanswered && !received && buffered != BufferedResponse::ResetDuringBackpressure + ); + assert_eq!(work.in_flight_contains(block::Height(1)), received); + assert_eq!(sequencer_recv.try_recv().is_ok(), received); + assert!(sequencer_recv.try_recv().is_err()); + } + #[test] fn repeated_no_progress_stall_disconnects_instead_of_parking_again() { assert_eq!( diff --git a/crates/zakura-network/src/zakura/block_sync/service.rs b/crates/zakura-network/src/zakura/block_sync/service.rs index 3f7b655c8f..143d945c6e 100644 --- a/crates/zakura-network/src/zakura/block_sync/service.rs +++ b/crates/zakura-network/src/zakura/block_sync/service.rs @@ -1,9 +1,9 @@ use super::{config::*, events::*, peer_registry::SessionAdmission, wire::*, *}; use crate::zakura::{ - handle_pipe_exit, spawn_supervised_pipe, FramedRecv, FramedSend, OrderedSendError, - OrderedSessionDemand, OrderedStreamOpening, OrderedStreamPolicy, Peer, PeerStreamSession, - Service, ServicePeerSnapshot, SinkReject, Stream, StreamMode, ZakuraBlockSyncCandidateState, - ZakuraConnId, ZakuraPeerId, FRAME_HEADER_BYTES, + handle_pipe_exit, spawn_supervised_pipe, FramedRecv, FramedSend, OrderedSendError, Peer, + PeerStreamSession, Service, ServicePeerSnapshot, SessionDemand, SessionOpening, SessionPolicy, + SinkReject, Stream, StreamMode, ZakuraBlockSyncCandidateState, ZakuraConnId, ZakuraPeerId, + FRAME_HEADER_BYTES, }; use std::{ sync::atomic::{AtomicU64, Ordering}, @@ -26,7 +26,7 @@ const BLOCK_SYNC_SERVICE_STREAMS: [Stream; 1] = [Stream { version: ZAKURA_BLOCK_SYNC_STREAM_VERSION, frame_cap: MAX_BS_FRAME_BYTES, capability: ZAKURA_CAP_BLOCK_SYNC, - mode: StreamMode::Ordered, + mode: StreamMode::Persistent, }]; /// Service-declared streams for native block sync. @@ -451,28 +451,28 @@ impl Service for BlockSyncService { block_sync_streams() } - fn ordered_stream_policy(&self, _kind: u16) -> OrderedStreamPolicy { - OrderedStreamPolicy { - opening: OrderedStreamOpening::EitherSide, + fn session_policy(&self) -> SessionPolicy { + SessionPolicy { + opening: SessionOpening::EitherSide, reopen: true, } } - fn ordered_session_demand( + fn session_demand( &self, conn_id: ZakuraConnId, peer: &ZakuraPeerId, _negotiated: u64, direction: ServicePeerDirection, - ) -> OrderedSessionDemand { + ) -> SessionDemand { if let Some(deadline) = self.peer_park_deadline(peer) { - return OrderedSessionDemand::RetryAt(deadline); + return SessionDemand::RetryAt(deadline); } let mut peer_snapshot = self.inner.peer_snapshot.clone(); peer_snapshot.borrow_and_update(); if !self.peer_slots_free(direction) { - return OrderedSessionDemand::WaitForChange(Box::pin(async move { + return SessionDemand::WaitForChange(Box::pin(async move { if peer_snapshot.changed().await.is_err() { std::future::pending::<()>().await; } @@ -495,7 +495,7 @@ impl Service for BlockSyncService { .is_empty() { let mut service_demand = self.service_demand.clone(); - return OrderedSessionDemand::WaitForChange(Box::pin(async move { + return SessionDemand::WaitForChange(Box::pin(async move { if let Some(demand) = service_demand.as_mut() { tokio::select! { changed = candidates.changed() => { @@ -516,7 +516,7 @@ impl Service for BlockSyncService { } } - OrderedSessionDemand::OpenNow + SessionDemand::OpenNow } fn wants_peer( @@ -711,7 +711,11 @@ impl Service for BlockSyncService { run_cancel, wiring.trace, ); - routine.run().await + tokio::select! { + biased; + () = connection_cancel_token.cancelled() => Ok(()), + result = routine.run() => result, + } } None => drain_inbound(recv, run_cancel).await, }; @@ -760,8 +764,8 @@ impl Service for BlockSyncService { return false; }; matches!( - self.ordered_session_demand(conn_id, peer, ZAKURA_CAP_BLOCK_SYNC, direction), - OrderedSessionDemand::OpenNow + self.session_demand(conn_id, peer, ZAKURA_CAP_BLOCK_SYNC, direction), + SessionDemand::OpenNow ) } diff --git a/crates/zakura-network/src/zakura/block_sync/tests.rs b/crates/zakura-network/src/zakura/block_sync/tests.rs index 462da5ca72..c113f8f403 100644 --- a/crates/zakura-network/src/zakura/block_sync/tests.rs +++ b/crates/zakura-network/src/zakura/block_sync/tests.rs @@ -32,8 +32,8 @@ use crate::zakura::{ framed_channel, testkit::{await_until, TraceCapture, TraceValue}, trace::BlockBodySource, - FramedRecv, FramedSend, OrderedSessionDemand, Peer, Service, ServicePeerSnapshot, - ServiceRegistry, StreamMode, ZakuraBlockSyncCandidateState, + FramedRecv, FramedSend, Peer, Service, ServicePeerSnapshot, ServiceRegistry, SessionDemand, + StreamMode, ZakuraBlockSyncCandidateState, }; use zakura_chain::{ fmt::HexDebug, @@ -6232,7 +6232,7 @@ fn block_sync_stream_declares_kind_capability_version_and_frame_cap() { assert_eq!(stream.kind, ZAKURA_STREAM_BLOCK_SYNC); assert_eq!(stream.version, ZAKURA_BLOCK_SYNC_STREAM_VERSION); assert_eq!(stream.capability, ZAKURA_CAP_BLOCK_SYNC); - assert_eq!(stream.mode, StreamMode::Ordered); + assert_eq!(stream.mode, StreamMode::Persistent); assert_eq!(stream.frame_cap, MAX_BS_FRAME_BYTES); } @@ -6255,14 +6255,14 @@ async fn service_registry_routes_block_sync_by_exact_capability_and_version() { .is_none()); assert_eq!( registry - .ordered_streams_for_negotiated(ZAKURA_CAP_BLOCK_SYNC) + .persistent_streams_for_negotiated(ZAKURA_CAP_BLOCK_SYNC) .iter() .map(|stream| stream.kind) .collect::>(), vec![ZAKURA_STREAM_BLOCK_SYNC] ); - assert!(registry.ordered_streams_for_negotiated(0).is_empty()); - assert!(registry.wants_ordered_stream( + assert!(registry.persistent_streams_for_negotiated(0).is_empty()); + assert!(registry.wants_session( ZAKURA_STREAM_BLOCK_SYNC, ZAKURA_CAP_BLOCK_SYNC, &peer, @@ -6536,6 +6536,78 @@ async fn add_peer_decode_failure_reports_malformed_and_cancels_connection() { .expect("malformed frame cancels the connection"); } +#[tokio::test] +async fn add_peer_connection_shutdown_cancels_pending_block_validation() { + use crate::zakura::transport::{OrderedStreamFailure, OrderedStreamFailureCause}; + + let config = ZakuraBlockSyncConfig::default(); + let (_tip_tx, tip_rx) = watch::channel((block::Height(0), block::Hash([0; 32]))); + let startup = BlockSyncStartup::new( + BlockSyncFrontiers { + finalized_height: block::Height(0), + verified_block_tip: block::Height(0), + verified_block_hash: block::Hash([0; 32]), + }, + (block::Height(0), block::Hash([0; 32])), + tip_rx, + config.clone(), + ); + let (handle, _actions, _reactor_task) = spawn_block_sync_reactor(startup); + let input = &handle.routine_wiring.as_ref().unwrap().sequencer_input; + let held_capacity: Vec<_> = (0..input.max_capacity()) + .map(|_| input.clone().try_reserve_owned().unwrap()) + .collect(); + let service = BlockSyncService::new_with_handle_for_test(config, handle.clone()); + let (inbound_tx, inbound_rx) = framed_channel(4); + let (outbound_tx, _outbound_rx) = framed_channel(4); + let cause = OrderedStreamFailureCause::default(); + let streams = HashMap::from([( + ZAKURA_STREAM_BLOCK_SYNC, + (inbound_rx.with_failure_cause(cause.clone()), outbound_tx), + )]); + let connection_cancel = CancellationToken::new(); + let remote = Peer::new( + peer(3), + None, + ZAKURA_CAP_BLOCK_SYNC, + streams, + connection_cancel.clone(), + ); + let session_cancel = remote.service_cancel_token(); + inbound_tx + .send(Frame { + message_type: u16::from(MSG_BS_BLOCK), + flags: 0, + payload: vec![MSG_BS_BLOCK], + }) + .await + .unwrap(); + service.add_peer(remote); + cause.record(OrderedStreamFailure::RemoteClose); + session_cancel.cancel(); + + tokio::time::timeout(Duration::from_secs(1), async { + while inbound_tx.capacity() != inbound_tx.max_capacity() { + tokio::task::yield_now().await; + } + }) + .await + .expect("the failed session takes its pending frame for validation"); + assert_eq!(service.peer_count(), 1, "validation still owns the session"); + assert!(!connection_cancel.is_cancelled()); + + connection_cancel.cancel(); + tokio::time::timeout(Duration::from_secs(1), async { + while service.peer_count() != 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("connection shutdown releases the session without decode capacity"); + assert_eq!(input.capacity(), 0); + drop(held_capacity); +} + #[tokio::test] async fn registry_add_peer_requires_negotiated_block_sync_capability() { let (service, mut events) = BlockSyncService::new_for_test(ZakuraBlockSyncConfig::default()); @@ -15311,24 +15383,24 @@ async fn parked_connection_cleanup_allows_a_fresh_connection_after_cooldown() { handle.park_session_for_test(&peer, old_conn_id, Duration::ZERO); assert!(matches!( - service.ordered_session_demand( + service.session_demand( old_conn_id, &peer, ZAKURA_CAP_BLOCK_SYNC, ServicePeerDirection::Outbound, ), - OrderedSessionDemand::WaitForChange(_), + SessionDemand::WaitForChange(_), )); service.remove_peer(&peer, old_conn_id); assert!(matches!( - service.ordered_session_demand( + service.session_demand( new_conn_id, &peer, ZAKURA_CAP_BLOCK_SYNC, ServicePeerDirection::Outbound, ), - OrderedSessionDemand::OpenNow + SessionDemand::OpenNow )); reactor_task.abort(); } @@ -15353,13 +15425,13 @@ async fn same_connection_block_sync_session_waits_at_tip_then_reopens_for_new_wo let conn_id = 17; handle.park_session_for_test(&peer, conn_id, Duration::ZERO); - let demand = service.ordered_session_demand( + let demand = service.session_demand( conn_id, &peer, ZAKURA_CAP_BLOCK_SYNC, ServicePeerDirection::Outbound, ); - let OrderedSessionDemand::WaitForChange(changed) = demand else { + let SessionDemand::WaitForChange(changed) = demand else { panic!("a locally parked session must stay absent while block sync is at tip"); }; @@ -15375,13 +15447,13 @@ async fn same_connection_block_sync_session_waits_at_tip_then_reopens_for_new_wo .expect("new block work wakes the parked session demand"); assert!(matches!( - service.ordered_session_demand( + service.session_demand( conn_id, &peer, ZAKURA_CAP_BLOCK_SYNC, ServicePeerDirection::Outbound, ), - OrderedSessionDemand::OpenNow, + SessionDemand::OpenNow, )); reactor_task.abort(); } @@ -15420,7 +15492,7 @@ async fn serving_only_coordinator_demand_keeps_block_session_available_during_fa let conn_id = 18; handle.park_session_for_test(&peer, conn_id, Duration::ZERO); - let OrderedSessionDemand::WaitForChange(changed) = service.ordered_session_demand( + let SessionDemand::WaitForChange(changed) = service.session_demand( conn_id, &peer, ZAKURA_CAP_BLOCK_SYNC, @@ -15440,13 +15512,13 @@ async fn serving_only_coordinator_demand_keeps_block_session_available_during_fa .await .expect("fallback service demand wakes the parked ordered session"); assert!(matches!( - service.ordered_session_demand( + service.session_demand( conn_id, &peer, ZAKURA_CAP_BLOCK_SYNC, ServicePeerDirection::Outbound, ), - OrderedSessionDemand::OpenNow, + SessionDemand::OpenNow, )); reactor_task.abort(); } diff --git a/crates/zakura-network/src/zakura/discovery/service.rs b/crates/zakura-network/src/zakura/discovery/service.rs index 5c6d490774..f3ecaca450 100644 --- a/crates/zakura-network/src/zakura/discovery/service.rs +++ b/crates/zakura-network/src/zakura/discovery/service.rs @@ -24,9 +24,9 @@ use tokio_util::sync::CancellationToken; use crate::zakura::{ handle_pipe_exit, spawn_supervised_peer_task, spawn_supervised_pipe, BlockSyncHandle, CloseCause, Event, Flow, Frame, FramedRecv, FramedSend, HeaderSyncHandle, OrderedSendError, - OrderedSessionDemand, OrderedStreamOpening, OrderedStreamPolicy, Peer, PeerStreamSession, Pipe, - Service, ServiceAdmissionDecision, ServicePeerDirection, SinkReject, Stream, StreamMode, - ZakuraConnId, ZakuraPeerId, LOCAL_MAX_CONTROL_FRAME_BYTES, ZAKURA_CAP_DISCOVERY, + Peer, PeerStreamSession, Pipe, Service, ServiceAdmissionDecision, ServicePeerDirection, + SessionDemand, SessionOpening, SessionPolicy, SinkReject, Stream, StreamMode, ZakuraConnId, + ZakuraPeerId, LOCAL_MAX_CONTROL_FRAME_BYTES, ZAKURA_CAP_DISCOVERY, }; #[cfg(test)] @@ -47,7 +47,7 @@ const DISCOVERY_SERVICE_STREAMS: [Stream; 1] = [Stream { version: ZAKURA_DISCOVERY_STREAM_VERSION, frame_cap: LOCAL_MAX_CONTROL_FRAME_BYTES, capability: ZAKURA_CAP_DISCOVERY, - mode: StreamMode::Ordered, + mode: StreamMode::Persistent, }]; /// Service-declared streams for native discovery. @@ -305,20 +305,20 @@ impl Service for DiscoveryService { discovery_streams() } - fn ordered_stream_policy(&self, _kind: u16) -> OrderedStreamPolicy { - OrderedStreamPolicy { - opening: OrderedStreamOpening::InitiatorOnly, + fn session_policy(&self) -> SessionPolicy { + SessionPolicy { + opening: SessionOpening::InitiatorOnly, reopen: true, } } - fn ordered_session_demand( + fn session_demand( &self, conn_id: ZakuraConnId, peer: &ZakuraPeerId, _negotiated: u64, direction: ServicePeerDirection, - ) -> OrderedSessionDemand { + ) -> SessionDemand { if self .session_states .lock() @@ -326,7 +326,7 @@ impl Service for DiscoveryService { .get(&(peer.clone(), conn_id)) .is_some_and(|record| record.state == DiscoverySessionState::Retired) { - return OrderedSessionDemand::Retire; + return SessionDemand::Retire; } let mut peers = self.handle.subscribe_peer_snapshot(); @@ -336,14 +336,14 @@ impl Service for DiscoveryService { ServicePeerDirection::Outbound => snapshot.outbound_slots_free, }; if slots_free == 0 { - return OrderedSessionDemand::WaitForChange(Box::pin(async move { + return SessionDemand::WaitForChange(Box::pin(async move { if peers.changed().await.is_err() { std::future::pending::<()>().await; } })); } - OrderedSessionDemand::OpenNow + SessionDemand::OpenNow } fn wants_peer( @@ -1970,13 +1970,13 @@ mod tests { ); assert_eq!(record.state, DiscoverySessionState::Retired); assert!(matches!( - service.ordered_session_demand( + service.session_demand( 0, &peer_id, ZAKURA_CAP_DISCOVERY, ServicePeerDirection::Inbound, ), - OrderedSessionDemand::Retire, + SessionDemand::Retire, )); Ok(()) } @@ -2313,13 +2313,13 @@ mod tests { // A finished discovery-only exchange retires the session so the handler // cannot reopen the stream while the connection tears down. assert!(matches!( - service.ordered_session_demand( + service.session_demand( 0, &peer_id, ZAKURA_CAP_DISCOVERY | ZAKURA_CAP_BLOCK_SYNC, ServicePeerDirection::Inbound, ), - OrderedSessionDemand::Retire, + SessionDemand::Retire, )); wait_for_discovery_inbound_peers(&handle, 0).await; @@ -2557,13 +2557,13 @@ mod tests { // The refresh loop still owns this session. // The handler must keep the stream eligible. assert!(!matches!( - service.ordered_session_demand( + service.session_demand( 0, &peer_id, ZAKURA_CAP_DISCOVERY | ZAKURA_CAP_HEADER_SYNC, ServicePeerDirection::Inbound, ), - OrderedSessionDemand::Retire, + SessionDemand::Retire, )); connection_cancel.cancel(); diff --git a/crates/zakura-network/src/zakura/handler.rs b/crates/zakura-network/src/zakura/handler.rs index caba3794e4..ba22759bd0 100644 --- a/crates/zakura-network/src/zakura/handler.rs +++ b/crates/zakura-network/src/zakura/handler.rs @@ -1,6 +1,8 @@ //! Zakura P2P v2 endpoint, protocol handler, and bounded connection serving. +mod service_session; mod trace; +use service_session::{spawn_service_session, PendingSessions, PreparedStream, SetupIo}; use std::{ collections::{HashMap, HashSet}, @@ -16,7 +18,7 @@ use std::{ }; use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt}; -use futures::StreamExt as _; +use futures::{future::BoxFuture, StreamExt as _}; use iroh::{ endpoint::{Connection, Endpoint, QuicTransportConfig, RecvStream, SendStream, VarInt}, protocol::{AcceptError, ProtocolHandler, Router}, @@ -41,7 +43,10 @@ use zakura_chain::{ use self::trace::ZakuraConnTrace; use super::discovery::{self, native_dial_supervised, spawn_native_bootstrap_dialer, RedialPolicy}; use super::trace::{reject_reason_label, ZakuraTrace}; -use super::transport::{worker_framed_channel, FramedWorkerRecv, QueuedFrame}; +use super::transport::{ + worker_framed_channel, FramedWorkerRecv, OrderedStreamFailure, OrderedStreamFailureCause, + SessionLayout, +}; #[cfg(any(test, feature = "zakura-testkit"))] use crate::zakura::drive_header_sync_actions; #[cfg(any(test, feature = "zakura-testkit"))] @@ -54,9 +59,9 @@ use crate::{ AuthenticatedPeerRegistration, BlockSyncAction, BlockSyncFrontiers, BlockSyncHandle, BlockSyncService, BlockSyncStartup, BoxRunFuture, Clock, CloseCause, Frame, FramedRecv, FramedSend, FullStateFrontiers, HeaderSyncPassthroughService, HeaderSyncService, - HeaderSyncStartup, OrderedSessionDemand, OrderedStreamOpening, OrderedStreamPolicy, Peer, - RealClock, Service, ServicePeerDirection, ServiceRegistry, ServiceStream, SinkReject, - Stream, StreamMode, StreamPrelude, ZakuraAcceptedLimits, ZakuraBlockSyncConfig, + HeaderSyncStartup, Peer, RealClock, Service, ServicePeerDirection, ServiceRegistry, + ServiceStream, SessionDemand, SessionOpening, SessionPolicy, SinkReject, Stream, + StreamMode, StreamPrelude, StreamWritePolicy, ZakuraAcceptedLimits, ZakuraBlockSyncConfig, ZakuraConnId, ZakuraControlAck, ZakuraControlHello, ZakuraControlRole, ZakuraControlValidation, ZakuraHandshakeConfig, ZakuraHandshakePath, ZakuraHeaderSyncConfig, ZakuraInitialLimits, ZakuraLimits, ZakuraPeerId, @@ -145,8 +150,9 @@ pub const ZAKURA_DUPLICATE_EVICT_MIN_AGE: Duration = Duration::from_secs(300); /// resolution (milliseconds) so a genuine race keeps the transcript-tiebreak /// winner instead of flapping. pub const ZAKURA_SAME_IP_DUPLICATE_EVICT_MIN_AGE: Duration = Duration::from_secs(5); -/// QUIC stream receive window used by Zakura endpoints. -pub const DEFAULT_ZAKURA_STREAM_RECEIVE_WINDOW: u32 = 32 * 1024 * 1024; +/// A paused stream may consume at most half the connection receive window, +/// leaving credit for another service stream. +pub const DEFAULT_ZAKURA_STREAM_RECEIVE_WINDOW: u32 = 16 * 1024 * 1024; /// QUIC connection receive window used by Zakura endpoints. pub const DEFAULT_ZAKURA_RECEIVE_WINDOW: u32 = 32 * 1024 * 1024; /// QUIC send window used by Zakura endpoints. @@ -472,6 +478,10 @@ impl ZakuraLocalLimits { /// Returns the QUIC transport config matching these local limits. pub fn transport_config(&self) -> QuicTransportConfig { + self.transport_config_builder().build() + } + + fn transport_config_builder(&self) -> iroh::endpoint::QuicTransportConfigBuilder { QuicTransportConfig::builder() .max_remote_nat_traversal_addresses(0) .max_concurrent_bidi_streams(VarInt::from_u32(u32::from(self.max_open_streams))) @@ -487,7 +497,6 @@ impl ZakuraLocalLimits { .keep_alive_interval(self.keep_alive_interval) .datagram_receive_buffer_size(None) .datagram_send_buffer_size(0) - .build() } } @@ -1528,7 +1537,18 @@ impl Drop for RegisteredPeerCleanupGuard { } } +struct IncomingStreamSetup { + io: SetupIo, + permit: OwnedSemaphorePermit, + stream_id: u64, + stream: Stream, + prelude: StreamPrelude, + session: Option<(SessionLayout, u64)>, +} + struct StreamAdmission<'a> { + is_initiator: bool, + direction: ServicePeerDirection, conn: ZakuraConnTrace, peer_id: &'a ZakuraPeerId, stream_sem: &'a Arc, @@ -1579,6 +1599,11 @@ struct StreamWorkerContext { _permit: OwnedSemaphorePermit, limits: ZakuraConnectionLimits, inbound_frame_cap: u32, + message_payload_limits: &'static [(u16, usize)], + message_types: Option<&'static [u16]>, + queue_depths: Option<(usize, usize)>, + write_policy: StreamWritePolicy, + session_resources: Option>, outbound_frame_cap: u32, message_bucket: SharedMessageBucket, connection_token: CancellationToken, @@ -1587,42 +1612,67 @@ struct StreamWorkerContext { freshness_tx: watch::Sender, } -struct AdmittedOrderedSession { +struct AdmittedSession { kind: u16, - version: u16, session_id: u64, + cancel_token: CancellationToken, + streams: Vec, +} + +struct ServiceStreamRole { + kind: u16, + version: u16, recv: FramedRecv, send: FramedSend, - cancel_token: CancellationToken, +} + +impl AdmittedSession { + fn into_service_streams(self) -> HashMap { + self.streams + .into_iter() + .map(|role| { + ( + role.kind, + ServiceStream::new( + self.session_id, + role.version, + role.recv, + role.send, + self.cancel_token.clone(), + ), + ) + }) + .collect() + } } #[derive(Copy, Clone, Debug)] -struct OrderedSessionExit { +struct SessionExit { stream: Stream, session_id: u64, opened_locally: bool, } #[derive(Copy, Clone, Debug, Eq, PartialEq)] -enum OrderedSessionWaitReason { +enum SessionWaitReason { Demand, Transport, } #[derive(Copy, Clone, Debug, Default, Eq, PartialEq)] -enum OrderedSessionReopenState { +enum SessionReopenState { #[default] Idle, - Waiting(OrderedSessionWaitReason), + Waiting(SessionWaitReason), Retired, } -type OrderedSessionWait = futures::stream::Once>; -type OrderedSessionWaits = StreamMap; +type SessionWait = futures::stream::Once>; +type SessionWaits = StreamMap; /// Connection-owned lifecycle for one negotiated ordered stream kind. #[derive(Debug)] -struct OrderedSessionState { +struct ServiceSessionState { stream: Stream, local_session_id: Option, remote_session_id: Option, @@ -1630,17 +1680,17 @@ struct OrderedSessionState { // service can immediately park the stream again, and resetting here would // turn persistent refusal into a base-delay reopen loop. reopen_attempts: u32, - reopen_state: OrderedSessionReopenState, + reopen_state: SessionReopenState, } -impl OrderedSessionState { +impl ServiceSessionState { fn new(stream: Stream) -> Self { Self { stream, local_session_id: None, remote_session_id: None, reopen_attempts: 0, - reopen_state: OrderedSessionReopenState::Idle, + reopen_state: SessionReopenState::Idle, } } @@ -1662,17 +1712,17 @@ impl OrderedSessionState { true } - fn schedule_demand(&mut self, waits: &mut OrderedSessionWaits, demand: OrderedSessionDemand) { + fn schedule_demand(&mut self, waits: &mut SessionWaits, demand: SessionDemand) { match demand { - OrderedSessionDemand::OpenNow => { + SessionDemand::OpenNow => { self.cancel_wait(waits); self.install_wait( waits, - OrderedSessionWaitReason::Demand, + SessionWaitReason::Demand, Box::pin(future::ready(())), ); } - OrderedSessionDemand::RetryAt(at) => { + SessionDemand::RetryAt(at) => { self.install_demand_wait( waits, Box::pin(async move { @@ -1680,18 +1730,18 @@ impl OrderedSessionState { }), ); } - OrderedSessionDemand::WaitForChange(changed) => { + SessionDemand::WaitForChange(changed) => { self.install_demand_wait(waits, changed); } - OrderedSessionDemand::Retire => { + SessionDemand::Retire => { waits.remove(&self.stream.kind); - self.reopen_state = OrderedSessionReopenState::Retired; + self.reopen_state = SessionReopenState::Retired; } } } - fn schedule_transport_backoff(&mut self, waits: &mut OrderedSessionWaits) -> Option { - if self.reopen_state != OrderedSessionReopenState::Idle { + fn schedule_transport_backoff(&mut self, waits: &mut SessionWaits) -> Option { + if self.reopen_state != SessionReopenState::Idle { return None; } debug_assert!( @@ -1703,48 +1753,43 @@ impl OrderedSessionState { self.reopen_attempts = self.reopen_attempts.saturating_add(1); self.install_wait( waits, - OrderedSessionWaitReason::Transport, + SessionWaitReason::Transport, Box::pin(tokio::time::sleep(delay)), ); Some(delay) } - fn finish_wait(&mut self, waits: &mut OrderedSessionWaits) { + fn finish_wait(&mut self, waits: &mut SessionWaits) { let removed = waits.remove(&self.stream.kind); debug_assert!( removed.is_some(), "yielded ordered session wait remains keyed until it is handled" ); - if matches!(self.reopen_state, OrderedSessionReopenState::Waiting(_)) { - self.reopen_state = OrderedSessionReopenState::Idle; + if matches!(self.reopen_state, SessionReopenState::Waiting(_)) { + self.reopen_state = SessionReopenState::Idle; } } - fn cancel_wait(&mut self, waits: &mut OrderedSessionWaits) { + fn cancel_wait(&mut self, waits: &mut SessionWaits) { waits.remove(&self.stream.kind); - if matches!(self.reopen_state, OrderedSessionReopenState::Waiting(_)) { - self.reopen_state = OrderedSessionReopenState::Idle; + if matches!(self.reopen_state, SessionReopenState::Waiting(_)) { + self.reopen_state = SessionReopenState::Idle; } } - fn install_demand_wait( - &mut self, - waits: &mut OrderedSessionWaits, - wait: BoxRunFuture<'static, ()>, - ) { - if self.reopen_state == OrderedSessionReopenState::Waiting(OrderedSessionWaitReason::Demand) - { + fn install_demand_wait(&mut self, waits: &mut SessionWaits, wait: BoxRunFuture<'static, ()>) { + if self.reopen_state == SessionReopenState::Waiting(SessionWaitReason::Demand) { return; } self.cancel_wait(waits); - self.install_wait(waits, OrderedSessionWaitReason::Demand, wait); + self.install_wait(waits, SessionWaitReason::Demand, wait); } fn install_wait( &mut self, - waits: &mut OrderedSessionWaits, - reason: OrderedSessionWaitReason, + waits: &mut SessionWaits, + reason: SessionWaitReason, wait: BoxRunFuture<'static, ()>, ) { let replaced = waits.insert(self.stream.kind, futures::stream::once(wait)); @@ -1752,7 +1797,7 @@ impl OrderedSessionState { replaced.is_none(), "ordered session wait is cancelled before replacement" ); - self.reopen_state = OrderedSessionReopenState::Waiting(reason); + self.reopen_state = SessionReopenState::Waiting(reason); } } @@ -1761,8 +1806,8 @@ impl OrderedSessionState { /// The initiator opens ordinary ordered streams. Block sync is symmetric, so /// either side may open it and simultaneous offers use the connection's /// deterministic collision tiebreak. -fn may_open_ordered_stream(policy: OrderedStreamPolicy, is_initiator: bool) -> bool { - is_initiator || policy.opening == OrderedStreamOpening::EitherSide +fn may_open_ordered_stream(policy: SessionPolicy, is_initiator: bool) -> bool { + is_initiator || policy.opening == SessionOpening::EitherSide } /// Whether this endpoint proactively opens an ordered stream. @@ -1772,13 +1817,13 @@ fn may_open_ordered_stream(policy: OrderedStreamPolicy, is_initiator: bool) -> b /// accept either-side streams, preserving compatibility with peers that race /// an offer during upgrade. fn opens_ordered_stream_locally( - policy: OrderedStreamPolicy, + policy: SessionPolicy, is_initiator: bool, i_open_collision_winner: bool, ) -> bool { match policy.opening { - OrderedStreamOpening::InitiatorOnly => is_initiator, - OrderedStreamOpening::EitherSide => i_open_collision_winner, + SessionOpening::InitiatorOnly => is_initiator, + SessionOpening::EitherSide => i_open_collision_winner, } } @@ -1786,13 +1831,13 @@ fn opens_ordered_stream_locally( /// sibling services remain healthy. The endpoint designated by the transport /// policy keeps offering a replacement with bounded backoff. fn should_reopen_ordered_session( - exited: OrderedSessionExit, - policy: OrderedStreamPolicy, + exited: SessionExit, + policy: SessionPolicy, is_initiator: bool, i_open_collision_winner: bool, connection_cancelled: bool, ) -> bool { - exited.stream.mode == StreamMode::Ordered + exited.stream.mode == StreamMode::Persistent && policy.reopen && opens_ordered_stream_locally(policy, is_initiator, i_open_collision_winner) && !connection_cancelled @@ -2280,18 +2325,25 @@ impl ZakuraProtocolHandler { let accepted_capabilities = context.accepted_capabilities; let stream_sem = Arc::new(Semaphore::new(usize::from(limits.max_open_streams))); let mut workers = JoinSet::new(); - let (ordered_session_exit_tx, mut ordered_session_exit_rx) = mpsc::unbounded_channel(); - let mut ordered_session_waits = OrderedSessionWaits::new(); + let (session_exit_tx, mut session_exit_rx) = mpsc::unbounded_channel(); + let mut session_waits = SessionWaits::new(); + let mut pending_sessions = PendingSessions::default(); + let mut incoming_setup: Option>> = None; let mut open_limiter = TokenBucket::new(limits.stream_open_rate_per_second); let mut message_buckets = MessageRateBuckets::new(); let (freshness_tx, freshness_rx) = watch::channel(Instant::now()); let negotiated_ordered_streams = self .registry - .ordered_streams_for_negotiated(accepted_capabilities); - let mut ordered_sessions: HashMap = negotiated_ordered_streams + .persistent_streams_for_negotiated(accepted_capabilities); + let mut service_sessions: HashMap = negotiated_ordered_streams .iter() .copied() - .map(|stream| (stream.kind, OrderedSessionState::new(stream))) + .filter(|stream| { + self.registry + .session_layout(*stream) + .is_some_and(|layout| layout.primary() == *stream) + }) + .map(|stream| (stream.kind, ServiceSessionState::new(stream))) .collect(); // The dialer opens ordinary ordered streams. For symmetric block sync, // the mirror-stable node-id tiebreak designates one proactive opener, @@ -2301,7 +2353,14 @@ impl ZakuraProtocolHandler { let mut ordered_streams = Vec::new(); let mut deferred_ordered_streams = Vec::new(); for stream in negotiated_ordered_streams.iter().copied() { - let policy = self.registry.ordered_stream_policy(stream.kind); + if self + .registry + .session_layout(stream) + .is_some_and(|layout| stream != layout.primary()) + { + continue; + } + let policy = self.registry.session_policy(stream.kind); if !opens_ordered_stream_locally( policy, context.is_initiator, @@ -2310,14 +2369,14 @@ impl ZakuraProtocolHandler { continue; } - match self.registry.ordered_session_demand( + match self.registry.session_demand( stream.kind, conn_id, accepted_capabilities, &peer_id, context.direction, ) { - OrderedSessionDemand::OpenNow => ordered_streams.push(stream), + SessionDemand::OpenNow => ordered_streams.push(stream), demand => deferred_ordered_streams.push((stream, demand)), } } @@ -2325,10 +2384,19 @@ impl ZakuraProtocolHandler { .registry .request_response_streams_for_negotiated(accepted_capabilities) .len(); - if ordered_streams.len() > usize::from(limits.max_open_streams) { + // Count every required stream in the sessions that local demand will open. + let opening_stream_count: usize = ordered_streams + .iter() + .map(|stream| { + self.registry + .session_layout(*stream) + .map_or(1, |layout| layout.streams.len()) + }) + .sum(); + if opening_stream_count > usize::from(limits.max_open_streams) { debug!( max_open_streams = limits.max_open_streams, - ordered_stream_count = ordered_streams.len(), + ordered_stream_count = opening_stream_count, "closing Zakura peer because negotiated ordered streams exceed max-open-streams" ); connection.close(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE), b"ordered streams"); @@ -2371,9 +2439,8 @@ impl ZakuraProtocolHandler { } else if !connection_token.is_cancelled() { let mut opened_capabilities = 0; for stream in ordered_streams { - opened_capabilities |= stream.capability; let admitted = match self - .open_ordered_service_stream( + .open_service_session( &connection, stream, &mut workers, @@ -2386,11 +2453,21 @@ impl ZakuraProtocolHandler { freshness_tx.clone(), conn.clone(), peer_id.clone(), - ordered_session_exit_tx.clone(), + context.direction, + session_exit_tx.clone(), ) .await { Ok(admitted) => admitted, + Err(ZakuraHandlerError::SessionFull) => { + // Demand is advisory: another connection can reserve the + // last slot before this open. Retry only this service. + service_sessions + .get_mut(&stream.kind) + .expect("selected stream has negotiated session state") + .schedule_transport_backoff(&mut session_waits); + continue; + } Err(error) => { debug!( ?error, @@ -2406,20 +2483,12 @@ impl ZakuraProtocolHandler { break; } }; - ordered_sessions + opened_capabilities |= stream.capability; + service_sessions .get_mut(&admitted.kind) .expect("opened ordered stream was selected from negotiated session state") .local_session_id = Some(admitted.session_id); - service_streams.insert( - admitted.kind, - ServiceStream::new( - admitted.session_id, - admitted.version, - admitted.recv, - admitted.send, - admitted.cancel_token, - ), - ); + service_streams.extend(admitted.into_service_streams()); } if !connection_token.is_cancelled() { // Escalation is already narrowed to opened ordered services. @@ -2449,32 +2518,39 @@ impl ZakuraProtocolHandler { ?demand, "deferring ordered service session according to reactor demand" ); - ordered_sessions + service_sessions .get_mut(&stream.kind) .expect("deferred ordered stream has negotiated session state") - .schedule_demand(&mut ordered_session_waits, demand); + .schedule_demand(&mut session_waits, demand); } } loop { + let session_deadline = pending_sessions.deadline(); tokio::select! { biased; _ = connection_token.cancelled() => break, + _ = async { + match session_deadline { + Some(deadline) => tokio::time::sleep_until(deadline).await, + None => future::pending().await, + } + } => pending_sessions.expire(Instant::now()), _ = freshness_reaper(freshness_rx.clone(), limits.idle_timeout), if run_freshness_reaper => { connection.close(VarInt::from_u32(ZAKURA_CLOSE_NEUTRAL), b"idle"); close_cause.record("idle_timeout"); break; } - Some(exited) = ordered_session_exit_rx.recv() => { + Some(exited) = session_exit_rx.recv() => { // Stop tracking the dead generation, or a legitimate reopen of the // same kind would look like a duplicate stream and kill the whole // connection, taking block sync and gossip down with header sync. - let Some(session) = ordered_sessions.get_mut(&exited.stream.kind) else { + let Some(session) = service_sessions.get_mut(&exited.stream.kind) else { continue; }; let removed_active_session = session.remove_active_session(exited.opened_locally, exited.session_id); - let policy = self.registry.ordered_stream_policy(exited.stream.kind); + let policy = self.registry.session_policy(exited.stream.kind); if removed_active_session && !session.has_active_session() && should_reopen_ordered_session( @@ -2486,7 +2562,7 @@ impl ZakuraProtocolHandler { ) { if let Some(delay) = - session.schedule_transport_backoff(&mut ordered_session_waits) + session.schedule_transport_backoff(&mut session_waits) { debug!( stream_kind = exited.stream.kind, @@ -2498,13 +2574,13 @@ impl ZakuraProtocolHandler { } } } - Some((kind, ())) = ordered_session_waits.next(), if !ordered_session_waits.is_empty() => { - let Some(session) = ordered_sessions.get_mut(&kind) else { + Some((kind, ())) = session_waits.next(), if !session_waits.is_empty() => { + let Some(session) = service_sessions.get_mut(&kind) else { continue; }; - session.finish_wait(&mut ordered_session_waits); + session.finish_wait(&mut session_waits); let stream = session.stream; - let policy = self.registry.ordered_stream_policy(stream.kind); + let policy = self.registry.session_policy(stream.kind); if connection_token.is_cancelled() || !opens_ordered_stream_locally( policy, @@ -2516,7 +2592,7 @@ impl ZakuraProtocolHandler { continue; } - let demand = self.registry.ordered_session_demand( + let demand = self.registry.session_demand( stream.kind, conn_id, accepted_capabilities, @@ -2524,15 +2600,15 @@ impl ZakuraProtocolHandler { context.direction, ); match demand { - OrderedSessionDemand::OpenNow => {} + SessionDemand::OpenNow => {} demand => { - session.schedule_demand(&mut ordered_session_waits, demand); + session.schedule_demand(&mut session_waits, demand); continue; } } match self - .open_ordered_service_stream( + .open_service_session( &connection, stream, &mut workers, @@ -2545,26 +2621,18 @@ impl ZakuraProtocolHandler { freshness_tx.clone(), conn.clone(), peer_id.clone(), - ordered_session_exit_tx.clone(), + context.direction, + session_exit_tx.clone(), ) .await { Ok(admitted) => { - let session = ordered_sessions + let session = service_sessions .get_mut(&admitted.kind) .expect("opened ordered stream has negotiated session state"); - session.cancel_wait(&mut ordered_session_waits); + session.cancel_wait(&mut session_waits); session.local_session_id = Some(admitted.session_id); - let service_streams = HashMap::from([( - admitted.kind, - ServiceStream::new( - admitted.session_id, - admitted.version, - admitted.recv, - admitted.send, - admitted.cancel_token, - ), - )]); + let service_streams = admitted.into_service_streams(); let admitted_capabilities = self.registry.add_escalated_peer( Peer::new_with_service_streams( conn_id, @@ -2580,10 +2648,10 @@ impl ZakuraProtocolHandler { cleanup_guard.add_admitted_capabilities(admitted_capabilities); } Err(error) => { - let delay = ordered_sessions + let delay = service_sessions .get_mut(&stream.kind) .expect("failed ordered stream open has negotiated session state") - .schedule_transport_backoff(&mut ordered_session_waits); + .schedule_transport_backoff(&mut session_waits); debug!( ?error, stream_kind = stream.kind, @@ -2598,10 +2666,156 @@ impl ZakuraProtocolHandler { debug!(?error, "Zakura stream worker exited unexpectedly"); } } - accepted = connection.accept_bi() => { + setup = async { + incoming_setup.as_mut().expect("setup branch is enabled only while reading").await + }, if incoming_setup.is_some() => { + incoming_setup = None; + let Some(setup) = setup else { continue; }; + // Active and retiring remote sessions still own their slot. + // Reject duplicates before reserving capacity for any member. + if setup.session.as_ref().is_some_and(|(layout, _)| { + service_sessions + .get(&layout.primary().kind) + .is_some_and(|session| session.remote_session_id.is_some()) + }) { + continue; + } + let mut admission = StreamAdmission { + is_initiator: context.is_initiator, + direction: context.direction, + conn: conn.clone(), + peer_id: &peer_id, + stream_sem: &stream_sem, + open_limiter: &mut open_limiter, + message_buckets: &mut message_buckets, + workers: &mut workers, + limits, + accepted_capabilities, + connection_token: connection_token.clone(), + close_cause: close_cause.clone(), + freshness_tx: freshness_tx.clone(), + }; + if let Some(admitted) = self.finish_bi_stream_setup( + setup, &mut admission, per_stream_queue_depth, + session_exit_tx.clone(), &mut pending_sessions, + ) { + let kind = admitted.kind; + + // A stream kind we never negotiated is a protocol + // fault: tear the connection down (unchanged + // strictness). + if !service_sessions.contains_key(&kind) { + debug!( + stream_kind = kind, + "closing peer after unexpected ordered stream" + ); + close_cause.record("unexpected_stream"); + connection_token.cancel(); + continue; + } + let is_collision = + service_sessions[&kind].local_session_id.is_some(); + // The deterministic winner keeps its own stream + // and parks the peer's. The loser falls through + // and adopts the peer's stream. + if is_collision && context.i_open_collision_winner { + debug!( + stream_kind = kind, + "winning ordered-stream collision; parking peer's duplicate" + ); + admitted.cancel_token.cancel(); + continue; + } + + service_sessions + .get_mut(&kind) + .expect("accepted ordered stream has negotiated session state") + .remote_session_id = Some(admitted.session_id); + + // Honour the owning service's current per-peer + // demand. For a collision we lost, demand is + // implied because we already opened the same kind. + let demand = (!is_collision).then(|| { + self.registry.reserved_session_demand( + kind, conn_id, accepted_capabilities, &peer_id, context.direction, + ) + }); + if demand + .as_ref() + .is_some_and(|demand| { + !matches!(demand, SessionDemand::OpenNow) + }) + { + metrics::counter!( + "zakura.p2p.stream.parked.no_demand", + "stream_kind" => stream_kind_label(kind), + ) + .increment(1); + info!( + ?peer_id, + stream_kind = kind, + "locally parking ordered service stream because the service has no demand" + ); + // We are not adopting it after all, so this kind + // is free again -- both for a later re-offer by + // the peer and for the demand re-check below. + let session = service_sessions + .get_mut(&kind) + .expect("parked ordered stream has negotiated session state"); + session.remove_active_session(false, admitted.session_id); + // If this side may open the stream, re-check its + // local demand. Otherwise the entitled remote + // opener observes the parked stream and retries. + let policy = self.registry.session_policy(kind); + if policy.reopen + && opens_ordered_stream_locally( + policy, + context.is_initiator, + context.i_open_collision_winner, + ) + { + session.schedule_demand( + &mut session_waits, + demand.expect("non-open demand exists because this branch checked it"), + ); + } + admitted.cancel_token.cancel(); + continue; + } + + service_sessions + .get_mut(&kind) + .expect("adopted ordered stream has negotiated session state") + .cancel_wait(&mut session_waits); + let service_streams = admitted.into_service_streams(); + // We keep the full accepted capability context so + // discovery can make cross-service ownership + // decisions; disconnect fanout still uses the + // registry's returned admitted mask. When this is + // a lost collision, `add_escalated_peer` → + // `add_peer` replaces our own opened session for + // this peer (see `can_admit_peer`). + let admitted_capabilities = self.registry.add_escalated_peer( + Peer::new_with_service_streams( + conn_id, + peer_id.clone(), + remote_ip, + accepted_capabilities, + context.direction, + service_streams, + connection_token.clone(), + close_cause.clone(), + ), + ); + cleanup_guard.add_admitted_capabilities(admitted_capabilities); + } + } + accepted = connection.accept_bi(), if incoming_setup.is_none() => { match accepted { Ok((send, recv)) => { let mut admission = StreamAdmission { + is_initiator: context.is_initiator, + direction: context.direction, conn: conn.clone(), peer_id: &peer_id, stream_sem: &stream_sem, @@ -2614,169 +2828,7 @@ impl ZakuraProtocolHandler { close_cause: close_cause.clone(), freshness_tx: freshness_tx.clone(), }; - if let Some(admitted) = self - .admit_bi_stream( - send, - recv, - &mut admission, - per_stream_queue_depth, - ordered_session_exit_tx.clone(), - ) - .await - { - let kind = admitted.kind; - - // A stream kind we never negotiated is a protocol - // fault: tear the connection down (unchanged - // strictness). - if !ordered_sessions.contains_key(&kind) { - debug!( - stream_kind = kind, - "closing peer after unexpected ordered stream" - ); - close_cause.record("unexpected_stream"); - connection_token.cancel(); - continue; - } - // Reject a stream opened by a peer whose - // connection role is not entitled to open this - // kind, even if it races our own offer. - let policy = self.registry.ordered_stream_policy(kind); - if !may_open_ordered_stream(policy, !context.is_initiator) { - debug!( - stream_kind = kind, - "closing peer after unexpected ordered stream" - ); - close_cause.record("unexpected_stream"); - connection_token.cancel(); - continue; - } - - let is_collision = - ordered_sessions[&kind].local_session_id.is_some(); - // The deterministic winner keeps its own stream - // and parks the peer's. The loser falls through - // and adopts the peer's stream. - if is_collision && context.i_open_collision_winner { - debug!( - stream_kind = kind, - "winning ordered-stream collision; parking peer's duplicate" - ); - admitted.cancel_token.cancel(); - continue; - } - - // We intend to adopt the peer's stream. A second - // accepted stream of the same kind is a real - // duplicate and a fault. - if ordered_sessions[&kind].remote_session_id.is_some() { - debug!( - stream_kind = kind, - "closing peer after duplicate ordered stream" - ); - close_cause.record("duplicate_stream"); - connection_token.cancel(); - continue; - } - ordered_sessions - .get_mut(&kind) - .expect("accepted ordered stream has negotiated session state") - .remote_session_id = Some(admitted.session_id); - - // Honour the owning service's current per-peer - // demand. For a collision we lost, demand is - // implied because we already opened the same kind. - let demand = (!is_collision).then(|| { - self.registry.ordered_session_demand( - kind, - conn_id, - accepted_capabilities, - &peer_id, - context.direction, - ) - }); - if demand - .as_ref() - .is_some_and(|demand| { - !matches!(demand, OrderedSessionDemand::OpenNow) - }) - { - metrics::counter!( - "zakura.p2p.stream.parked.no_demand", - "stream_kind" => stream_kind_label(kind), - ) - .increment(1); - info!( - ?peer_id, - stream_kind = kind, - "locally parking ordered service stream because the service has no demand" - ); - // We are not adopting it after all, so this kind - // is free again -- both for a later re-offer by - // the peer and for the demand re-check below. - let session = ordered_sessions - .get_mut(&kind) - .expect("parked ordered stream has negotiated session state"); - session.remove_active_session(false, admitted.session_id); - // If this side may open the stream, re-check its - // local demand. Otherwise the entitled remote - // opener observes the parked stream and retries. - if policy.reopen - && opens_ordered_stream_locally( - policy, - context.is_initiator, - context.i_open_collision_winner, - ) - { - session.schedule_demand( - &mut ordered_session_waits, - demand.expect("non-open demand exists because this branch checked it"), - ); - } - admitted.cancel_token.cancel(); - continue; - } - - ordered_sessions - .get_mut(&kind) - .expect("adopted ordered stream has negotiated session state") - .cancel_wait(&mut ordered_session_waits); - let service_streams = HashMap::from([( - kind, - ServiceStream::new( - admitted.session_id, - admitted.version, - admitted.recv, - admitted.send, - admitted.cancel_token.clone(), - ), - )]); - // Current ordered services own one ordered stream - // each, so we fan out accepted streams one at a - // time. Batch here if a service gains multiple - // ordered streams. - // - // We keep the full accepted capability context so - // discovery can make cross-service ownership - // decisions; disconnect fanout still uses the - // registry's returned admitted mask. When this is - // a lost collision, `add_escalated_peer` → - // `add_peer` replaces our own opened session for - // this peer (see `can_admit_peer`). - let admitted_capabilities = self.registry.add_escalated_peer( - Peer::new_with_service_streams( - conn_id, - peer_id.clone(), - remote_ip, - accepted_capabilities, - context.direction, - service_streams, - connection_token.clone(), - close_cause.clone(), - ), - ); - cleanup_guard.add_admitted_capabilities(admitted_capabilities); - } + incoming_setup = self.begin_bi_stream_setup(send, recv, &mut admission); } Err(error) => { debug!(?error, "Zakura connection stopped accepting streams"); @@ -2817,6 +2869,8 @@ impl ZakuraProtocolHandler { &connection, limits, stream, + self.registry.message_payload_limits(stream), + self.registry.message_types(stream), request_id, message_type, flags, @@ -2847,6 +2901,8 @@ impl ZakuraProtocolHandler { } connection_token.cancel(); + drop(incoming_setup); + drop(pending_sessions); while let Some(joined) = timeout(STREAM_WORKER_DRAIN_TIMEOUT, workers.join_next()) .await .ok() @@ -2869,7 +2925,7 @@ impl ZakuraProtocolHandler { } #[allow(clippy::too_many_arguments)] - async fn open_ordered_service_stream( + async fn open_service_session( &self, connection: &Connection, stream: Stream, @@ -2883,81 +2939,60 @@ impl ZakuraProtocolHandler { freshness_tx: watch::Sender, conn: ZakuraConnTrace, peer_id: ZakuraPeerId, - ordered_session_exit_tx: mpsc::UnboundedSender, - ) -> Result { - let stream_id = self.next_stream_id.fetch_add(1, Ordering::Relaxed); - let permit = stream_sem - .clone() - .try_acquire_owned() - .map_err(|_| ZakuraHandlerError::ResourceLimit("ordered stream permit"))?; - let (mut send, recv) = timeout(OUTBOUND_STREAM_WRITE_TIMEOUT, connection.open_bi()) - .await - .map_err(|_| ZakuraHandlerError::Timeout("open ordered service stream"))??; - let prelude = StreamPrelude { - magic: STREAM_PRELUDE_MAGIC, - stream_kind: stream.kind, - stream_version: stream.version, - request_id: None, - max_frame_bytes: inbound_frame_cap_for_stream(&limits, stream), - }; - let prelude_bytes = prelude.encode()?; - timeout( - OUTBOUND_STREAM_WRITE_TIMEOUT, - send.write_all(&prelude_bytes), - ) - .await - .map_err(|_| ZakuraHandlerError::Timeout("ordered stream prelude write"))??; - - let message_bucket = message_bucket_for( - message_buckets, - stream.kind, - limits.message_rate_per_second, - RealClock, - ); - let stream_token = connection_token.child_token(); - let context = StreamWorkerContext { - conn: conn.clone(), - peer_id, - stream_id, - _permit: permit, - limits, - inbound_frame_cap: prelude.max_frame_bytes, - outbound_frame_cap: application_frame_cap(&limits, stream), - message_bucket, - connection_token, - stream_token, - close_cause, - freshness_tx, - }; - - metrics::counter!( - "zakura.p2p.stream.accepted", - "stream_kind" => stream_kind_label(stream.kind), - ) - .increment(1); - conn.trace_stream("accepted", stream_id, Some(stream_kind_label(stream.kind))); - - Ok(spawn_persistent_stream_worker( + direction: ServicePeerDirection, + session_exit_tx: mpsc::UnboundedSender, + ) -> Result { + let layout = self + .registry + .session_layout(stream) + .expect("only negotiated persistent streams open service sessions"); + if layout.primary() != stream { + return Err(ZakuraHandlerError::InvalidServiceSession); + } + let resources = self + .registry + .service_for_kind(stream.kind) + .expect("a selected stream has an owning service") + .reserve_session(direction) + .map_err(|_| ZakuraHandlerError::SessionFull)?; + let wire_id = layout.is_multi_stream().then(random_stream_session_seed); + let mut prepared = Vec::with_capacity(layout.streams.len()); + for member in layout.streams.iter().copied() { + let mut member = self + .prepare_ordered_stream( + connection, + member, + wire_id, + stream_sem, + message_buckets, + limits, + connection_token.clone(), + close_cause.clone(), + freshness_tx.clone(), + conn.clone(), + peer_id.clone(), + ) + .await?; + member.set_session_resources(resources.clone()); + prepared.push(member); + } + Ok(spawn_service_session( workers, - send, - recv, - stream, - prelude, - context, + prepared, per_stream_queue_depth, true, - ordered_session_exit_tx, + session_exit_tx, )) } - async fn admit_bi_stream( + /// Start at most one peer-paced setup read per connection. The returned + /// future owns its stream permit but no service reservation. + fn begin_bi_stream_setup( &self, mut send: SendStream, - mut recv: RecvStream, + recv: RecvStream, admission: &mut StreamAdmission<'_>, - per_stream_queue_depth: usize, - ordered_session_exit_tx: mpsc::UnboundedSender, - ) -> Option { + ) -> Option>> { let stream_id = self.next_stream_id.fetch_add(1, Ordering::Relaxed); let Ok(permit) = admission.stream_sem.clone().try_acquire_owned() else { let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE)); @@ -2984,129 +3019,194 @@ impl ZakuraProtocolHandler { return None; } - let prelude = match read_stream_prelude(&mut recv, admission.limits.prelude_timeout).await { - Ok(prelude) => prelude, - Err(error) => { - debug!(?error, "rejecting Zakura stream with bad prelude"); + let mut io = SetupIo::new(send, recv); + let registry = self.registry.clone(); + let conn = admission.conn.clone(); + let limits = admission.limits; + let accepted_capabilities = admission.accepted_capabilities; + let is_initiator = admission.is_initiator; + let close_cause = admission.close_cause.clone(); + let connection_token = admission.connection_token.clone(); + Some(Box::pin(async move { + let (send, recv) = io.streams(); + let prelude = match read_stream_prelude(recv, limits.prelude_timeout).await { + Ok(prelude) => prelude, + Err(error) => { + debug!(?error, "rejecting Zakura stream with bad prelude"); + let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE)); + metrics::counter!("zakura.p2p.stream.rejected.prelude").increment(1); + conn.trace_stream("rejected.prelude", stream_id, None); + return None; + } + }; + let stream_kind = stream_kind_label(prelude.stream_kind); + + let Some(stream) = registry.stream(prelude.stream_kind, prelude.stream_version) else { + debug!( + stream_kind = prelude.stream_kind, + stream_version = prelude.stream_version, + "rejecting Zakura stream with unknown kind or unsupported version" + ); + let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_UNKNOWN_STREAM)); + metrics::counter!( + "zakura.p2p.stream.rejected.unknown_kind", + "stream_kind" => stream_kind, + ) + .increment(1); + conn.trace_stream("rejected.unknown_kind", stream_id, Some(stream_kind)); + return None; + }; + + if accepted_capabilities & stream.capability != stream.capability { + debug!( + stream_kind = prelude.stream_kind, + stream_version = prelude.stream_version, + accepted_capabilities = accepted_capabilities, + "rejecting Zakura stream that was not negotiated for this peer" + ); + let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_UNKNOWN_STREAM)); + metrics::counter!( + "zakura.p2p.stream.rejected.unnegotiated_capability", + "stream_kind" => stream_kind, + ) + .increment(1); + conn.trace_stream( + "rejected.unnegotiated_capability", + stream_id, + Some(stream_kind), + ); + return None; + } + + if stream.mode == StreamMode::Persistent + && registry + .persistent_streams_for_negotiated(accepted_capabilities) + .iter() + .all(|selected| *selected != stream) + { + debug!( + stream_kind = prelude.stream_kind, + stream_version = prelude.stream_version, + "rejecting a lower ordered-stream version than the mutually selected version" + ); + let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_UNKNOWN_STREAM)); + metrics::counter!( + "zakura.p2p.stream.rejected.unselected_version", + "stream_kind" => stream_kind, + ) + .increment(1); + conn.trace_stream("rejected.unselected_version", stream_id, Some(stream_kind)); + return None; + } + + if stream.mode != StreamMode::RequestResponse && prelude.request_id.is_some() { + debug!("rejecting non-request Zakura stream with request id"); let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE)); - metrics::counter!("zakura.p2p.stream.rejected.prelude").increment(1); - admission - .conn - .trace_stream("rejected.prelude", stream_id, None); + close_cause.record("unexpected_request_id"); + connection_token.cancel(); + metrics::counter!("zakura.p2p.stream.rejected.unexpected_request_id").increment(1); + conn.trace_stream( + "rejected.unexpected_request_id", + stream_id, + Some(stream_kind), + ); return None; } - }; - let stream_kind = stream_kind_label(prelude.stream_kind); - let Some(stream) = self - .registry - .stream(prelude.stream_kind, prelude.stream_version) - else { - debug!( - stream_kind = prelude.stream_kind, - stream_version = prelude.stream_version, - "rejecting Zakura stream with unknown kind or unsupported version" - ); - let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_UNKNOWN_STREAM)); - metrics::counter!( - "zakura.p2p.stream.rejected.unknown_kind", - "stream_kind" => stream_kind, - ) - .increment(1); - admission - .conn - .trace_stream("rejected.unknown_kind", stream_id, Some(stream_kind)); - return None; - }; + if stream.mode == StreamMode::RequestResponse && prelude.request_id.is_none() { + debug!("rejecting Zakura request stream without request id"); + let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE)); + close_cause.record("request_without_id"); + connection_token.cancel(); + metrics::counter!("zakura.p2p.stream.rejected.request_without_id").increment(1); + conn.trace_stream("rejected.request_without_id", stream_id, Some(stream_kind)); + return None; + } - if admission.accepted_capabilities & stream.capability != stream.capability { - debug!( - stream_kind = prelude.stream_kind, - stream_version = prelude.stream_version, - accepted_capabilities = admission.accepted_capabilities, - "rejecting Zakura stream that was not negotiated for this peer" - ); - let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_UNKNOWN_STREAM)); - metrics::counter!( - "zakura.p2p.stream.rejected.unnegotiated_capability", - "stream_kind" => stream_kind, - ) - .increment(1); - admission.conn.trace_stream( - "rejected.unnegotiated_capability", - stream_id, - Some(stream_kind), - ); - return None; - } + if stream.mode == StreamMode::Persistent + && !may_open_ordered_stream(registry.session_policy(stream.kind), !is_initiator) + { + debug!( + stream_kind = stream.kind, + "rejecting an ineligible ordered-stream opener" + ); + let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE)); + let _ = recv.stop(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE)); + close_cause.record("unexpected_stream"); + connection_token.cancel(); + return None; + } - if stream.mode == StreamMode::Ordered - && self - .registry - .ordered_streams_for_negotiated(admission.accepted_capabilities) - .into_iter() - .find(|selected| selected.kind == stream.kind) - .is_some_and(|selected| selected.version != stream.version) - { - debug!( - stream_kind = prelude.stream_kind, - stream_version = prelude.stream_version, - "rejecting a lower ordered-stream version than the mutually selected version" - ); - let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_UNKNOWN_STREAM)); metrics::counter!( - "zakura.p2p.stream.rejected.unselected_version", + "zakura.p2p.stream.accepted", "stream_kind" => stream_kind, ) .increment(1); - admission.conn.trace_stream( - "rejected.unselected_version", - stream_id, - Some(stream_kind), - ); - return None; - } - - if stream.mode != StreamMode::RequestResponse && prelude.request_id.is_some() { - debug!("rejecting non-request Zakura stream with request id"); - let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE)); - admission.close_cause.record("unexpected_request_id"); - admission.connection_token.cancel(); - metrics::counter!("zakura.p2p.stream.rejected.unexpected_request_id").increment(1); - admission.conn.trace_stream( - "rejected.unexpected_request_id", - stream_id, - Some(stream_kind), - ); - return None; - } - - if stream.mode == StreamMode::RequestResponse && prelude.request_id.is_none() { - debug!("rejecting Zakura request stream without request id"); - let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE)); - admission.close_cause.record("request_without_id"); - admission.connection_token.cancel(); - metrics::counter!("zakura.p2p.stream.rejected.request_without_id").increment(1); - admission.conn.trace_stream( - "rejected.request_without_id", + conn.trace_stream("accepted", stream_id, Some(stream_kind)); + + let session = if let Some(layout) = registry.session_layout(stream) { + let mut bytes = [0u8; 8]; + if layout.is_multi_stream() + && !matches!( + timeout(limits.prelude_timeout, recv.read_exact(&mut bytes)).await, + Ok(Ok(())) + ) + { + let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE)); + let _ = recv.stop(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE)); + return None; + } + Some((layout, u64::from_le_bytes(bytes))) + } else { + None + }; + Some(IncomingStreamSetup { + io, + permit, stream_id, - Some(stream_kind), - ); - return None; - } - - metrics::counter!( - "zakura.p2p.stream.accepted", - "stream_kind" => stream_kind, - ) - .increment(1); - admission - .conn - .trace_stream("accepted", stream_id, Some(stream_kind)); + stream, + prelude, + session, + }) + })) + } + /// Admit a decoded prelude without waiting on the peer. Lifecycle events + /// stay reactor-owned while the setup future reads bytes independently. + fn finish_bi_stream_setup( + &self, + incoming: IncomingStreamSetup, + admission: &mut StreamAdmission<'_>, + per_stream_queue_depth: usize, + session_exit_tx: mpsc::UnboundedSender, + pending_sessions: &mut PendingSessions, + ) -> Option { + let IncomingStreamSetup { + io, + permit, + stream_id, + stream, + prelude, + session, + } = incoming; + let (mut send, mut recv) = io.take(); + let resources = if let Some((layout, _)) = &session { + match pending_sessions.reserve_or_share(layout, &self.registry, admission.direction) { + Ok(resources) => resources, + Err(_) => { + let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE)); + let _ = recv.stop(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE)); + return None; + } + } + } else { + None + }; let message_bucket = message_bucket_for( admission.message_buckets, - prelude.stream_kind, + session + .as_ref() + .map_or(prelude.stream_kind, |(layout, _)| layout.primary().kind), admission.limits.message_rate_per_second, RealClock, ); @@ -3119,6 +3219,11 @@ impl ZakuraProtocolHandler { _permit: permit, limits: admission.limits, inbound_frame_cap: inbound_frame_cap_for_stream(&admission.limits, stream), + message_payload_limits: self.registry.message_payload_limits(stream), + message_types: self.registry.message_types(stream), + queue_depths: self.registry.stream_queue_depths(stream), + write_policy: self.registry.stream_write_policy(stream), + session_resources: resources, outbound_frame_cap: peer_accepted_frame_cap( &admission.limits, stream, @@ -3131,28 +3236,33 @@ impl ZakuraProtocolHandler { freshness_tx: admission.freshness_tx.clone(), }; - if stream.mode == StreamMode::RequestResponse { - admission.workers.spawn(request_stream_worker( - send, - recv, - prelude, - context, - self.registry.clone(), - )); - None - } else { - Some(spawn_persistent_stream_worker( - admission.workers, - send, - recv, - stream, - prelude, - context, - per_stream_queue_depth, - false, - ordered_session_exit_tx, - )) + if let Some((layout, wire_id)) = session { + let prepared = PreparedStream::new(send, recv, stream, prelude, context); + return match pending_sessions.insert(&layout, wire_id, prepared) { + Ok(Some(streams)) => Some(spawn_service_session( + admission.workers, + streams, + per_stream_queue_depth, + false, + session_exit_tx, + )), + Ok(None) => None, + Err(error) => { + debug!(?error, "rejecting mismatched service session"); + admission.close_cause.record("invalid_service_session"); + admission.connection_token.cancel(); + None + } + }; } + admission.workers.spawn(request_stream_worker( + send, + recv, + prelude, + context, + self.registry.clone(), + )); + None } async fn register_and_serve( @@ -3176,7 +3286,7 @@ impl ZakuraProtocolHandler { // Count each symmetric session regardless of which endpoint opens it. let ordered_stream_count = self .registry - .ordered_streams_for_escalation( + .persistent_streams_for_escalation( context.accepted_capabilities, &peer_id, context.direction, @@ -3938,6 +4048,7 @@ async fn run_native_initiator_handshake( } #[allow(clippy::too_many_arguments)] +#[cfg(test)] fn spawn_persistent_stream_worker( workers: &mut JoinSet<()>, send: SendStream, @@ -3947,42 +4058,59 @@ fn spawn_persistent_stream_worker( context: StreamWorkerContext, queue_depth: usize, opened_locally: bool, - ordered_session_exit_tx: mpsc::UnboundedSender, -) -> AdmittedOrderedSession { - let (to_service_tx, to_service_rx) = mpsc::channel(queue_depth); - let (from_service_tx, from_service_rx) = worker_framed_channel(queue_depth); - let admitted = AdmittedOrderedSession { - kind: prelude.stream_kind, - version: prelude.stream_version, - session_id: context.stream_id, - recv: FramedRecv::new(to_service_rx), - send: from_service_tx, - cancel_token: context.stream_token.clone(), - }; - - let exit = OrderedSessionExit { - stream, - session_id: admitted.session_id, + session_exit_tx: mpsc::UnboundedSender, +) -> AdmittedSession { + spawn_service_session( + workers, + vec![PreparedStream::new(send, recv, stream, prelude, context)], + queue_depth, opened_locally, - }; - workers.spawn(async move { - persistent_stream_worker( - send, - recv, - prelude, - context, - to_service_tx, - from_service_rx, - queue_depth, + session_exit_tx, + ) +} + +fn bounded_stream_queue_depths( + transport_depth: usize, + service: Option<(usize, usize)>, +) -> (usize, usize) { + service.map_or((transport_depth, transport_depth), |(inbound, outbound)| { + ( + transport_depth.min(inbound.max(1)), + transport_depth.min(outbound.max(1)), ) - .await; - let _ = ordered_session_exit_tx.send(exit); - }); + }) +} - admitted +#[derive(Debug, Error)] +#[error("Zakura outbound frame write timed out")] +struct OrderedFrameWriteTimeout; + +#[cfg(test)] +async fn persistent_stream_worker( + send: SendStream, + recv: RecvStream, + prelude: StreamPrelude, + context: StreamWorkerContext, + inbound_tx: mpsc::Sender, + outbound_rx: FramedWorkerRecv, + queue_depth_limit: usize, +) { + persistent_stream_worker_with_policy( + send, + recv, + prelude, + context, + inbound_tx, + outbound_rx, + queue_depth_limit, + None, + None, + ) + .await; } -async fn persistent_stream_worker( +#[allow(clippy::too_many_arguments)] +async fn persistent_stream_worker_with_policy( mut send: SendStream, recv: RecvStream, prelude: StreamPrelude, @@ -3990,24 +4118,23 @@ async fn persistent_stream_worker( inbound_tx: mpsc::Sender, outbound_rx: FramedWorkerRecv, queue_depth_limit: usize, + failure_cause: Option, + application_drain: Option>, ) { let context = Arc::new(context); let stream_kind = prelude.stream_kind; - // The inbound reader runs in its own task, not as a `select!` branch racing - // the outbound writer below. `read_frame` is NOT cancellation-safe: it - // consumes the fixed frame header, then awaits the (multi-packet) payload. If - // it shared this `select!` with the outbound arm, an outbound frame becoming - // ready mid-read would drop the `read_frame` future and discard the header - // bytes it already consumed, desyncing the stream forever -- the next read - // decodes body bytes as a header, yielding a garbage multi-GiB `payload_len`, - // an `OversizeFrame` error, and a stream reset. Heavy concurrent body-sync - // (inbound bodies + outbound `GetBlocks`) made this fire constantly. Reading - // in a dedicated task removes the write/read race; the main loop only ever - // *receives* fully-read frames over a channel, which is cancellation-safe. - let (frame_tx, mut frame_rx) = mpsc::channel::>(1); + // Reading and forwarding may both block. Keep them independent of writes + // so waiting for inbound channel space cannot stall an outgoing response. + // A dedicated reader also preserves partial frame reads across outbound writes. + let (error_tx, mut error_rx) = mpsc::channel::(1); + let inbound_closed = inbound_tx.clone(); let reader_context = Arc::clone(&context); - let reader = tokio::spawn(async move { + let reader_failure_cause = failure_cause.clone(); + let reader = tokio_util::task::AbortOnDropHandle::new(tokio::spawn(async move { + // Session teardown must interrupt a blocked writer without waiting for it + // to poll the reader's terminal event. Preserve the close cause first. + let _cancel_session_on_exit = reader_context.stream_token.clone().drop_guard(); let mut recv = recv; loop { let frame = tokio::select! { @@ -4017,6 +4144,8 @@ async fn persistent_stream_worker( frame = read_frame( &mut recv, reader_context.inbound_frame_cap, + reader_context.message_payload_limits, + reader_context.message_types, reader_context.limits.idle_timeout, // A persistent ordered stream is legitimately quiet between // frames; do not let an inter-frame gap time out and cancel @@ -4026,20 +4155,36 @@ async fn persistent_stream_worker( None, ) => frame, }; - // Admit (rate/oversize) at ingress, the instant a frame is read, so - // throttling never trails behind the main loop draining queued - // outbound writes or forwarding an earlier frame to a service that - // might disconnect first. The main loop only ever receives frames - // that already cleared admission, plus terminal errors it maps to a - // reset code (it owns the send half). Admission is charged exactly - // once, here. - let message = match frame { + // Charge ingress admission once, before forwarding to the service. + // A full service channel stops this reader from granting QUIC credit. + let error = match frame { Ok(frame) => { let _ = reader_context.freshness_tx.send(Instant::now()); match admit_inbound_message(frame.payload.len(), &reader_context, stream_kind) { - InboundMessageAdmission::Admit => Ok(frame), - InboundMessageAdmission::Oversize => Err(ZakuraHandlerError::Oversize), - InboundMessageAdmission::Throttled => Err(ZakuraHandlerError::RateLimited), + InboundMessageAdmission::Admit => { + let forwarded = tokio::select! { + biased; + _ = reader_context.connection_token.cancelled() => break, + _ = reader_context.stream_token.cancelled() => break, + result = inbound_tx.send(frame) => result, + }; + if forwarded.is_err() { + // Local receiver closure leaves retained senders and + // queued writes alive. Keep bounded ingress checks and + // real peer-close detection until the writers retire. + continue; + } + // Queue depth is capped by the negotiated u16 limit, so + // this integer is exactly representable as f64. + metrics::gauge!( + "zakura.p2p.queue.depth", + "stream_kind" => stream_kind_label(stream_kind), + ) + .set(queue_depth_limit.saturating_sub(inbound_tx.capacity()) as f64); + continue; + } + InboundMessageAdmission::Oversize => ZakuraHandlerError::Oversize, + InboundMessageAdmission::Throttled => ZakuraHandlerError::RateLimited, } } Err(error) => { @@ -4058,7 +4203,7 @@ async fn persistent_stream_worker( Some(max_frame_bytes), ); } - Err(error) + error } }; // Any error is terminal. `Closed` is a clean peer-initiated close; @@ -4067,26 +4212,44 @@ async fn persistent_stream_worker( // map it to a reset code, then cancel the connection ourselves so // the disconnect is guaranteed even if the main loop tore the worker // down for a stopped outbound write before processing it. - let is_terminal = message.is_err(); - let must_disconnect = - matches!(&message, Err(error) if !matches!(error, ZakuraHandlerError::Closed)); - let forward_failed = frame_tx.send(message).await.is_err(); + let must_disconnect = !matches!(error, ZakuraHandlerError::Closed); + if !must_disconnect && !reader_context.stream_token.is_cancelled() { + if let Some(cause) = &reader_failure_cause { + // Publish the cause before waking either the service EOF or + // the sibling worker's cancellation path. + cause.record(OrderedStreamFailure::RemoteClose); + } + } + let _ = error_tx.send(error).await; if must_disconnect { reader_context.close_cause.record("ordered_read_error"); reader_context.connection_token.cancel(); } - if forward_failed || is_terminal { - break; - } + break; } - }); + })); let mut outbound_rx = Some(outbound_rx); + let mut drained = false; loop { tokio::select! { biased; _ = context.connection_token.cancelled() => break, _ = context.stream_token.cancelled() => break, + _ = inbound_closed.closed(), if outbound_rx.is_none() => { + drained = true; + // An idle member must not cancel another member's queued writes + // or retained application handles. Faults still interrupt draining. + if let Some(barrier) = &application_drain { + tokio::select! { + biased; + _ = context.connection_token.cancelled() => {}, + _ = context.stream_token.cancelled() => {}, + _ = barrier.wait() => {}, + } + } + break; + } outbound = async { match outbound_rx.as_mut() { Some(outbound_rx) => outbound_rx.recv().await, @@ -4095,17 +4258,41 @@ async fn persistent_stream_worker( } => { match outbound { Some(queued_frame) => { - if let Err(error) = write_queued_ordered_frame( - &mut send, - queued_frame, - context.limits, - context.outbound_frame_cap, - ).await { + // Cancellation resets the stream before a replacement can write. + let result = tokio::select! { + biased; + _ = context.connection_token.cancelled() => break, + _ = context.stream_token.cancelled() => break, + result = queued_frame.write_with(|frame| async { + let result = write_ordered_frame_with_policy( + &mut send, frame, context.limits, + context.outbound_frame_cap, context.write_policy, + ).await; + // A failed request claim can cancel the session on + // drop. Record the cause while it is still alive. + if !context.stream_token.is_cancelled() { + if let (Err(error), Some(cause)) = (&result, &failure_cause) { + if error.is::() { + cause.record(OrderedStreamFailure::WriteTimeout); + } else if ordered_stream_write_was_stopped(error) { + cause.record(OrderedStreamFailure::RemoteClose); + } + } + } + result + }) => result, + }; + if let Err(error) = result { + if error.is::() { + debug!(stream_kind, stream_id = context.stream_id, + "retiring Zakura service session after stream write timeout"); + break; + } if ordered_stream_write_was_stopped(&error) { debug!(?error, "closing Zakura ordered stream after peer stopped receiving"); break; } - debug!(?error, "closing Zakura ordered stream writer"); + debug!(stream_kind, stream_id = context.stream_id, ?error, "closing Zakura ordered stream writer"); let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE)); context.close_cause.record("ordered_write_error"); context.connection_token.cancel(); @@ -4117,41 +4304,25 @@ async fn persistent_stream_worker( } } } - inbound = frame_rx.recv() => { - match inbound { - // Frames here already cleared ingress admission in the reader. - Some(Ok(frame)) => { - if inbound_tx.send(frame).await.is_err() { - debug!( - stream_kind, - "closing Zakura ordered stream after local service receiver dropped" - ); - break; - } - metrics::gauge!( - "zakura.p2p.queue.depth", - "stream_kind" => stream_kind_label(stream_kind), - ) - .set(queue_depth_limit.saturating_sub(inbound_tx.capacity()) as f64); - } - // The reader signalled an oversize message: disconnect it. - Some(Err(ZakuraHandlerError::Oversize)) => { + error = error_rx.recv() => { + match error { + Some(ZakuraHandlerError::Oversize) => { let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_OVERSIZE)); context.close_cause.record("ordered_oversize"); context.connection_token.cancel(); break; } - Some(Err(ZakuraHandlerError::RateLimited)) => { + Some(ZakuraHandlerError::RateLimited) => { let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_RATE_LIMIT)); context.close_cause.record("ordered_rate_limited"); context.connection_token.cancel(); break; } - Some(Err(ZakuraHandlerError::Closed)) | None => { + Some(ZakuraHandlerError::Closed) | None => { break; } // The reader already emitted any oversize-desync diagnostic. - Some(Err(error)) => { + Some(error) => { debug!(?error, "closing Zakura stream worker"); let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_BAD_PRELUDE)); context.close_cause.record("ordered_read_error"); @@ -4163,10 +4334,17 @@ async fn persistent_stream_worker( } } - // Stop the reader: it also observes the cancellation tokens, but abort - // guarantees a prompt exit on the paths that break without cancelling one - // (e.g. a peer that stopped receiving, or the local service receiver closing). + // Never leave a partial frame followed by a graceful FIN. + if drained { + // A FIN preserves completed writes in QUIC's send buffer. + let _ = send.finish(); + } else { + let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_NEUTRAL)); + } + context.stream_token.cancel(); reader.abort(); + // Keep the stream permit until the reader has actually dropped its buffers. + let _ = reader.await; } fn ordered_stream_write_was_stopped(error: &BoxError) -> bool { @@ -4196,6 +4374,8 @@ async fn request_stream_worker( frame = read_frame( &mut recv, context.inbound_frame_cap, + context.message_payload_limits, + context.message_types, context.limits.idle_timeout, // A request stream carries its request frame immediately after the // prelude, so a peer that opens one and then goes silent is treated @@ -4367,6 +4547,8 @@ async fn read_stream_prelude( async fn read_frame( recv: &mut RecvStream, max_frame_bytes: u32, + message_payload_limits: &[(u16, usize)], + message_types: Option<&[u16]>, read_timeout: Duration, first_byte_timeout: Option, ) -> Result { @@ -4396,11 +4578,22 @@ async fn read_frame( } let mut reader = &header[..]; let message_type = reader.read_u16::()?; + if message_types.is_some_and(|types| !types.contains(&message_type)) { + return Err(ZakuraHandlerError::InvalidMessageType(message_type)); + } let flags = reader.read_u16::()?; let payload_len = usize::try_from(reader.read_u32::()?) .expect("u32 payload lengths fit usize on supported targets"); let max_frame_bytes = usize::try_from(max_frame_bytes).expect("u32 frame cap fits usize on supported targets"); + // A service may declare a tighter limit for this message. Apply it before + // allocating the payload; it can never enlarge the negotiated stream cap. + let max_frame_bytes = message_payload_limits + .iter() + .find(|(kind, _)| *kind == message_type) + .map_or(max_frame_bytes, |(_, max_payload_bytes)| { + max_frame_bytes.min(max_payload_bytes.saturating_add(FRAME_HEADER_BYTES)) + }); let frame_len = FRAME_HEADER_BYTES.saturating_add(payload_len); if frame_len > max_frame_bytes { metrics::counter!("zakura.p2p.ratelimit.frame.oversize").increment(1); @@ -4411,9 +4604,7 @@ async fn read_frame( }); } let mut payload = vec![0; payload_len]; - timeout(read_timeout, recv.read_exact(&mut payload)) - .await - .map_err(|_| ZakuraHandlerError::Timeout("frame payload"))??; + read_frame_payload(recv, &mut payload, read_timeout).await?; Ok(Frame { message_type, flags, @@ -4421,6 +4612,25 @@ async fn read_frame( }) } +async fn read_frame_payload( + recv: &mut RecvStream, + payload: &mut [u8], + read_timeout: Duration, +) -> Result<(), ZakuraHandlerError> { + match timeout(read_timeout, recv.read_exact(payload)) + .await + .map_err(|_| ZakuraHandlerError::Timeout("frame payload"))? + { + Ok(()) => Ok(()), + // A peer can reset a pair during a partial frame. Reset ends that stream; + // a normal FIN with a truncated payload still reports a protocol error. + Err(iroh::endpoint::ReadExactError::ReadError(iroh::endpoint::ReadError::Reset(_))) => { + Err(ZakuraHandlerError::Closed) + } + Err(error) => Err(error.into()), + } +} + async fn read_control_payload( recv: &mut RecvStream, max_bytes: u32, @@ -4467,11 +4677,29 @@ async fn write_control_payload( Ok(()) } +#[cfg(test)] async fn write_ordered_frame( send: &mut SendStream, frame: Frame, limits: ZakuraConnectionLimits, max_frame_bytes: u32, +) -> Result<(), BoxError> { + write_ordered_frame_with_policy( + send, + frame, + limits, + max_frame_bytes, + StreamWritePolicy::Timeout(OUTBOUND_STREAM_WRITE_TIMEOUT), + ) + .await +} + +async fn write_ordered_frame_with_policy( + send: &mut SendStream, + frame: Frame, + limits: ZakuraConnectionLimits, + max_frame_bytes: u32, + write_policy: StreamWritePolicy, ) -> Result<(), BoxError> { // Mirror `write_response_frame`: a persistent ordered-stream frame whose // payload exceeds the peer's negotiated `max_message_bytes` would be @@ -4488,31 +4716,24 @@ async fn write_ordered_frame( .into()); } let frame = frame.encode(max_frame_bytes)?; - timeout(OUTBOUND_STREAM_WRITE_TIMEOUT, send.write_all(&frame)) - .await - .map_err(|_| -> BoxError { "Zakura outbound frame write timed out".into() })??; + match write_policy { + StreamWritePolicy::UntilCancelled => send.write_all(&frame).await?, + StreamWritePolicy::Timeout(duration) => { + timeout(duration, send.write_all(&frame)) + .await + .map_err(|_| -> BoxError { Box::new(OrderedFrameWriteTimeout) })?? + } + } Ok(()) } -/// Write one queued frame while retaining its byte-accounting lease. -/// -/// The lease remains owned until the QUIC write succeeds, fails, times out, or -/// this future is dropped with the worker. -async fn write_queued_ordered_frame( - send: &mut SendStream, - queued_frame: QueuedFrame, - limits: ZakuraConnectionLimits, - max_frame_bytes: u32, -) -> Result<(), BoxError> { - queued_frame - .write_with(|frame| write_ordered_frame(send, frame, limits, max_frame_bytes)) - .await -} - +#[allow(clippy::too_many_arguments)] async fn write_outbound_request_frame( connection: &Connection, limits: ZakuraConnectionLimits, stream: Stream, + message_payload_limits: &'static [(u16, usize)], + message_types: Option<&'static [u16]>, request_id: u64, message_type: u16, flags: u16, @@ -4524,6 +4745,8 @@ async fn write_outbound_request_frame( connection, limits, stream, + message_payload_limits, + message_types, request_id, message_type, flags, @@ -4534,10 +4757,13 @@ async fn write_outbound_request_frame( .map_err(|_| OutboundRequestError::Local("Zakura outbound request/response timed out".into()))? } +#[allow(clippy::too_many_arguments)] async fn write_outbound_request_frame_inner( connection: &Connection, limits: ZakuraConnectionLimits, stream: Stream, + message_payload_limits: &'static [(u16, usize)], + message_types: Option<&'static [u16]>, request_id: u64, message_type: u16, flags: u16, @@ -4591,6 +4817,8 @@ async fn write_outbound_request_frame_inner( match read_frame( &mut recv, inbound_frame_cap, + message_payload_limits, + message_types, limits.idle_timeout, // This is the requester side of a one-shot legacy request/response: // the responder streams its frames promptly, so a silent gap before @@ -5347,6 +5575,12 @@ impl TokenBucket { /// Errors produced by the Zakura protocol handler. #[derive(Debug, Error)] pub enum ZakuraHandlerError { + /// The frame header names a message that is invalid on this stream role. + #[error("invalid message type {0} for this stream role")] + InvalidMessageType(u16), + /// Two ordered stream roles failed to name one complete session. + #[error("invalid Zakura service session")] + InvalidServiceSession, /// A bounded read/write timed out. #[error("Zakura {0} timed out")] Timeout(&'static str), @@ -5381,6 +5615,9 @@ pub enum ZakuraHandlerError { /// A local resource cap rejected the operation. #[error("Zakura resource limit exceeded: {0}")] ResourceLimit(&'static str), + /// Another connection reserved the last service session slot. + #[error("ordered service session capacity is full")] + SessionFull, /// The peer exceeded its per-kind inbound message rate. #[error("Zakura message rate exceeded")] RateLimited, @@ -5431,6 +5668,7 @@ impl ZakuraHandlerError { #[cfg(test)] mod tests { + pub(super) mod connection; mod quic_progress; use super::*; use crate::{ @@ -5478,7 +5716,8 @@ mod tests { #[tokio::test] async fn parked_block_sync_peer_gets_a_stream_when_its_cooldown_lapses() -> Result<(), BoxError> { - const COOLDOWN: Duration = Duration::from_secs(3); + const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); + const COOLDOWN: Duration = Duration::from_secs(15); let _guard = zakura_test::init(); @@ -5497,20 +5736,28 @@ mod tests { let dialer = node(140).await?; let listener = node(141).await?; - let listener_peer = ZakuraPeerId::new(listener.node_addr().await.id.as_bytes().to_vec())?; + let listener_addr = listener.node_addr().await; + let listener_peer = ZakuraPeerId::new(listener_addr.id.as_bytes().to_vec())?; let block_sync = dialer .block_sync() .expect("the header-sync driver spawns the block-sync reactor"); // Block sync evicts and parks the peer after its no-progress deadline. - // The transport redials during the cooldown. + // Keep the cooldown longer than connection setup so the first assertion + // tests a live park even when the dial is slow. + let parked_at = std::time::Instant::now(); block_sync.park_peer_for_test(&listener_peer, COOLDOWN); dialer - .connect_native(&listener, Duration::from_secs(10)) + .connect_native_to_addr(listener_addr, CONNECT_TIMEOUT) .await?; // The park remains active. // Withhold block sync from this connection. + assert!( + parked_at.elapsed() < COOLDOWN, + "connection setup outlasted the test cooldown: {:?}", + parked_at.elapsed(), + ); assert_eq!( block_sync.peer_snapshot().outbound_peers, 0, @@ -6241,7 +6488,7 @@ mod tests { version: 1, frame_cap: 1024, capability: ZAKURA_CAP_LEGACY_GOSSIP, - mode: StreamMode::Ordered, + mode: StreamMode::Persistent, }; let service = GenerationGuardedRecordingService::new(vec![stream]); let registry = Arc::new( @@ -6356,7 +6603,7 @@ mod tests { version: 1, frame_cap: 1024, capability: ZAKURA_CAP_LEGACY_GOSSIP, - mode: StreamMode::Ordered, + mode: StreamMode::Persistent, }; let service = GenerationGuardedRecordingService::new(vec![stream]); let registry = Arc::new( @@ -6833,7 +7080,7 @@ mod tests { version: 1, frame_cap: 64 * 1024, capability: 1 << 16, - mode: StreamMode::Ordered, + mode: StreamMode::Persistent, }; let _guard = zakura_test::init(); @@ -6964,8 +7211,7 @@ mod tests { } #[tokio::test] - async fn header_ordered_session_waits_for_coordinator_capability_epoch() -> Result<(), BoxError> - { + async fn header_session_waits_for_coordinator_capability_epoch() -> Result<(), BoxError> { use zakura_node_services::sync_lifecycle::{ BlockServiceDemand, HeaderServiceDemand, LifecycleEpoch, SyncServiceDemand, }; @@ -6989,7 +7235,7 @@ mod tests { ZAKURA_CAP_HEADER_SYNC, ServicePeerDirection::Outbound, )); - let OrderedSessionDemand::WaitForChange(changed) = service.ordered_session_demand( + let SessionDemand::WaitForChange(changed) = service.session_demand( test_conn_id(), &peer, ZAKURA_CAP_HEADER_SYNC, @@ -7452,15 +7698,15 @@ mod tests { version: ZAKURA_HEADER_SYNC_STREAM_VERSION, frame_cap: 1, capability: ZAKURA_CAP_HEADER_SYNC, - mode: StreamMode::Ordered, + mode: StreamMode::Persistent, }; - let exit = OrderedSessionExit { + let exit = SessionExit { stream: initiator_opened, session_id: 1, opened_locally: true, }; - let initiator_policy = OrderedStreamPolicy { - opening: OrderedStreamOpening::InitiatorOnly, + let initiator_policy = SessionPolicy { + opening: SessionOpening::InitiatorOnly, reopen: true, }; assert!(should_reopen_ordered_session( @@ -7486,7 +7732,7 @@ mod tests { )); assert!(!should_reopen_ordered_session( exit, - OrderedStreamPolicy::default(), + SessionPolicy::default(), true, false, false, @@ -7495,7 +7741,7 @@ mod tests { // The authorized side replaces an accepted stream too; replacement does // not depend on which physical generation just exited. assert!(should_reopen_ordered_session( - OrderedSessionExit { + SessionExit { opened_locally: false, ..exit }, @@ -7505,15 +7751,15 @@ mod tests { false, )); - let either_peer = OrderedSessionExit { + let either_peer = SessionExit { stream: Stream { kind: ZAKURA_STREAM_BLOCK_SYNC, ..initiator_opened }, ..exit }; - let either_policy = OrderedStreamPolicy { - opening: OrderedStreamOpening::EitherSide, + let either_policy = SessionPolicy { + opening: SessionOpening::EitherSide, reopen: true, }; assert!(should_reopen_ordered_session( @@ -7533,7 +7779,7 @@ mod tests { assert!(opens_ordered_stream_locally(either_policy, false, true)); assert!(!opens_ordered_stream_locally(either_policy, true, false)); - let request_response = OrderedSessionExit { + let request_response = SessionExit { stream: Stream { kind: 102, mode: StreamMode::RequestResponse, @@ -7552,17 +7798,17 @@ mod tests { #[test] fn either_side_session_has_one_proactive_opener_across_connection_roles() { - let policy = OrderedStreamPolicy { - opening: OrderedStreamOpening::EitherSide, + let policy = SessionPolicy { + opening: SessionOpening::EitherSide, reopen: true, }; - let exit = OrderedSessionExit { + let exit = SessionExit { stream: Stream { kind: ZAKURA_STREAM_BLOCK_SYNC, version: ZAKURA_BLOCK_SYNC_STREAM_VERSION, frame_cap: 1, capability: ZAKURA_CAP_BLOCK_SYNC, - mode: StreamMode::Ordered, + mode: StreamMode::Persistent, }, session_id: 1, opened_locally: false, @@ -7597,9 +7843,9 @@ mod tests { version: ZAKURA_HEADER_SYNC_STREAM_VERSION, frame_cap: 1, capability: ZAKURA_CAP_HEADER_SYNC, - mode: StreamMode::Ordered, + mode: StreamMode::Persistent, }; - let mut session = OrderedSessionState::new(stream); + let mut session = ServiceSessionState::new(stream); session.remote_session_id = Some(2); assert!(!session.remove_active_session(false, 1)); @@ -7614,33 +7860,33 @@ mod tests { version: ZAKURA_HEADER_SYNC_STREAM_VERSION, frame_cap: 1, capability: ZAKURA_CAP_HEADER_SYNC, - mode: StreamMode::Ordered, + mode: StreamMode::Persistent, } } #[tokio::test(start_paused = true)] - async fn ordered_session_waits_deduplicate_demand() { + async fn session_waits_deduplicate_demand() { let stream = test_ordered_stream(); - let mut session = OrderedSessionState::new(stream); - let mut waits = OrderedSessionWaits::new(); + let mut session = ServiceSessionState::new(stream); + let mut waits = SessionWaits::new(); let (change_tx, change_rx) = watch::channel(()); session.schedule_demand( &mut waits, - OrderedSessionDemand::WaitForChange(Box::pin(async move { + SessionDemand::WaitForChange(Box::pin(async move { let mut change_rx = change_rx; let _ = change_rx.changed().await; })), ); session.schedule_demand( &mut waits, - OrderedSessionDemand::RetryAt(std::time::Instant::now() + Duration::from_secs(60)), + SessionDemand::RetryAt(std::time::Instant::now() + Duration::from_secs(60)), ); assert_eq!(waits.len(), 1); assert_eq!( session.reopen_state, - OrderedSessionReopenState::Waiting(OrderedSessionWaitReason::Demand) + SessionReopenState::Waiting(SessionWaitReason::Demand) ); assert_eq!( change_tx.receiver_count(), @@ -7650,10 +7896,10 @@ mod tests { } #[tokio::test(start_paused = true)] - async fn ordered_session_demand_replaces_transport_backoff() { + async fn session_demand_replaces_transport_backoff() { let stream = test_ordered_stream(); - let mut session = OrderedSessionState::new(stream); - let mut waits = OrderedSessionWaits::new(); + let mut session = ServiceSessionState::new(stream); + let mut waits = SessionWaits::new(); assert_eq!( session.schedule_transport_backoff(&mut waits), @@ -7661,26 +7907,26 @@ mod tests { ); session.schedule_demand( &mut waits, - OrderedSessionDemand::RetryAt(std::time::Instant::now() + Duration::from_secs(60)), + SessionDemand::RetryAt(std::time::Instant::now() + Duration::from_secs(60)), ); assert_eq!(session.reopen_attempts, 1); assert_eq!(waits.len(), 1); assert_eq!( session.reopen_state, - OrderedSessionReopenState::Waiting(OrderedSessionWaitReason::Demand) + SessionReopenState::Waiting(SessionWaitReason::Demand) ); } #[tokio::test(start_paused = true)] async fn ordered_session_exit_keeps_existing_demand_wait() { let stream = test_ordered_stream(); - let mut session = OrderedSessionState::new(stream); - let mut waits = OrderedSessionWaits::new(); + let mut session = ServiceSessionState::new(stream); + let mut waits = SessionWaits::new(); session.schedule_demand( &mut waits, - OrderedSessionDemand::RetryAt(std::time::Instant::now() + Duration::from_secs(60)), + SessionDemand::RetryAt(std::time::Instant::now() + Duration::from_secs(60)), ); assert_eq!(session.schedule_transport_backoff(&mut waits), None); @@ -7688,24 +7934,24 @@ mod tests { assert_eq!(waits.len(), 1); assert_eq!( session.reopen_state, - OrderedSessionReopenState::Waiting(OrderedSessionWaitReason::Demand) + SessionReopenState::Waiting(SessionWaitReason::Demand) ); } #[tokio::test] async fn ordered_session_retirement_cancels_and_blocks_transport_waits() { let stream = test_ordered_stream(); - let mut session = OrderedSessionState::new(stream); - let mut waits = OrderedSessionWaits::new(); + let mut session = ServiceSessionState::new(stream); + let mut waits = SessionWaits::new(); session.schedule_demand( &mut waits, - OrderedSessionDemand::RetryAt(std::time::Instant::now() + Duration::from_secs(60)), + SessionDemand::RetryAt(std::time::Instant::now() + Duration::from_secs(60)), ); - session.schedule_demand(&mut waits, OrderedSessionDemand::Retire); + session.schedule_demand(&mut waits, SessionDemand::Retire); assert!(waits.is_empty()); - assert_eq!(session.reopen_state, OrderedSessionReopenState::Retired); + assert_eq!(session.reopen_state, SessionReopenState::Retired); assert_eq!(session.schedule_transport_backoff(&mut waits), None); assert_eq!(session.reopen_attempts, 0); } @@ -7713,26 +7959,26 @@ mod tests { #[tokio::test] async fn ordered_session_adoption_drops_pending_wait() { let stream = test_ordered_stream(); - let mut session = OrderedSessionState::new(stream); - let mut waits = OrderedSessionWaits::new(); + let mut session = ServiceSessionState::new(stream); + let mut waits = SessionWaits::new(); - session.schedule_demand(&mut waits, OrderedSessionDemand::OpenNow); + session.schedule_demand(&mut waits, SessionDemand::OpenNow); session.cancel_wait(&mut waits); assert!(waits.is_empty()); - assert_eq!(session.reopen_state, OrderedSessionReopenState::Idle); + assert_eq!(session.reopen_state, SessionReopenState::Idle); } #[tokio::test] - async fn ordered_session_demand_change_yields_exactly_one_reopen() { + async fn session_demand_change_yields_exactly_one_reopen() { let stream = test_ordered_stream(); - let mut session = OrderedSessionState::new(stream); - let mut waits = OrderedSessionWaits::new(); + let mut session = ServiceSessionState::new(stream); + let mut waits = SessionWaits::new(); let (change_tx, mut change_rx) = watch::channel(()); session.schedule_demand( &mut waits, - OrderedSessionDemand::WaitForChange(Box::pin(async move { + SessionDemand::WaitForChange(Box::pin(async move { let _ = change_rx.changed().await; })), ); @@ -7751,14 +7997,14 @@ mod tests { } #[tokio::test(start_paused = true)] - async fn ordered_session_demand_deadline_yields_exactly_one_reopen() { + async fn session_demand_deadline_yields_exactly_one_reopen() { let stream = test_ordered_stream(); - let mut session = OrderedSessionState::new(stream); - let mut waits = OrderedSessionWaits::new(); + let mut session = ServiceSessionState::new(stream); + let mut waits = SessionWaits::new(); session.schedule_demand( &mut waits, - OrderedSessionDemand::RetryAt(std::time::Instant::now() + Duration::from_secs(60)), + SessionDemand::RetryAt(std::time::Instant::now() + Duration::from_secs(60)), ); tokio::time::advance(Duration::from_secs(60)).await; let (kind, ()) = waits @@ -7774,13 +8020,13 @@ mod tests { #[test] fn ordered_session_connection_teardown_drops_pending_waits() { let stream = test_ordered_stream(); - let mut session = OrderedSessionState::new(stream); - let mut waits = OrderedSessionWaits::new(); + let mut session = ServiceSessionState::new(stream); + let mut waits = SessionWaits::new(); let (change_tx, mut change_rx) = watch::channel(()); session.schedule_demand( &mut waits, - OrderedSessionDemand::WaitForChange(Box::pin(async move { + SessionDemand::WaitForChange(Box::pin(async move { let _ = change_rx.changed().await; })), ); @@ -7814,13 +8060,205 @@ mod tests { ); } + /// Exercise the production worker with a full application receive channel. + /// Sending must stay live, QUIC must eventually stop the bounded sender, and + /// another stream must retain connection credit until the application resumes. + #[tokio::test] + async fn paused_ordered_reads_preserve_writes_and_sibling_credit() -> Result<(), BoxError> { + const ALPN: &[u8] = b"/zakura/testkit/paused-reader/0"; + const FRAME_COUNT: usize = 24; + let local = ZakuraLocalLimits::from_config(&Config::default()); + let server = LocalEndpointFactory::with_transport_config(local.transport_config()) + .endpoint(52) + .await?; + let (conn_tx, _conn_rx) = mpsc::channel(1); + let (stream_tx, mut stream_rx) = mpsc::channel(2); + let router = Router::builder(server) + .accept( + ALPN, + CaptureConnection { + connection_tx: conn_tx, + stream_tx, + }, + ) + .spawn(); + // A small sender buffer makes write completion reflect receiver credit. + let client_transport = local + .transport_config_builder() + .send_window(64 * 1024) + .build(); + let client = LocalEndpointFactory::with_transport_config(client_transport) + .endpoint(53) + .await?; + let address = router.endpoint().addr(); + let connection = timeout(Duration::from_secs(10), client.connect(address, ALPN)).await??; + let (mut sender, mut receiver) = connection.open_bi().await?; + let frame = Frame { + message_type: 3, + flags: 0, + payload: vec![0; 1024 * 1024], + }; + let limits = test_connection_limits(); + let stream = Stream { + kind: ZAKURA_STREAM_BLOCK_SYNC, + version: ZAKURA_BLOCK_SYNC_STREAM_VERSION, + frame_cap: 2_000_009, + capability: ZAKURA_CAP_BLOCK_SYNC, + mode: StreamMode::Persistent, + }; + let encoded = frame.encode(stream.frame_cap)?; + // Opening a QUIC stream becomes visible to the receiver after the first bytes. + sender.write_all(&encoded[..1]).await?; + let (send, recv) = timeout(Duration::from_secs(5), stream_rx.recv()) + .await? + .unwrap(); + let cancel = CancellationToken::new(); + let stream_cancel = cancel.child_token(); + let (freshness_tx, _freshness_rx) = watch::channel(Instant::now()); + let context = StreamWorkerContext { + conn: ZakuraConnTrace::without_peer(1), + peer_id: test_peer(52), + stream_id: 1, + _permit: Arc::new(Semaphore::new(1)).try_acquire_owned().unwrap(), + limits, + inbound_frame_cap: stream.frame_cap, + message_payload_limits: &[], + message_types: None, + queue_depths: None, + write_policy: StreamWritePolicy::Timeout(OUTBOUND_STREAM_WRITE_TIMEOUT), + session_resources: None, + outbound_frame_cap: stream.frame_cap, + message_bucket: Arc::new(std::sync::Mutex::new(TokenBucket::new(128))), + connection_token: cancel.clone(), + stream_token: stream_cancel.clone(), + close_cause: CloseCause::new(), + freshness_tx, + }; + let prelude = StreamPrelude { + magic: STREAM_PRELUDE_MAGIC, + stream_kind: stream.kind, + stream_version: stream.version, + request_id: None, + max_frame_bytes: stream.frame_cap, + }; + let (inbound_tx, mut inbound_rx) = mpsc::channel(1); + let (outbound_tx, outbound_rx) = worker_framed_channel(1); + let mut worker = tokio_util::task::AbortOnDropHandle::new(tokio::spawn( + persistent_stream_worker(send, recv, prelude, context, inbound_tx, outbound_rx, 1), + )); + let (progress_tx, mut progress_rx) = watch::channel(0); + let mut sending = tokio_util::task::AbortOnDropHandle::new(tokio::spawn(async move { + sender.write_all(&encoded[1..]).await.unwrap(); + progress_tx.send_replace(1); + for sent in 2..=FRAME_COUNT { + sender.write_all(&encoded).await.unwrap(); + progress_tx.send_replace(sent); + } + sender + })); + timeout(Duration::from_secs(10), async { + while *progress_rx.borrow_and_update() < 16 { + progress_rx.changed().await.unwrap(); + } + }) + .await + .expect("the sender fills the default stream receive window"); + assert!( + timeout(Duration::from_millis(100), &mut sending) + .await + .is_err(), + "QUIC must stop the sender while application reads are paused" + ); + assert_eq!(inbound_rx.len(), 1); + + let response = Frame { + message_type: 4, + flags: 0, + payload: vec![7; 9], + }; + let producer = crate::zakura::regulation::SlotBudget::new(1).unwrap(); + let ownership = Arc::new(producer.try_reserve().unwrap()); + outbound_tx.try_reserve_guarded().unwrap().send( + response.clone(), + crate::zakura::transport::FrameGuard::new(ownership), + ); + let written = timeout( + Duration::from_secs(2), + read_frame( + &mut receiver, + stream.frame_cap, + &[], + None, + Duration::from_secs(2), + None, + ), + ) + .await??; + assert_eq!(written, response); + assert_eq!( + producer.reserved(), + 0, + "the response write releases serving capacity" + ); + + let (mut sibling_send, _sibling_recv) = connection.open_bi().await?; + timeout(Duration::from_secs(2), sibling_send.write_all(b"sibling")).await??; + let (_, mut sibling_read) = timeout(Duration::from_secs(2), stream_rx.recv()) + .await? + .unwrap(); + let mut bytes = [0; 7]; + timeout(Duration::from_secs(2), sibling_read.read_exact(&mut bytes)).await??; + assert_eq!(&bytes, b"sibling"); + + timeout(Duration::from_secs(10), async { + for _ in 0..FRAME_COUNT { + assert_eq!(inbound_rx.recv().await.unwrap(), frame); + } + }) + .await + .expect("resuming application reads drains all frames in order"); + let mut sender = timeout(Duration::from_secs(2), &mut sending).await??; + assert_eq!(*progress_rx.borrow(), FRAME_COUNT); + // Cancellation must also work while the reader waits on a full channel. + let empty = Frame { + message_type: 2, + flags: 0, + payload: vec![0; 9], + } + .encode(stream.frame_cap)?; + sender.write_all(&empty).await?; + sender.write_all(&empty).await?; + timeout(Duration::from_secs(2), async { + while inbound_rx.len() != 1 { + tokio::task::yield_now().await; + } + }) + .await?; + stream_cancel.cancel(); + timeout(Duration::from_secs(2), &mut worker).await??; + assert!( + !cancel.is_cancelled(), + "local stream cancellation preserves the connection" + ); + connection.close(0u32.into(), b"done"); + client.close().await; + router.shutdown().await?; + Ok(()) + } + #[tokio::test] async fn stream_cancel_closes_ordered_worker_without_connection_cancel() -> Result<(), BoxError> { const ALPN: &[u8] = b"/zakura/testkit/stream-cancel/0"; let _guard = zakura_test::init(); - let server = LocalEndpointFactory::new().endpoint(50).await?; + let server_transport = ZakuraLocalLimits::from_config(&Config::default()) + .transport_config_builder() + .send_window(64 * 1024) + .build(); + let server = LocalEndpointFactory::with_transport_config(server_transport) + .endpoint(50) + .await?; let (conn_tx, mut conn_rx) = mpsc::channel(1); let (stream_tx, mut stream_rx) = mpsc::channel(2); let router = Router::builder(server) @@ -7832,7 +8270,13 @@ mod tests { }, ) .spawn(); - let client = LocalEndpointFactory::new().endpoint(51).await?; + let client_transport = ZakuraLocalLimits::from_config(&Config::default()) + .transport_config_builder() + .stream_receive_window(VarInt::from_u32(64 * 1024)) + .build(); + let client = LocalEndpointFactory::with_transport_config(client_transport) + .endpoint(51) + .await?; let server_addr = router.endpoint().addr(); let client_conn = timeout(Duration::from_secs(10), client.connect(server_addr, ALPN)) @@ -7843,7 +8287,7 @@ mod tests { .expect("server connection is captured") .expect("capture handler sends the accepted connection"); drop(server_conn); - let (mut client_send, _client_recv) = + let (mut client_send, mut client_recv) = timeout(Duration::from_secs(1), client_conn.open_bi()) .await .expect("client opens the worker stream")?; @@ -7863,7 +8307,7 @@ mod tests { let mut limits = test_connection_limits(); limits.idle_timeout = Duration::from_millis(50); - let stream_kind = DISCOVERY_STREAM_KIND; + let stream_kind = ZAKURA_STREAM_BLOCK_SYNC; let connection_token = CancellationToken::new(); let stream_token = connection_token.child_token(); let (freshness_tx, _freshness_rx) = watch::channel(Instant::now()); @@ -7873,9 +8317,9 @@ mod tests { let stream = Stream { kind: stream_kind, version: ZAKURA_STREAM_VERSION_1, - frame_cap: LOCAL_MAX_CONTROL_FRAME_BYTES, - capability: ZAKURA_CAP_DISCOVERY, - mode: StreamMode::Ordered, + frame_cap: 2_000_009, + capability: ZAKURA_CAP_BLOCK_SYNC, + mode: StreamMode::Persistent, }; let context = StreamWorkerContext { conn: ZakuraConnTrace::without_peer(1), @@ -7884,6 +8328,11 @@ mod tests { _permit: permit, limits, inbound_frame_cap: inbound_frame_cap_for_stream(&limits, stream), + message_payload_limits: &[], + message_types: None, + queue_depths: None, + write_policy: StreamWritePolicy::Timeout(OUTBOUND_STREAM_WRITE_TIMEOUT), + session_resources: None, outbound_frame_cap: application_frame_cap(&limits, stream), message_bucket: Arc::new(std::sync::Mutex::new(TokenBucket::new(128))), connection_token: connection_token.clone(), @@ -7899,7 +8348,7 @@ mod tests { max_frame_bytes: inbound_frame_cap_for_stream(&limits, stream), }; let mut workers = JoinSet::new(); - let (ordered_session_exit_tx, mut ordered_session_exit_rx) = mpsc::unbounded_channel(); + let (session_exit_tx, mut session_exit_rx) = mpsc::unbounded_channel(); let admitted = spawn_persistent_stream_worker( &mut workers, server_send, @@ -7909,13 +8358,34 @@ mod tests { context, 1, true, - ordered_session_exit_tx, + session_exit_tx, ); + let response = Frame { + message_type: 3, + flags: 0, + payload: vec![7; 1024 * 1024], + }; + admitted.streams[0].send.try_send(response.clone()).unwrap(); + // Reading the header proves the write has started. The much smaller + // QUIC windows keep the remaining payload blocked until we drain it. + let mut header = [0; FRAME_HEADER_BYTES]; + timeout(Duration::from_secs(2), client_recv.read_exact(&mut header)).await??; + assert_eq!( + &header[..], + &response.encode(stream.frame_cap)?[..FRAME_HEADER_BYTES] + ); admitted.cancel_token.cancel(); + let mut payload = vec![0; response.payload.len()]; + assert!( + timeout(Duration::from_secs(2), client_recv.read_exact(&mut payload)) + .await? + .is_err(), + "session cancellation resets an unfinished frame" + ); // The exit must be reported, or the connection loop never prunes the dead // generation and never reopens the stream. - let exited = timeout(Duration::from_secs(1), ordered_session_exit_rx.recv()) + let exited = timeout(Duration::from_secs(1), session_exit_rx.recv()) .await .expect("stream cancellation reports worker exit") .expect("exit channel stays open"); @@ -8055,7 +8525,7 @@ mod tests { version: ZAKURA_STREAM_VERSION_1, frame_cap: LOCAL_MAX_CONTROL_FRAME_BYTES, capability: ZAKURA_CAP_LEGACY_GOSSIP, - mode: StreamMode::Ordered, + mode: StreamMode::Persistent, }; let context = StreamWorkerContext { conn: ZakuraConnTrace::without_peer(1), @@ -8064,6 +8534,11 @@ mod tests { _permit: permit, limits, inbound_frame_cap: inbound_frame_cap_for_stream(&limits, stream), + message_payload_limits: &[], + message_types: None, + queue_depths: None, + write_policy: StreamWritePolicy::Timeout(OUTBOUND_STREAM_WRITE_TIMEOUT), + session_resources: None, outbound_frame_cap: application_frame_cap(&limits, stream), message_bucket: Arc::new(std::sync::Mutex::new(TokenBucket::new( limits.message_rate_per_second, @@ -8194,7 +8669,7 @@ mod tests { version: ZAKURA_STREAM_VERSION_1, frame_cap: LOCAL_MAX_CONTROL_FRAME_BYTES, capability: ZAKURA_CAP_DISCOVERY, - mode: StreamMode::Ordered, + mode: StreamMode::Persistent, }; let frame_cap = application_frame_cap(&limits, stream); @@ -8341,6 +8816,166 @@ mod tests { Ok(()) } + #[tokio::test] + async fn message_payload_limits_apply_before_payload_reads() -> Result<(), BoxError> { + use crate::zakura::{block_sync_streams, BlockSyncMessage}; + use zakura_chain::{block, serialization::ZcashDeserializeInto}; + + #[derive(Debug)] + struct PayloadLimitedService(Stream); + + impl Service for PayloadLimitedService { + fn name(&self) -> &'static str { + "payload-limited" + } + + fn streams(&self) -> &[Stream] { + std::slice::from_ref(&self.0) + } + + fn message_payload_limits(&self, stream: Stream) -> &'static [(u16, usize)] { + if stream == self.0 { + &[(2, 9)] + } else { + &[] + } + } + + fn add_peer(&self, _peer: Peer) {} + + fn remove_peer(&self, _peer: &ZakuraPeerId, _conn_id: ZakuraConnId) {} + } + + let stream = block_sync_streams()[0]; + let registry = ServiceRegistry::new(vec![Arc::new(PayloadLimitedService(stream))])?; + let payload_limits = registry.message_payload_limits(stream); + assert_eq!(payload_limits, &[(2, 9)]); + assert!(registry + .message_payload_limits(Stream { + version: stream.version + 1, + ..stream + }) + .is_empty()); + assert!(registry + .message_payload_limits(Stream { + kind: u16::MAX, + ..stream + }) + .is_empty()); + assert!(NoopService.message_payload_limits(stream).is_empty()); + + const ALPN: &[u8] = b"/zakura/testkit/message-payload-limits/0"; + let _guard = zakura_test::init(); + let server = LocalEndpointFactory::new().endpoint(79).await?; + let (conn_tx, _conn_rx) = mpsc::channel(8); + let (stream_tx, mut stream_rx) = mpsc::channel(8); + let router = Router::builder(server) + .accept( + ALPN, + CaptureConnection { + connection_tx: conn_tx, + stream_tx, + }, + ) + .spawn(); + let client = LocalEndpointFactory::new().endpoint(80).await?; + let server_addr = router.endpoint().addr(); + + // Only send headers and keep the send sides open. The reader must reject + // before waiting for a payload that the peer has not supplied. + for (message_type, payload_len, frame_cap, expected_cap) in [ + // Eight frame-header bytes plus the nine-byte GetBlocks payload cap. + (2u16, 10u32, stream.frame_cap, 17usize), + (2, u32::MAX, stream.frame_cap, 17), + // A tighter stream cap still applies to an otherwise legal request. + (2, 9, 16, 16), + // Block has no message-specific cap yet; its stream cap still applies. + (3, 100, 107, 107), + ] { + let connection = timeout( + Duration::from_secs(5), + client.connect(server_addr.clone(), ALPN), + ) + .await??; + let (mut send, _recv) = timeout(Duration::from_secs(2), connection.open_bi()).await??; + let mut header = Vec::with_capacity(FRAME_HEADER_BYTES); + header.extend_from_slice(&message_type.to_le_bytes()); + header.extend_from_slice(&0u16.to_le_bytes()); + header.extend_from_slice(&payload_len.to_le_bytes()); + timeout(Duration::from_secs(2), send.write_all(&header)).await??; + let (_, mut recv) = timeout(Duration::from_secs(2), stream_rx.recv()) + .await? + .unwrap(); + let result = timeout( + Duration::from_secs(1), + read_frame( + &mut recv, + frame_cap, + payload_limits, + None, + Duration::from_secs(5), + Some(Duration::from_secs(5)), + ), + ) + .await + .expect("an oversized header is rejected without waiting for payload bytes"); + assert!( + matches!(result, Err(ZakuraHandlerError::OversizeFrame { max_frame_bytes, .. }) if max_frame_bytes == expected_cap) + ); + } + + // The generic payload gate preserves the independent Block allowance. + // Role enforcement is covered by the paired-stream tests. + let connection = + timeout(Duration::from_secs(5), client.connect(server_addr, ALPN)).await??; + let (mut send, _recv) = timeout(Duration::from_secs(2), connection.open_bi()).await??; + let messages = [ + BlockSyncMessage::GetBlocks { + start_height: block::Height(1), + count: 1, + }, + BlockSyncMessage::Block( + zakura_test::vectors::BLOCK_MAINNET_1_BYTES.zcash_deserialize_into()?, + ), + ]; + let first = messages[0].encode_frame()?; + timeout( + Duration::from_secs(2), + send.write_all(&first.encode(stream.frame_cap)?), + ) + .await??; + let (_, mut recv) = timeout(Duration::from_secs(2), stream_rx.recv()) + .await? + .unwrap(); + for (index, message) in messages.into_iter().enumerate() { + if index > 0 { + let frame = message.encode_frame()?; + assert!(frame.payload.len() > 9); + timeout( + Duration::from_secs(2), + send.write_all(&frame.encode(stream.frame_cap)?), + ) + .await??; + } + let frame = timeout( + Duration::from_secs(2), + read_frame( + &mut recv, + stream.frame_cap, + payload_limits, + None, + Duration::from_secs(2), + Some(Duration::from_secs(2)), + ), + ) + .await??; + assert_eq!(BlockSyncMessage::decode_frame(frame)?, message); + } + client.close().await; + router.shutdown().await?; + Ok(()) + } + // claude-late-message-cap-allocation: read_frame checks only // frame_len > max_frame_bytes before `vec![0; payload_len]`, while the smaller // max_message_bytes is enforced later in admit_inbound_message. A peer can @@ -8366,7 +9001,7 @@ mod tests { version: ZAKURA_STREAM_VERSION_1, frame_cap: LOCAL_MAX_CONTROL_FRAME_BYTES, capability: ZAKURA_CAP_LEGACY_GOSSIP, - mode: StreamMode::Ordered, + mode: StreamMode::Persistent, }; let limits = ZakuraConnectionLimits { @@ -8439,6 +9074,8 @@ mod tests { let rejected = read_frame( &mut s1_recv, inbound_cap, + &[], + None, Duration::from_secs(2), Some(Duration::from_secs(2)), ) @@ -8472,6 +9109,8 @@ mod tests { let allocated = read_frame( &mut s2_recv, raw_cap, + &[], + None, Duration::from_secs(2), Some(Duration::from_secs(2)), ) @@ -8510,6 +9149,8 @@ mod tests { let frame = read_frame( &mut s3_recv, inbound_cap, + &[], + None, Duration::from_secs(2), Some(Duration::from_secs(2)), ) @@ -8605,6 +9246,8 @@ mod tests { let (freshness_tx, _freshness_rx) = watch::channel(Instant::now()); let mut admission = StreamAdmission { + is_initiator: false, + direction: ServicePeerDirection::Inbound, conn: ZakuraConnTrace::placeholder(), peer_id: &peer_id, stream_sem: &stream_sem, @@ -8617,14 +9260,15 @@ mod tests { close_cause: CloseCause::new(), freshness_tx, }; - let (ordered_session_exit_tx, _ordered_session_exit_rx) = mpsc::unbounded_channel(); + let (session_exit_tx, _session_exit_rx) = mpsc::unbounded_channel(); let admitted = handler .admit_bi_stream( server_send, server_recv, &mut admission, 16, - ordered_session_exit_tx, + session_exit_tx, + &mut PendingSessions::default(), ) .await; @@ -8748,46 +9392,53 @@ mod tests { #[test] fn supported_stream_accepts_registered_kinds_at_declared_version_only() { - let registry = ServiceRegistry::new(vec![Arc::new(DeclaredStreamService { - streams: vec![ - Stream { - kind: LEGACY_GOSSIP_STREAM_KIND, - version: ZAKURA_STREAM_VERSION_1, - frame_cap: 1024, - capability: ZAKURA_CAP_LEGACY_GOSSIP, - mode: StreamMode::Ordered, - }, - Stream { - kind: LEGACY_REQUEST_STREAM_KIND, - version: ZAKURA_STREAM_VERSION_1, - frame_cap: 1024, - capability: ZAKURA_CAP_LEGACY_GOSSIP, - mode: StreamMode::RequestResponse, - }, - Stream { - kind: DISCOVERY_STREAM_KIND, - version: ZAKURA_STREAM_VERSION_1, - frame_cap: 1024, - capability: ZAKURA_CAP_DISCOVERY, - mode: StreamMode::Ordered, - }, - Stream { - kind: HEADER_SYNC_STREAM_KIND, - version: ZAKURA_HEADER_SYNC_STREAM_VERSION, - frame_cap: 1024, - capability: ZAKURA_CAP_HEADER_SYNC, - mode: StreamMode::Ordered, - }, - Stream { - kind: ZAKURA_STREAM_BLOCK_SYNC, - version: ZAKURA_STREAM_VERSION_1, - frame_cap: MAX_BS_FRAME_BYTES, - capability: crate::zakura::ZAKURA_CAP_BLOCK_SYNC, - mode: StreamMode::Ordered, - }, - ], - }) as Arc]) - .expect("test registry declares unique stream kinds"); + let streams = vec![ + Stream { + kind: LEGACY_GOSSIP_STREAM_KIND, + version: ZAKURA_STREAM_VERSION_1, + frame_cap: 1024, + capability: ZAKURA_CAP_LEGACY_GOSSIP, + mode: StreamMode::Persistent, + }, + Stream { + kind: LEGACY_REQUEST_STREAM_KIND, + version: ZAKURA_STREAM_VERSION_1, + frame_cap: 1024, + capability: ZAKURA_CAP_LEGACY_GOSSIP, + mode: StreamMode::RequestResponse, + }, + Stream { + kind: DISCOVERY_STREAM_KIND, + version: ZAKURA_STREAM_VERSION_1, + frame_cap: 1024, + capability: ZAKURA_CAP_DISCOVERY, + mode: StreamMode::Persistent, + }, + Stream { + kind: HEADER_SYNC_STREAM_KIND, + version: ZAKURA_HEADER_SYNC_STREAM_VERSION, + frame_cap: 1024, + capability: ZAKURA_CAP_HEADER_SYNC, + mode: StreamMode::Persistent, + }, + Stream { + kind: ZAKURA_STREAM_BLOCK_SYNC, + version: ZAKURA_STREAM_VERSION_1, + frame_cap: MAX_BS_FRAME_BYTES, + capability: crate::zakura::ZAKURA_CAP_BLOCK_SYNC, + mode: StreamMode::Persistent, + }, + ]; + let services = streams + .into_iter() + .map(|stream| -> Arc { + Arc::new(DeclaredStreamService { + streams: vec![stream], + }) + }) + .collect(); + let registry = + ServiceRegistry::new(services).expect("test registry declares unique stream kinds"); for (kind, version) in [ (LEGACY_GOSSIP_STREAM_KIND, ZAKURA_STREAM_VERSION_1), @@ -8823,7 +9474,9 @@ mod tests { "the predecessor header-sync stream version is rejected" ); assert!( - registry.ordered_streams_for_negotiated(1 << 4).is_empty(), + registry + .persistent_streams_for_negotiated(1 << 4) + .is_empty(), "the retired predecessor capability opens no header-sync stream" ); diff --git a/crates/zakura-network/src/zakura/handler/service_session.rs b/crates/zakura-network/src/zakura/handler/service_session.rs new file mode 100644 index 0000000000..282df6dd20 --- /dev/null +++ b/crates/zakura-network/src/zakura/handler/service_session.rs @@ -0,0 +1,353 @@ +//! Set up every required persistent stream before admitting a service session. +//! +//! Each member has independent queues and workers. The session shares a wire +//! identifier, admission reservation, message budget, and cancellation scope. +//! Single-stream sessions retain their existing prelude without a wire identifier. + +use super::*; +use crate::zakura::transport::SessionLayout; + +/// Own both directions of one QUIC stream until setup hands them to a new owner. +/// +/// While the handles are present, dropping this guard resets sending and stops +/// receiving. This also cleans up failed or cancelled setup, including a partly +/// written prelude (the header identifying the stream's role and version). +/// `take` transfers ownership and leaves `None`, disabling this guard's cleanup. +pub(super) struct SetupIo(Option<(SendStream, RecvStream)>); + +impl SetupIo { + pub(super) fn new(send: SendStream, recv: RecvStream) -> Self { + Self(Some((send, recv))) + } + + /// Borrow the handles for setup I/O while retaining responsibility for cleanup. + pub(super) fn streams(&mut self) -> (&mut SendStream, &mut RecvStream) { + let (send, recv) = self + .0 + .as_mut() + .expect("setup owns both stream halves until handoff"); + (send, recv) + } + + /// Hand both handles to the caller, which becomes responsible for their lifetime. + pub(super) fn take(mut self) -> (SendStream, RecvStream) { + self.0 + .take() + .expect("setup owns both stream halves until handoff") + } +} + +impl Drop for SetupIo { + fn drop(&mut self) { + if let Some((send, recv)) = &mut self.0 { + let _ = send.reset(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE)); + let _ = recv.stop(VarInt::from_u32(ZAKURA_CLOSE_RESOURCE)); + } + } +} + +/// One stream's I/O, setup metadata, and resources, before its worker starts. +/// +/// Retaining this value keeps its transport stream permit and any service +/// reservation charged, including while it waits for the remaining members. Dropping +/// it stops the stream through `SetupIo` and releases its resource ownership. +pub(super) struct PreparedStream { + io: SetupIo, + stream: Stream, + prelude: StreamPrelude, + context: StreamWorkerContext, +} + +impl PreparedStream { + /// Attach the session's shared service reservation before starting its workers. + pub(super) fn set_session_resources( + &mut self, + resources: Option>, + ) { + self.context.session_resources = resources; + } + + pub(super) fn new( + send: SendStream, + recv: RecvStream, + stream: Stream, + prelude: StreamPrelude, + context: StreamWorkerContext, + ) -> Self { + Self { + io: SetupIo(Some((send, recv))), + stream, + prelude, + context, + } + } +} + +/// Incomplete setup retains every arrived stream under the first arrival's deadline. +struct PendingSession { + id: u64, + streams: Vec, + deadline: Instant, +} + +/// At most one incomplete offer per service session on this connection. +#[derive(Default)] +pub(super) struct PendingSessions { + sessions: HashMap, + retry_after: HashMap, +} + +impl PendingSessions { + /// Charge the service once, then share its reservation across all members. + pub(super) fn reserve_or_share( + &self, + layout: &SessionLayout, + registry: &ServiceRegistry, + direction: ServicePeerDirection, + ) -> Result>, crate::zakura::SessionFull> { + let kind = layout.primary().kind; + if self + .retry_after + .get(&kind) + .is_some_and(|retry| *retry > Instant::now()) + { + return Err(crate::zakura::SessionFull); + } + match self.sessions.get(&kind) { + Some(pending) => Ok(pending.streams[0].context.session_resources.clone()), + None => registry + .service_for_kind(kind) + .expect("a selected session has an owning service") + .reserve_session(direction), + } + } + + pub(super) fn deadline(&self) -> Option { + self.sessions.values().map(|session| session.deadline).min() + } + + /// Expiry releases every arrived member and briefly defers another offer. + pub(super) fn expire(&mut self, now: Instant) { + self.retry_after.retain(|_, retry| *retry > now); + for (kind, session) in self + .sessions + .extract_if(|_, session| session.deadline <= now) + { + self.retry_after.insert( + kind, + now + session.streams[0].context.limits.prelude_timeout, + ); + } + } + + /// Admit the complete layout in kind order, regardless of arrival order. + pub(super) fn insert( + &mut self, + layout: &SessionLayout, + id: u64, + incoming: PreparedStream, + ) -> Result>, ZakuraHandlerError> { + let kind = layout.primary().kind; + // Remove first so every invalid continuation releases the existing offer. + let pending = self.sessions.remove(&kind); + if (layout.is_multi_stream() && id == 0) || !layout.streams.contains(&incoming.stream) { + return Err(ZakuraHandlerError::InvalidServiceSession); + } + let mut pending = pending.unwrap_or_else(|| PendingSession { + id, + streams: Vec::with_capacity(layout.streams.len()), + deadline: Instant::now() + incoming.context.limits.prelude_timeout, + }); + if pending.id != id { + // A retry can arrive before the abandoned offer's reset. Reject both + // offers and back off without closing unrelated service sessions. + self.retry_after.insert( + kind, + Instant::now() + incoming.context.limits.prelude_timeout, + ); + return Ok(None); + } + if Instant::now() >= pending.deadline + || pending.streams.iter().any(|s| { + s.stream.kind == incoming.stream.kind || !layout.streams.contains(&s.stream) + }) + { + return Err(ZakuraHandlerError::InvalidServiceSession); + } + pending.streams.push(incoming); + if pending.streams.len() == layout.streams.len() { + pending.streams.sort_unstable_by_key(|s| s.stream.kind); + Ok(Some(pending.streams)) + } else { + self.sessions.insert(kind, pending); + Ok(None) + } + } +} + +/// Run all members independently and report exit only after every worker finishes. +pub(super) fn spawn_service_session( + workers: &mut JoinSet<()>, + streams: Vec, + queue_depth: usize, + opened_locally: bool, + exits: mpsc::UnboundedSender, +) -> AdmittedSession { + let primary = streams + .first() + .expect("a validated session has at least one stream"); + if let Some(resources) = &primary.context.session_resources { + resources.admitted(); + } + let cancel = primary.context.connection_token.child_token(); + let failure_cause = OrderedStreamFailureCause::default(); + let mut admitted = AdmittedSession { + kind: primary.stream.kind, + session_id: primary.context.stream_id, + cancel_token: cancel.clone(), + streams: Vec::with_capacity(streams.len()), + }; + let exit = SessionExit { + stream: primary.stream, + session_id: admitted.session_id, + opened_locally, + }; + let mut running = futures::stream::FuturesUnordered::new(); + let application_drain = Arc::new(tokio::sync::Barrier::new(streams.len())); + for mut prepared in streams { + prepared.context.stream_token = cancel.clone(); + let (inbound_depth, outbound_depth) = + bounded_stream_queue_depths(queue_depth, prepared.context.queue_depths); + let (inbound_tx, inbound_rx) = mpsc::channel(inbound_depth); + let (sender, outbound_rx) = worker_framed_channel(outbound_depth); + admitted.streams.push(ServiceStreamRole { + kind: prepared.stream.kind, + version: prepared.stream.version, + recv: FramedRecv::new(inbound_rx).with_failure_cause(failure_cause.clone()), + send: sender.with_session_resources(prepared.context.session_resources.clone()), + }); + let failure_cause = failure_cause.clone(); + let application_drain = application_drain.clone(); + running.push(async move { + let (send, recv) = prepared.io.take(); + persistent_stream_worker_with_policy( + send, + recv, + prepared.prelude, + prepared.context, + inbound_tx, + outbound_rx, + inbound_depth, + Some(failure_cause), + Some(application_drain), + ) + .await; + }); + } + workers.spawn(async move { + let _cancel_on_exit = cancel.clone().drop_guard(); + while running.next().await.is_some() { + cancel.cancel(); + } + let _ = exits.send(exit); + }); + admitted +} + +impl ZakuraProtocolHandler { + /// Complete local setup of one outbound stream without starting its worker. + /// + /// Reserve a transport stream slot, open a bidirectional stream, and write its + /// prelude under bounded waits. For a multi-stream session, the caller supplies the same + /// nonzero `session_id` for every member and attaches their shared service resources + /// afterward. The returned value retains the stream slot until handoff or drop. + #[allow(clippy::too_many_arguments)] + pub(super) async fn prepare_ordered_stream( + &self, + connection: &Connection, + stream: Stream, + session_id: Option, + stream_sem: &Arc, + message_buckets: &mut MessageRateBuckets, + limits: ZakuraConnectionLimits, + connection_token: CancellationToken, + close_cause: CloseCause, + freshness_tx: watch::Sender, + conn: ZakuraConnTrace, + peer_id: ZakuraPeerId, + ) -> Result { + let stream_id = self.next_stream_id.fetch_add(1, Ordering::Relaxed); + let permit = stream_sem + .clone() + .try_acquire_owned() + .map_err(|_| ZakuraHandlerError::ResourceLimit("ordered stream permit"))?; + let io = timeout(OUTBOUND_STREAM_WRITE_TIMEOUT, connection.open_bi()) + .await + .map_err(|_| ZakuraHandlerError::Timeout("open ordered service stream"))??; + let mut io = SetupIo(Some(io)); + let prelude = StreamPrelude { + magic: STREAM_PRELUDE_MAGIC, + stream_kind: stream.kind, + stream_version: stream.version, + request_id: None, + max_frame_bytes: inbound_frame_cap_for_stream(&limits, stream), + }; + let mut bytes = prelude.encode()?; + // The session ID follows the ordinary prelude. It matches persistent roles; + // it is separate from the prelude's per-request `request_id` field. + if let Some(id) = session_id { + bytes.extend_from_slice(&id.to_le_bytes()); + } + timeout( + OUTBOUND_STREAM_WRITE_TIMEOUT, + io.0.as_mut() + .expect("setup retains its stream halves") + .0 + .write_all(&bytes), + ) + .await + .map_err(|_| ZakuraHandlerError::Timeout("ordered stream prelude write"))??; + // All members spend the same message-rate budget, so splitting a service + // across several streams does not double its allowance. + let bucket_kind = self + .registry + .session_layout(stream) + .map_or(stream.kind, |layout| layout.primary().kind); + let message_bucket = message_bucket_for( + message_buckets, + bucket_kind, + limits.message_rate_per_second, + RealClock, + ); + let context = StreamWorkerContext { + conn: conn.clone(), + peer_id, + stream_id, + _permit: permit, + limits, + inbound_frame_cap: prelude.max_frame_bytes, + message_payload_limits: self.registry.message_payload_limits(stream), + message_types: self.registry.message_types(stream), + queue_depths: self.registry.stream_queue_depths(stream), + write_policy: self.registry.stream_write_policy(stream), + session_resources: None, + outbound_frame_cap: application_frame_cap(&limits, stream), + message_bucket, + stream_token: connection_token.child_token(), + connection_token, + close_cause, + freshness_tx, + }; + metrics::counter!("zakura.p2p.stream.accepted", "stream_kind" => stream_kind_label(stream.kind)).increment(1); + conn.trace_stream("accepted", stream_id, Some(stream_kind_label(stream.kind))); + Ok(PreparedStream { + io, + stream, + prelude, + context, + }) + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/zakura-network/src/zakura/handler/service_session/tests.rs b/crates/zakura-network/src/zakura/handler/service_session/tests.rs new file mode 100644 index 0000000000..2dd4951345 --- /dev/null +++ b/crates/zakura-network/src/zakura/handler/service_session/tests.rs @@ -0,0 +1,1822 @@ +use super::*; +use crate::zakura::testkit::LocalEndpointFactory; +use tokio_util::task::AbortOnDropHandle; + +const DATA: Stream = Stream { + kind: 64, + version: 1, + frame_cap: 2 * 1024 * 1024, + capability: 1 << 16, + mode: StreamMode::Persistent, +}; +const REQUESTS: Stream = Stream { kind: 65, ..DATA }; +const EVENTS: Stream = Stream { kind: 67, ..DATA }; +const ONE_SHOT: Stream = Stream { + kind: 68, + mode: StreamMode::RequestResponse, + ..DATA +}; +const SIBLING: Stream = Stream { + kind: 66, + capability: 1 << 17, + ..DATA +}; +const ALPN: &[u8] = b"/zakura/test/ordered-pair/1"; +const TEST_DATA_WRITE_TIMEOUT: Duration = Duration::from_secs(32); +const TEST_TIMEOUT: Duration = Duration::from_secs(30); + +#[derive(Debug)] +struct SessionService { + streams: &'static [Stream], + opening: SessionOpening, + fail_first_reservation: std::sync::atomic::AtomicBool, + capacity: Option>, + sessions: mpsc::Sender, +} + +#[derive(Debug)] +struct SessionSlot { + _permit: OwnedSemaphorePermit, +} + +impl crate::zakura::SessionResources for SessionSlot { + fn admitted(&self) {} +} + +#[derive(Debug)] +struct SiblingService { + sessions: mpsc::Sender, + retired: bool, +} + +impl Service for SiblingService { + fn name(&self) -> &'static str { + "test-sibling" + } + fn streams(&self) -> &[Stream] { + &[SIBLING] + } + fn wants_peer(&self, _: &ZakuraPeerId, _: u64, _: ServicePeerDirection) -> bool { + !self.retired + } + fn add_peer(&self, peer: Peer) { + let cancel = peer.service_cancel_token(); + if self.sessions.try_send(peer).is_err() { + cancel.cancel(); + } + } + fn remove_peer(&self, _: &ZakuraPeerId, _: ZakuraConnId) {} +} + +impl Service for SessionService { + fn reserve_session( + &self, + _: ServicePeerDirection, + ) -> Result>, crate::zakura::SessionFull> { + // Inject the outcome of another connection taking the final slot after + // this connection's advisory OpenNow check. A retry can succeed. + if self.fail_first_reservation.swap(false, Ordering::SeqCst) { + Err(crate::zakura::SessionFull) + } else if let Some(capacity) = &self.capacity { + let permit = capacity + .clone() + .try_acquire_owned() + .map_err(|_| crate::zakura::SessionFull)?; + Ok(Some(Arc::new(SessionSlot { _permit: permit }))) + } else { + Ok(None) + } + } + fn name(&self) -> &'static str { + "test-pair" + } + fn streams(&self) -> &[Stream] { + self.streams + } + fn stream_write_policy(&self, stream: Stream) -> StreamWritePolicy { + if stream == REQUESTS { + StreamWritePolicy::UntilCancelled + } else { + StreamWritePolicy::Timeout(TEST_DATA_WRITE_TIMEOUT) + } + } + fn stream_queue_depths(&self, stream: Stream) -> Option<(usize, usize)> { + Some(if stream == EVENTS { (3, 3) } else { (1, 1) }) + } + fn as_request_response(&self) -> Option<&dyn crate::zakura::RequestResponseService> { + Some(self) + } + fn session_policy(&self) -> SessionPolicy { + SessionPolicy { + opening: self.opening, + reopen: true, + } + } + fn add_peer(&self, peer: Peer) { + let cancel = peer.service_cancel_token(); + if self.sessions.try_send(peer).is_err() { + cancel.cancel(); + } + } + fn remove_peer(&self, _: &ZakuraPeerId, _: ZakuraConnId) {} +} + +impl crate::zakura::RequestResponseService for SessionService { + fn request_frame<'a>( + &'a self, + _: ZakuraPeerId, + _: u16, + _: u64, + _: u32, + _: u32, + frame: Frame, + ) -> BoxRunFuture<'a, Result, SinkReject>> { + Box::pin(async move { Ok(vec![frame]) }) + } +} + +struct Session { + id: u64, + conn_id: u64, + data_recv: FramedRecv, + data_send: FramedSend, + request_recv: FramedRecv, + request_send: FramedSend, + cancel: CancellationToken, + connection_cancel: CancellationToken, +} + +impl Session { + async fn receive(receiver: &mut mpsc::Receiver) -> Result { + let mut peer = timeout(TEST_TIMEOUT, receiver.recv()) + .await? + .ok_or("pair admission channel closed")?; + let (data_id, data_version, data_recv, data_send) = peer + .take_versioned_stream_with_session_id(DATA.kind) + .expect("complete data role"); + let (request_id, request_version, request_recv, request_send) = peer + .take_versioned_stream_with_session_id(REQUESTS.kind) + .expect("complete request role"); + assert_eq!(data_id, request_id); + assert_ne!(data_id, 0); + assert_eq!(data_version, DATA.version); + assert_eq!(request_version, REQUESTS.version); + Ok(Self { + id: data_id, + conn_id: peer.conn_id, + data_recv, + data_send, + request_recv, + request_send, + cancel: peer.service_cancel_token(), + connection_cancel: peer.cancel_token(), + }) + } +} + +struct Fixture { + router: Router, + client: Endpoint, + connection: Connection, + serving: AbortOnDropHandle>, + server_sessions: mpsc::Receiver, + client_sessions: mpsc::Receiver, + server_siblings: mpsc::Receiver, + client_siblings: mpsc::Receiver, +} + +impl Fixture { + async fn start() -> Result { + Self::start_with_reservation_race(false).await + } + + async fn start_with_reservation_race(fail_first_reservation: bool) -> Result { + Self::start_with_config(fail_first_reservation, None, false).await + } + + async fn start_with_config( + fail_first_reservation: bool, + max_open_streams: Option, + retired_sibling: bool, + ) -> Result { + Self::start_with_streams( + fail_first_reservation, + max_open_streams, + retired_sibling, + &[DATA, REQUESTS], + ) + .await + } + + async fn start_with_streams( + fail_first_reservation: bool, + max_open_streams: Option, + retired_sibling: bool, + streams: &'static [Stream], + ) -> Result { + let mut local = ZakuraLocalLimits::from_config(&Config::default()); + if let Some(max_open_streams) = max_open_streams { + local.max_open_streams = max_open_streams; + } + let server = LocalEndpointFactory::with_transport_config(local.transport_config()) + .endpoint(93101) + .await?; + let client = LocalEndpointFactory::with_transport_config(local.transport_config()) + .endpoint(93102) + .await?; + let (server_tx, server_sessions) = mpsc::channel(2); + let (client_tx, client_sessions) = mpsc::channel(2); + let (server_sibling_tx, server_siblings) = mpsc::channel(1); + let (client_sibling_tx, client_siblings) = mpsc::channel(1); + let handler = |sessions, siblings, endpoint: Endpoint, fail_reservation| { + ZakuraProtocolHandler::new_with_registry( + ZakuraSupervisorHandle::new(16), + Network::Mainnet, + ZakuraHandshakeConfig::for_network(&Network::Mainnet), + local.clone(), + Arc::new( + ServiceRegistry::new(vec![ + Arc::new(SessionService { + streams, + capacity: None, + opening: SessionOpening::EitherSide, + sessions, + fail_first_reservation: std::sync::atomic::AtomicBool::new( + fail_reservation, + ), + }), + Arc::new(SiblingService { + sessions: siblings, + retired: retired_sibling, + }), + ]) + .unwrap(), + ), + ) + .with_endpoint(endpoint) + }; + let server_opens = i_open_collision_winner(&server.id(), &client.id()); + let server_handler = handler( + server_tx, + server_sibling_tx, + server.clone(), + fail_first_reservation && server_opens, + ); + let client_handler = handler( + client_tx, + client_sibling_tx, + client.clone(), + fail_first_reservation && !server_opens, + ); + let router = Router::builder(server).accept(ALPN, server_handler).spawn(); + let address = LocalEndpointFactory::node_addr(router.endpoint()).await; + let (connection, serving) = super::super::tests::connection::connect_and_serve( + &client, + address, + client_handler, + local, + ALPN, + TEST_TIMEOUT, + ) + .await?; + Ok(Self { + router, + client, + connection, + serving, + server_sessions, + client_sessions, + server_siblings, + client_siblings, + }) + } + + async fn sessions(&mut self) -> Result<(Session, Session), BoxError> { + tokio::try_join!( + Session::receive(&mut self.client_sessions), + Session::receive(&mut self.server_sessions) + ) + } + + async fn close(self) -> Result<(), BoxError> { + self.connection.close(0u32.into(), b"test complete"); + timeout(TEST_TIMEOUT, self.serving).await???; + timeout(TEST_TIMEOUT, self.client.close()).await?; + timeout(TEST_TIMEOUT, self.router.shutdown()).await??; + Ok(()) + } +} + +fn frame(kind: u16, byte: u8, bytes: usize) -> Frame { + Frame { + message_type: kind, + flags: 0, + payload: vec![byte; bytes], + } +} + +async fn exchange(client: &mut Session, server: &mut Session) -> Result<(), BoxError> { + let request = frame(1, 17, 9); + timeout(TEST_TIMEOUT, client.request_send.send(request.clone())).await??; + assert_eq!( + timeout(TEST_TIMEOUT, server.request_recv.recv()).await?, + Some(request) + ); + let response = frame(2, 43, 1024 * 1024); + timeout(TEST_TIMEOUT, server.data_send.send(response.clone())).await??; + assert_eq!( + timeout(TEST_TIMEOUT, client.data_recv.recv()).await?, + Some(response) + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn paired_data_timeout_preserves_sibling_and_reopens_pair() -> Result<(), BoxError> { + let _guard = zakura_test::init(); + let mut fixture = Fixture::start().await?; + let (client, server) = fixture.sessions().await?; + let mut client_sibling = timeout(TEST_TIMEOUT, fixture.client_siblings.recv()) + .await + .expect("the client admits the sibling service") + .ok_or("missing client sibling")?; + let mut server_sibling = timeout(TEST_TIMEOUT, fixture.server_siblings.recv()) + .await + .expect("the server admits the sibling service") + .ok_or("missing server sibling")?; + let (mut sibling_recv, _send) = client_sibling.take_stream(SIBLING.kind).unwrap(); + let (_recv, sibling_send) = server_sibling.take_stream(SIBLING.kind).unwrap(); + + // Keep the peer's data consumer paused past the actual production deadline. + // More than both transport windows ensures a data write must wait. + let started = Instant::now(); + let sender = client.data_send.clone(); + let writes = AbortOnDropHandle::new(tokio::spawn(async move { + for _ in 0..80 { + sender.send(frame(2, 43, 1024 * 1024)).await?; + } + Ok::<_, BoxError>(()) + })); + timeout(TEST_DATA_WRITE_TIMEOUT + Duration::from_secs(10), async { + loop { + let ping = frame(1, 17, 64); + sibling_send.send(ping.clone()).await?; + assert_eq!(sibling_recv.recv().await, Some(ping)); + tokio::select! { + () = client.cancel.cancelled() => break, + () = tokio::time::sleep(Duration::from_millis(100)) => {}, + } + } + Ok::<_, BoxError>(()) + }) + .await + .expect("the paired data writer retires its session at the write deadline")?; + assert!(started.elapsed() >= TEST_DATA_WRITE_TIMEOUT); + assert_eq!( + client.data_recv.failure(), + Some(OrderedStreamFailure::WriteTimeout) + ); + assert_eq!( + client.request_recv.failure(), + Some(OrderedStreamFailure::WriteTimeout) + ); + timeout(TEST_TIMEOUT, server.cancel.cancelled()).await?; + assert!(!client.connection_cancel.is_cancelled()); + assert!(!server.connection_cancel.is_cancelled()); + assert!(!client_sibling.service_cancel_token().is_cancelled()); + assert!(!server_sibling.service_cancel_token().is_cancelled()); + assert!(fixture.connection.close_reason().is_none()); + assert!(timeout(TEST_TIMEOUT, writes).await??.is_err()); + + let (mut replacement_client, mut replacement_server) = fixture.sessions().await?; + assert_eq!(replacement_client.conn_id, client.conn_id); + assert_ne!(replacement_client.id, client.id); + exchange(&mut replacement_client, &mut replacement_server).await?; + let ping = frame(1, 19, 64); + timeout(TEST_TIMEOUT, sibling_send.send(ping.clone())).await??; + assert_eq!( + timeout(TEST_TIMEOUT, sibling_recv.recv()).await?, + Some(ping) + ); + fixture.close().await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn retired_sibling_does_not_prevent_full_capacity_pair() -> Result<(), BoxError> { + let _guard = zakura_test::init(); + // The pair needs both slots; the third negotiated role has no local demand. + let mut fixture = Fixture::start_with_config(false, Some(2), true).await?; + let (mut client, mut server) = fixture.sessions().await?; + exchange(&mut client, &mut server).await?; + assert!(matches!( + fixture.client_siblings.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + assert!(matches!( + fixture.server_siblings.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + + client.cancel.cancel(); + timeout(TEST_TIMEOUT, server.cancel.cancelled()).await?; + let (mut replacement_client, mut replacement_server) = fixture.sessions().await?; + assert_eq!(replacement_client.conn_id, client.conn_id); + assert_ne!(replacement_client.id, client.id); + exchange(&mut replacement_client, &mut replacement_server).await?; + assert!(fixture.connection.close_reason().is_none()); + fixture.close().await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn negotiated_pair_reopens_as_one_session_on_the_same_connection() -> Result<(), BoxError> { + let _guard = zakura_test::init(); + let mut fixture = Fixture::start().await?; + let (mut client, mut server) = fixture.sessions().await?; + let conn_id = client.conn_id; + for _ in 0..3 { + exchange(&mut client, &mut server).await?; + let old_id = client.id; + client.cancel.cancel(); + timeout(TEST_TIMEOUT, server.cancel.cancelled()).await?; + assert!(!client.connection_cancel.is_cancelled()); + assert!(!server.connection_cancel.is_cancelled()); + (client, server) = fixture.sessions().await?; + assert_ne!(client.id, old_id); + assert_eq!(client.conn_id, conn_id); + } + exchange(&mut client, &mut server).await?; + fixture.close().await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn request_backpressure_survives_write_timeout_and_pair_cancellation() -> Result<(), BoxError> +{ + let _guard = zakura_test::init(); + let mut fixture = Fixture::start().await?; + let (mut client, server) = fixture.sessions().await?; + let sender = client.request_send.clone(); + // Exceed both receive and send windows while the request consumer is paused. + // This custom service uses larger frames to reach the same transport state + // without encoding a million tiny GetBlocks requests. + let mut writes = AbortOnDropHandle::new(tokio::spawn(async move { + for _ in 0..80 { + sender.send(frame(1, 17, 1024 * 1024)).await?; + } + Ok::<_, mpsc::error::SendError>(()) + })); + let started = Instant::now(); + while started.elapsed() < OUTBOUND_STREAM_WRITE_TIMEOUT + Duration::from_secs(1) { + let response = frame(2, 43, 64); + timeout(TEST_TIMEOUT, server.data_send.send(response.clone())).await??; + assert_eq!( + timeout(TEST_TIMEOUT, client.data_recv.recv()).await?, + Some(response) + ); + tokio::time::sleep(Duration::from_millis(100)).await; + } + assert!( + !writes.is_finished(), + "request writes are still held by flow control" + ); + assert!( + !client.cancel.is_cancelled(), + "useful responses keep the pair valid" + ); + client.cancel.cancel(); + assert!(timeout(TEST_TIMEOUT, &mut writes).await??.is_err()); + timeout(TEST_TIMEOUT, server.cancel.cancelled()).await?; + assert!(!client.connection_cancel.is_cancelled()); + let (mut replacement_client, mut replacement_server) = fixture.sessions().await?; + exchange(&mut replacement_client, &mut replacement_server).await?; + fixture.close().await +} + +/// Capture only the connection so tests can send exact setup bytes. +#[derive(Debug)] +struct RawConnection(mpsc::Sender); + +impl ProtocolHandler for RawConnection { + async fn accept(&self, connection: Connection) -> Result<(), AcceptError> { + self.0 + .send(connection.clone()) + .await + .map_err(AcceptError::from_err)?; + connection.closed().await; + Ok(()) + } +} + +async fn raw_connection() -> Result<(Router, Endpoint, Connection, Connection), BoxError> { + let local = ZakuraLocalLimits::from_config(&Config::default()); + let transport = || { + local + .transport_config_builder() + .stream_receive_window(64_000u32.into()) + .receive_window(128_000u32.into()) + .send_window(64_000) + .build() + }; + let server = LocalEndpointFactory::with_transport_config(transport()) + .endpoint(93201) + .await?; + let client = LocalEndpointFactory::with_transport_config(transport()) + .endpoint(93202) + .await?; + let (accepted, mut connections) = mpsc::channel(1); + let router = Router::builder(server) + .accept(ALPN, RawConnection(accepted)) + .spawn(); + let connection = timeout( + TEST_TIMEOUT, + client.connect( + LocalEndpointFactory::node_addr(router.endpoint()).await, + ALPN, + ), + ) + .await??; + let remote = timeout(TEST_TIMEOUT, connections.recv()) + .await? + .ok_or("missing connection")?; + Ok((router, client, connection, remote)) +} + +/// Run the real post-handshake connection loop with a raw peer controlling setup bytes. +struct RawFixture { + router: Router, + client: Endpoint, + connection: Connection, + serving: AbortOnDropHandle>, + shutdown: CancellationToken, + service: Arc, + sessions: mpsc::Receiver, + siblings: mpsc::Receiver, +} + +#[tokio::test] +async fn three_stream_session_reopens_at_the_exact_stream_limit() -> Result<(), BoxError> { + let mut fixture = + Fixture::start_with_streams(false, Some(3), true, &[DATA, REQUESTS, EVENTS]).await?; + let mut previous_id = None; + for _ in 0..2 { + let mut client = timeout(TEST_TIMEOUT, fixture.client_sessions.recv()) + .await? + .ok_or("missing client session")?; + let mut server = timeout(TEST_TIMEOUT, fixture.server_sessions.recv()) + .await? + .ok_or("missing server session")?; + let mut client_streams = Vec::new(); + let mut server_streams = Vec::new(); + for stream in [DATA, REQUESTS, EVENTS] { + client_streams.push(client.take_stream_with_session_id(stream.kind).unwrap()); + server_streams.push(server.take_stream_with_session_id(stream.kind).unwrap()); + } + let client_id = client_streams[0].0; + let server_id = server_streams[0].0; + assert_ne!(previous_id, Some(client_id)); + previous_id = Some(client_id); + for ((id, _, send), (remote_id, recv, _)) in + client_streams.iter().zip(server_streams.iter_mut()) + { + assert_eq!(*id, client_id); + assert_eq!(*remote_id, server_id); + let message = frame(1, 21, 8); + timeout(TEST_TIMEOUT, send.send(message.clone())).await??; + assert_eq!(timeout(TEST_TIMEOUT, recv.recv()).await?, Some(message)); + } + client.service_cancel_token().cancel(); + timeout(TEST_TIMEOUT, server.service_cancel_token().cancelled()).await?; + assert!(!client.cancel_token().is_cancelled()); + } + fixture.close().await +} + +#[tokio::test] +async fn invalid_third_member_releases_the_entire_pending_session() -> Result<(), BoxError> { + for (stream, id) in [(REQUESTS, 9), (EVENTS, 10), (EVENTS, 0)] { + let mut fixture = + RawFixture::start_with_streams(1, Duration::from_secs(3), &[DATA, REQUESTS, EVENTS]) + .await?; + let _data = fixture.offer(DATA, Some(9)).await?; + let _requests = fixture.offer(REQUESTS, Some(9)).await?; + fixture.wait_for_slots(0, TEST_TIMEOUT).await?; + let _invalid = fixture.offer(stream, Some(id)).await?; + fixture.wait_for_slots(1, Duration::from_secs(1)).await?; + if id == 10 { + assert!(fixture.connection.close_reason().is_none()); + assert!(!fixture.serving.is_finished()); + } else { + timeout(Duration::from_secs(1), async { + while !fixture.serving.is_finished() { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await?; + } + assert!(fixture.sessions.try_recv().is_err()); + fixture.close().await?; + } + Ok(()) +} + +impl RawFixture { + async fn start(slots: usize, setup_timeout: Duration) -> Result { + Self::start_with_streams(slots, setup_timeout, &[DATA, REQUESTS]).await + } + + async fn start_with_streams( + slots: usize, + setup_timeout: Duration, + streams: &'static [Stream], + ) -> Result { + let (router, client, connection, remote) = raw_connection().await?; + let local = ZakuraLocalLimits::from_config(&Config::default()); + let (sessions_tx, sessions) = mpsc::channel(2); + let (siblings_tx, siblings) = mpsc::channel(1); + let service = Arc::new(SessionService { + streams, + opening: SessionOpening::InitiatorOnly, + fail_first_reservation: std::sync::atomic::AtomicBool::new(false), + capacity: Some(Arc::new(Semaphore::new(slots))), + sessions: sessions_tx, + }); + let handler = ZakuraProtocolHandler::new_with_registry( + ZakuraSupervisorHandle::new(16), + Network::Mainnet, + ZakuraHandshakeConfig::for_network(&Network::Mainnet), + local.clone(), + Arc::new(ServiceRegistry::new(vec![ + service.clone(), + Arc::new(SiblingService { + sessions: siblings_tx, + retired: false, + }), + ])?), + ); + let shutdown = handler.shutdown.clone(); + let mut limits = local.clamp(&local.initial_limits()); + limits.prelude_timeout = setup_timeout; + let peer_id = ZakuraPeerId::new(client.id().as_bytes().to_vec())?; + let transcript_hash = native_connection_transcript_hash( + ServicePeerDirection::Inbound, + &router.endpoint().id(), + &client.id(), + ); + let remote_ip = confirmed_remote_ip(&remote); + let serving = AbortOnDropHandle::new(tokio::spawn(async move { + handler + .register_and_serve( + remote, + peer_id, + remote_ip, + ConnectionServeContext { + limits, + accepted_capabilities: DATA.capability | SIBLING.capability, + role: "responder", + direction: ServicePeerDirection::Inbound, + transcript_hash, + i_open_collision_winner: false, + conn: ZakuraConnTrace::without_peer(1), + }, + ) + .await + })); + Ok(Self { + router, + client, + connection, + serving, + shutdown, + service, + sessions, + siblings, + }) + } + + fn capacity(&self) -> &Arc { + self.service.capacity.as_ref().unwrap() + } + + async fn wait_for_slots(&self, count: usize, deadline: Duration) -> Result<(), BoxError> { + timeout(deadline, async { + while self.capacity().available_permits() != count { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await?; + Ok(()) + } + + async fn offer( + &self, + stream: Stream, + pair_id: Option, + ) -> Result<(SendStream, RecvStream), BoxError> { + let (mut send, recv) = timeout(TEST_TIMEOUT, self.connection.open_bi()).await??; + let mut bytes = StreamPrelude { + magic: STREAM_PRELUDE_MAGIC, + stream_kind: stream.kind, + stream_version: stream.version, + request_id: None, + max_frame_bytes: stream.frame_cap, + } + .encode()?; + if let Some(id) = pair_id { + bytes.extend_from_slice(&id.to_le_bytes()); + } + timeout(TEST_TIMEOUT, send.write_all(&bytes)).await??; + Ok((send, recv)) + } + + async fn close(mut self) -> Result<(), BoxError> { + self.shutdown.cancel(); + timeout(Duration::from_millis(500), &mut self.serving).await???; + self.connection.close(0u32.into(), b"test complete"); + timeout(TEST_TIMEOUT, self.client.close()).await?; + timeout(TEST_TIMEOUT, self.router.shutdown()).await??; + Ok(()) + } +} + +#[tokio::test] +async fn three_stream_session_waits_for_every_member_and_leaves_requests_independent( +) -> Result<(), BoxError> { + let mut fixture = RawFixture::start_with_streams( + 1, + Duration::from_secs(3), + &[EVENTS, ONE_SHOT, REQUESTS, DATA], + ) + .await?; + let (mut sibling_send, _sibling_recv) = fixture.offer(SIBLING, None).await?; + let mut sibling = timeout(TEST_TIMEOUT, fixture.siblings.recv()) + .await? + .ok_or("missing sibling")?; + let (mut sibling_recv, _sibling_sender) = sibling.take_stream(SIBLING.kind).unwrap(); + + let mut events = fixture.offer(EVENTS, Some(42)).await?; + let mut requests = fixture.offer(REQUESTS, Some(42)).await?; + let event = frame(3, 11, 8); + let request = frame(1, 12, 8); + events.0.write_all(&event.encode(EVENTS.frame_cap)?).await?; + requests + .0 + .write_all(&request.encode(REQUESTS.frame_cap)?) + .await?; + fixture.wait_for_slots(0, TEST_TIMEOUT).await?; + assert!( + timeout(Duration::from_millis(100), fixture.sessions.recv()) + .await + .is_err(), + "two of three required streams cannot start the service" + ); + + let _data = fixture.offer(DATA, Some(42)).await?; + let mut peer = timeout(TEST_TIMEOUT, fixture.sessions.recv()) + .await? + .ok_or("missing session")?; + let cancel = peer.service_cancel_token(); + let (data_id, mut data_recv, _data_send) = peer.take_stream_with_session_id(DATA.kind).unwrap(); + let (request_id, mut request_recv, _request_send) = + peer.take_stream_with_session_id(REQUESTS.kind).unwrap(); + let (event_id, mut event_recv, event_send) = + peer.take_stream_with_session_id(EVENTS.kind).unwrap(); + assert_eq!((data_id, data_id), (request_id, event_id)); + assert_eq!( + event_send.max_capacity(), + 3, + "the service controls each stream's queue" + ); + assert_eq!(timeout(TEST_TIMEOUT, event_recv.recv()).await?, Some(event)); + assert_eq!( + timeout(TEST_TIMEOUT, request_recv.recv()).await?, + Some(request) + ); + assert!(peer.take_stream(ONE_SHOT.kind).is_none()); + + // A per-request stream carries no session identifier and completes independently. + let (mut send, mut recv) = fixture.connection.open_bi().await?; + let mut bytes = StreamPrelude { + magic: STREAM_PRELUDE_MAGIC, + stream_kind: ONE_SHOT.kind, + stream_version: ONE_SHOT.version, + request_id: Some(99), + max_frame_bytes: ONE_SHOT.frame_cap, + } + .encode()?; + let echo = frame(4, 13, 8).encode(ONE_SHOT.frame_cap)?; + bytes.extend_from_slice(&echo); + timeout(TEST_TIMEOUT, send.write_all(&bytes)).await??; + send.finish()?; + assert_eq!(timeout(TEST_TIMEOUT, recv.read_to_end(1024)).await??, echo); + assert!(!cancel.is_cancelled()); + + events.0.reset(0u32.into())?; + events.1.stop(0u32.into())?; + timeout(TEST_TIMEOUT, cancel.cancelled()).await?; + assert!(timeout(TEST_TIMEOUT, data_recv.recv()).await?.is_none()); + assert!(timeout(TEST_TIMEOUT, request_recv.recv()).await?.is_none()); + assert!(timeout(TEST_TIMEOUT, event_recv.recv()).await?.is_none()); + assert!(!peer.cancel_token().is_cancelled()); + let ping = frame(2, 14, 8); + sibling_send + .write_all(&ping.encode(SIBLING.frame_cap)?) + .await?; + assert_eq!( + timeout(TEST_TIMEOUT, sibling_recv.recv()).await?, + Some(ping) + ); + fixture.close().await +} + +#[tokio::test] +async fn incomplete_three_stream_session_releases_all_arrivals_on_expiry() -> Result<(), BoxError> { + let mut fixture = + RawFixture::start_with_streams(1, Duration::from_millis(300), &[DATA, REQUESTS, EVENTS]) + .await?; + let (_request_send, mut request_recv) = fixture.offer(REQUESTS, Some(1)).await?; + let (_event_send, mut event_recv) = fixture.offer(EVENTS, Some(1)).await?; + fixture.wait_for_slots(0, TEST_TIMEOUT).await?; + fixture.wait_for_slots(1, TEST_TIMEOUT).await?; + assert!(fixture.sessions.try_recv().is_err()); + assert!(timeout(TEST_TIMEOUT, request_recv.read_exact(&mut [0; 1])) + .await? + .is_err()); + assert!(timeout(TEST_TIMEOUT, event_recv.read_exact(&mut [0; 1])) + .await? + .is_err()); + assert!(fixture.connection.close_reason().is_none()); + fixture.close().await +} + +#[tokio::test] +async fn withheld_pair_id_does_not_reserve_service_capacity() -> Result<(), BoxError> { + let fixture = RawFixture::start(1, Duration::from_secs(2)).await?; + let (mut send, _recv) = fixture.offer(DATA, None).await?; + tokio::time::sleep(Duration::from_millis(200)).await; + assert_eq!(fixture.capacity().available_permits(), 1); + let outbound = fixture + .service + .reserve_session(ServicePeerDirection::Outbound)?; + drop(outbound); + send.write_all(&1u64.to_le_bytes()).await?; + fixture.wait_for_slots(0, TEST_TIMEOUT).await?; + fixture.close().await +} + +#[tokio::test] +async fn abandoned_half_pair_retry_preserves_the_connection() -> Result<(), BoxError> { + // Also cover a reset arriving after the retry, since QUIC streams can reorder. + for reset_before_retry in [true, false] { + let setup_timeout = Duration::from_secs(1); + let mut fixture = RawFixture::start(1, setup_timeout).await?; + let (mut sibling_send, _sibling_recv) = fixture.offer(SIBLING, None).await?; + let mut sibling = timeout(TEST_TIMEOUT, fixture.siblings.recv()) + .await? + .ok_or("missing sibling")?; + let (mut sibling_recv, _sibling_send) = sibling.take_stream(SIBLING.kind).unwrap(); + let (send, recv) = fixture.offer(DATA, Some(1)).await?; + let mut abandoned = Some(SetupIo::new(send, recv)); + fixture.wait_for_slots(0, TEST_TIMEOUT).await?; + if reset_before_retry { + drop(abandoned.take()); + } + + let (_retry_send, mut retry_recv) = fixture.offer(DATA, Some(2)).await?; + assert!(timeout(TEST_TIMEOUT, retry_recv.read_exact(&mut [0; 1])) + .await? + .is_err()); + fixture.wait_for_slots(1, TEST_TIMEOUT).await?; + drop(abandoned); + assert!( + !sibling.cancel_token().is_cancelled(), + "retrying an abandoned half pair must preserve sibling services" + ); + assert!(fixture.sessions.try_recv().is_err()); + let ping = frame(2, 14, 8); + sibling_send + .write_all(&ping.encode(SIBLING.frame_cap)?) + .await?; + assert_eq!( + timeout(TEST_TIMEOUT, sibling_recv.recv()).await?, + Some(ping) + ); + + let (_early_send, mut early_recv) = fixture.offer(DATA, Some(3)).await?; + assert!(timeout(TEST_TIMEOUT, early_recv.read_exact(&mut [0; 1])) + .await? + .is_err()); + assert_eq!(fixture.capacity().available_permits(), 1); + let outbound = fixture + .service + .reserve_session(ServicePeerDirection::Outbound)?; + drop(outbound); + + tokio::time::sleep(setup_timeout).await; + let _data = fixture.offer(DATA, Some(4)).await?; + let (mut requests_send, _requests_recv) = fixture.offer(REQUESTS, Some(4)).await?; + let mut replacement = Session::receive(&mut fixture.sessions).await?; + let request = frame(1, 7, 8); + requests_send + .write_all(&request.encode(REQUESTS.frame_cap)?) + .await?; + assert_eq!( + timeout(TEST_TIMEOUT, replacement.request_recv.recv()).await?, + Some(request) + ); + assert_eq!(fixture.capacity().available_permits(), 0); + assert!(!sibling.cancel_token().is_cancelled()); + fixture.close().await?; + } + Ok(()) +} + +#[tokio::test] +async fn slow_pair_setup_does_not_delay_expiry_or_shutdown() -> Result<(), BoxError> { + let fixture = RawFixture::start(1, Duration::from_secs(2)).await?; + let _first = fixture.offer(DATA, Some(1)).await?; + fixture.wait_for_slots(0, TEST_TIMEOUT).await?; + tokio::time::sleep(Duration::from_secs(1)).await; + let _slow = fixture.offer(REQUESTS, None).await?; + fixture + .wait_for_slots(1, Duration::from_millis(1500)) + .await + .expect("the earlier pair expires while the later pair ID is still missing"); + fixture.close().await +} + +#[tokio::test] +async fn expired_pair_cannot_reclaim_capacity_before_an_outgoing_session() -> Result<(), BoxError> { + let fixture = RawFixture::start(1, Duration::from_secs(1)).await?; + let _first = fixture.offer(DATA, Some(1)).await?; + fixture.wait_for_slots(0, TEST_TIMEOUT).await?; + fixture.wait_for_slots(1, TEST_TIMEOUT).await?; + for id in 2..5 { + let (_send, mut recv) = fixture.offer(DATA, Some(id)).await?; + assert!( + timeout(Duration::from_millis(500), recv.read_exact(&mut [0; 1])) + .await? + .is_err() + ); + assert_eq!(fixture.capacity().available_permits(), 1); + } + let outbound = fixture + .service + .reserve_session(ServicePeerDirection::Outbound)?; + assert_eq!(fixture.capacity().available_permits(), 0); + drop(outbound); + fixture.close().await +} + +#[tokio::test] +async fn duplicate_pair_offers_leave_capacity_for_other_peers() -> Result<(), BoxError> { + let mut fixture = RawFixture::start(2, TEST_TIMEOUT).await?; + let (mut sibling_send, _sibling_recv) = fixture.offer(SIBLING, None).await?; + let mut sibling = timeout(TEST_TIMEOUT, fixture.siblings.recv()) + .await? + .ok_or("missing sibling")?; + let (mut sibling_recv, _sibling_send) = sibling.take_stream(SIBLING.kind).unwrap(); + let (_data_send, mut data_recv) = fixture.offer(DATA, Some(1)).await?; + let (mut request_send, _request_recv) = fixture.offer(REQUESTS, Some(1)).await?; + let mut active = Session::receive(&mut fixture.sessions).await?; + assert_eq!(fixture.capacity().available_permits(), 1); + + // Either role can arrive first, but neither may reserve a second session. + for stream in [DATA, REQUESTS] { + let (_duplicate_send, mut duplicate_recv) = fixture.offer(stream, Some(2)).await?; + assert!(timeout( + Duration::from_secs(2), + duplicate_recv.read_exact(&mut [0; 1]) + ) + .await + .expect("duplicate offers must be rejected before the setup deadline") + .is_err()); + assert_eq!(fixture.capacity().available_permits(), 1); + let other_peer = fixture + .service + .reserve_session(ServicePeerDirection::Inbound)?; + assert_eq!(fixture.capacity().available_permits(), 0); + drop(other_peer); + + let request = frame(1, 7, 8); + request_send + .write_all(&request.encode(REQUESTS.frame_cap)?) + .await?; + assert_eq!( + timeout(TEST_TIMEOUT, active.request_recv.recv()).await?, + Some(request) + ); + let response = frame(2, 8, 64); + timeout(TEST_TIMEOUT, active.data_send.send(response.clone())).await??; + assert_eq!( + read_frame( + &mut data_recv, + DATA.frame_cap, + &[], + None, + TEST_TIMEOUT, + Some(TEST_TIMEOUT), + ) + .await?, + response + ); + let ping = frame(3, 9, 8); + sibling_send + .write_all(&ping.encode(SIBLING.frame_cap)?) + .await?; + assert_eq!( + timeout(TEST_TIMEOUT, sibling_recv.recv()).await?, + Some(ping) + ); + assert!(!active.cancel.is_cancelled()); + assert!(!sibling.cancel_token().is_cancelled()); + assert!(fixture.sessions.try_recv().is_err()); + } + fixture.close().await +} + +#[tokio::test] +async fn paired_replacement_during_cleanup_preserves_the_connection() -> Result<(), BoxError> { + let mut fixture = RawFixture::start(2, Duration::from_secs(3)).await?; + let (mut sibling_send, _sibling_recv) = fixture.offer(SIBLING, None).await?; + let mut sibling = timeout(TEST_TIMEOUT, fixture.siblings.recv()) + .await? + .ok_or("missing sibling")?; + let (mut sibling_recv, _sibling_send) = sibling.take_stream(SIBLING.kind).unwrap(); + let mut old_data = fixture.offer(DATA, Some(1)).await?; + let mut old_requests = fixture.offer(REQUESTS, Some(1)).await?; + let old = Session::receive(&mut fixture.sessions).await?; + + let (mut _new_data_send, mut new_data_recv) = fixture.offer(DATA, Some(2)).await?; + assert!(timeout( + Duration::from_secs(1), + new_data_recv.read_exact(&mut [0; 1]) + ) + .await? + .is_err()); + assert_eq!(fixture.capacity().available_permits(), 1); + let (mut new_requests_send, _new_requests_recv) = fixture.offer(REQUESTS, None).await?; + // Keep the new setup read pending while the old pair's workers exit. + tokio::time::sleep(Duration::from_millis(200)).await; + old_data.0.reset(0u32.into())?; + old_data.1.stop(0u32.into())?; + old_requests.0.reset(0u32.into())?; + old_requests.1.stop(0u32.into())?; + timeout(TEST_TIMEOUT, old.cancel.cancelled()).await?; + let old_id = old.id; + drop(old); + fixture.wait_for_slots(2, TEST_TIMEOUT).await?; + new_requests_send.write_all(&2u64.to_le_bytes()).await?; + fixture.wait_for_slots(1, TEST_TIMEOUT).await?; + (_new_data_send, new_data_recv) = fixture.offer(DATA, Some(2)).await?; + let mut replacement = Session::receive(&mut fixture.sessions).await?; + assert_ne!(replacement.id, old_id); + + let request = frame(1, 7, 8); + new_requests_send + .write_all(&request.encode(REQUESTS.frame_cap)?) + .await?; + assert_eq!( + timeout(TEST_TIMEOUT, replacement.request_recv.recv()).await?, + Some(request) + ); + let response = frame(2, 8, 64); + timeout(TEST_TIMEOUT, replacement.data_send.send(response.clone())).await??; + assert_eq!( + read_frame( + &mut new_data_recv, + DATA.frame_cap, + &[], + None, + TEST_TIMEOUT, + Some(TEST_TIMEOUT) + ) + .await?, + response + ); + let ping = frame(3, 9, 8); + sibling_send + .write_all(&ping.encode(SIBLING.frame_cap)?) + .await?; + assert_eq!( + timeout(TEST_TIMEOUT, sibling_recv.recv()).await?, + Some(ping) + ); + assert!(!replacement.connection_cancel.is_cancelled()); + assert!(!sibling.cancel_token().is_cancelled()); + fixture.close().await +} + +fn raw_worker_context(client: &Endpoint, slots: Arc) -> StreamWorkerContext { + let local = ZakuraLocalLimits::from_config(&Config::default()); + let cancel = CancellationToken::new(); + let (freshness_tx, _freshness_rx) = watch::channel(Instant::now()); + StreamWorkerContext { + conn: ZakuraConnTrace::without_peer(1), + peer_id: ZakuraPeerId::new(client.id().as_bytes().to_vec()).unwrap(), + stream_id: 1, + _permit: slots.try_acquire_owned().unwrap(), + limits: local.clamp(&local.initial_limits()), + inbound_frame_cap: DATA.frame_cap, + message_payload_limits: &[], + message_types: None, + queue_depths: None, + write_policy: StreamWritePolicy::UntilCancelled, + session_resources: None, + outbound_frame_cap: DATA.frame_cap, + message_bucket: Arc::new(std::sync::Mutex::new(TokenBucket::new(128))), + connection_token: cancel.clone(), + stream_token: cancel.child_token(), + close_cause: CloseCause::new(), + freshness_tx, + } +} + +#[tokio::test] +async fn paired_request_reader_close_interrupts_a_blocked_write() -> Result<(), BoxError> { + let _guard = zakura_test::init(); + let (router, client, connection, remote) = raw_connection().await?; + for (reset, drop_receiver) in [(false, false), (true, false), (false, true), (true, true)] { + let (mut peer_send, mut peer_recv) = connection.open_bi().await?; + peer_send + .write_all(&frame(1, 0, 0).encode(DATA.frame_cap)?) + .await?; + let (send, recv) = timeout(TEST_TIMEOUT, remote.accept_bi()).await??; + let slots = Arc::new(Semaphore::new(1)); + let context = raw_worker_context(&client, slots.clone()); + let connection_cancel = context.connection_token.clone(); + let pair_cancel = context.stream_token.clone(); + let failure_cause = OrderedStreamFailureCause::default(); + let prelude = StreamPrelude { + magic: STREAM_PRELUDE_MAGIC, + stream_kind: REQUESTS.kind, + stream_version: REQUESTS.version, + request_id: None, + max_frame_bytes: DATA.frame_cap, + }; + let (inbound_tx, mut inbound_rx) = mpsc::channel(1); + let (outbound_tx, outbound_rx) = worker_framed_channel(1); + let mut worker = + AbortOnDropHandle::new(tokio::spawn(persistent_stream_worker_with_policy( + send, + recv, + prelude, + context, + inbound_tx, + outbound_rx, + 1, + Some(failure_cause.clone()), + None, + ))); + assert_eq!( + timeout(TEST_TIMEOUT, inbound_rx.recv()).await?, + Some(frame(1, 0, 0)) + ); + let resources = Arc::new(Semaphore::new(1)); + outbound_tx.try_reserve_guarded().unwrap().send( + frame(1, 17, 1024 * 1024), + crate::zakura::FrameGuard::new(Arc::new(resources.clone().try_acquire_owned()?)), + ); + // Read only the first byte: the application write cannot finish within + // the smaller QUIC windows, independently of scheduling or elapsed time. + timeout(TEST_TIMEOUT, peer_recv.read_exact(&mut [0; 1])).await??; + assert_eq!(resources.available_permits(), 0); + if drop_receiver { + drop(inbound_rx); + peer_send + .write_all(&frame(1, 19, 8).encode(DATA.frame_cap)?) + .await?; + assert!(timeout(Duration::from_millis(100), pair_cancel.cancelled()) + .await + .is_err()); + } + if reset { + peer_send.reset(0u32.into())?; + } else { + peer_send.finish()?; + } + timeout(Duration::from_secs(2), &mut worker) + .await + .expect("closing the reader interrupts a flow-controlled request write")?; + assert_eq!(failure_cause.get(), Some(OrderedStreamFailure::RemoteClose)); + assert!(pair_cancel.is_cancelled()); + assert!(!connection_cancel.is_cancelled()); + assert_eq!(slots.available_permits(), 1); + assert_eq!(resources.available_permits(), 1); + } + connection.close(0u32.into(), b"done"); + timeout(TEST_TIMEOUT, client.close()).await?; + timeout(TEST_TIMEOUT, router.shutdown()).await??; + Ok(()) +} + +#[tokio::test] +async fn request_response_allowlists_reject_headers_before_payloads() -> Result<(), BoxError> { + let _guard = zakura_test::init(); + let (router, client, connection, remote) = raw_connection().await?; + let stream = Stream { + kind: LEGACY_REQUEST_STREAM_KIND, + mode: StreamMode::RequestResponse, + ..DATA + }; + let mut header = Vec::new(); + header.extend_from_slice(&u16::MAX.to_le_bytes()); + header.extend_from_slice(&0u16.to_le_bytes()); + header.extend_from_slice(&1024u32.to_le_bytes()); + let (mut peer_send, _peer_recv) = connection.open_bi().await?; + peer_send.write_all(&header).await?; + let (send, recv) = timeout(TEST_TIMEOUT, remote.accept_bi()).await??; + let mut context = raw_worker_context(&client, Arc::new(Semaphore::new(1))); + context.message_types = Some(&[LEGACY_REQUEST_PING, LEGACY_RESPONSE_PONG]); + let cancel = context.connection_token.clone(); + let limits = context.limits; + let types = context.message_types; + let prelude = StreamPrelude { + magic: STREAM_PRELUDE_MAGIC, + stream_kind: stream.kind, + stream_version: stream.version, + request_id: Some(42), + max_frame_bytes: stream.frame_cap, + }; + timeout( + Duration::from_secs(2), + request_stream_worker( + send, + recv, + prelude, + context, + Arc::new(ServiceRegistry::new(vec![])?), + ), + ) + .await + .expect("the request header is rejected without waiting for its absent payload"); + assert!(cancel.is_cancelled()); + + let response = write_outbound_request_frame( + &connection, + limits, + stream, + &[], + types, + 42, + LEGACY_REQUEST_PING, + 0, + Vec::new(), + ); + let responder = async { + let (mut send, mut recv) = remote.accept_bi().await?; + read_stream_prelude(&mut recv, TEST_TIMEOUT).await?; + let request = read_frame( + &mut recv, + stream.frame_cap, + &[], + types, + TEST_TIMEOUT, + Some(TEST_TIMEOUT), + ) + .await?; + assert_eq!(request.message_type, LEGACY_REQUEST_PING); + send.write_all(&header).await?; + Ok::<_, BoxError>((send, recv)) + }; + let (result, held_stream) = tokio::join!(timeout(Duration::from_secs(2), response), responder); + let _held_stream = held_stream?; + assert!( + matches!(result.expect("the response header is rejected before its absent payload"), + Err(OutboundRequestError::Fatal(error)) + if matches!(error.downcast_ref::(), Some(ZakuraHandlerError::InvalidMessageType(u16::MAX))) + ) + ); + connection.close(0u32.into(), b"done"); + timeout(TEST_TIMEOUT, client.close()).await?; + timeout(TEST_TIMEOUT, router.shutdown()).await??; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn incomplete_pairs_expire_and_mismatched_roles_release_stream_permits( +) -> Result<(), BoxError> { + let _guard = zakura_test::init(); + let (router, client, connection, remote) = raw_connection().await?; + let local = ZakuraLocalLimits::from_config(&Config::default()); + let permits = Arc::new(Semaphore::new(2)); + let (sessions, _sessions_rx) = mpsc::channel(1); + let handler = ZakuraProtocolHandler::new_with_registry( + ZakuraSupervisorHandle::new(16), + Network::Mainnet, + ZakuraHandshakeConfig::for_network(&Network::Mainnet), + local.clone(), + Arc::new(ServiceRegistry::new(vec![Arc::new(SessionService { + streams: &[DATA, REQUESTS], + capacity: None, + opening: SessionOpening::EitherSide, + fail_first_reservation: std::sync::atomic::AtomicBool::new(false), + sessions, + })])?), + ); + let mut limits = local.clamp(&local.initial_limits()); + limits.prelude_timeout = Duration::from_millis(100); + let peer = ZakuraPeerId::new(client.id().as_bytes().to_vec())?; + let (freshness, _freshness_rx) = watch::channel(Instant::now()); + let cancel = CancellationToken::new(); + let mut pending = PendingSessions::default(); + let mut workers = JoinSet::new(); + let mut buckets = MessageRateBuckets::new(); + let mut open_limiter = TokenBucket::new(100); + let (exits, _exit_rx) = mpsc::unbounded_channel(); + let mut admission = StreamAdmission { + is_initiator: false, + direction: ServicePeerDirection::Inbound, + conn: ZakuraConnTrace::without_peer(1), + peer_id: &peer, + stream_sem: &permits, + open_limiter: &mut open_limiter, + message_buckets: &mut buckets, + workers: &mut workers, + limits, + accepted_capabilities: DATA.capability, + connection_token: cancel.clone(), + close_cause: CloseCause::new(), + freshness_tx: freshness, + }; + let mut offers = Vec::new(); + // The request role may arrive first. It must not reach the service alone. + for (kind, id, should_expire) in [ + (REQUESTS.kind, 1u64, true), + (DATA.kind, 2, false), + (REQUESTS.kind, 3, false), + ] { + let (mut send, recv) = connection.open_bi().await?; + let mut bytes = StreamPrelude { + magic: STREAM_PRELUDE_MAGIC, + stream_kind: kind, + stream_version: 1, + request_id: None, + max_frame_bytes: DATA.frame_cap, + } + .encode()?; + bytes.extend_from_slice(&id.to_le_bytes()); + timeout(TEST_TIMEOUT, send.write_all(&bytes)).await??; + offers.push((send, recv)); + let (send, recv) = timeout(TEST_TIMEOUT, remote.accept_bi()).await??; + assert!(handler + .admit_bi_stream(send, recv, &mut admission, 2, exits.clone(), &mut pending) + .await + .is_none()); + if should_expire { + assert_eq!(permits.available_permits(), 1); + let deadline = pending.deadline().expect("incomplete pair owns a deadline"); + tokio::time::sleep_until(deadline).await; + pending.expire(Instant::now()); + assert_eq!(permits.available_permits(), 2); + assert!( + !cancel.is_cancelled(), + "incomplete setup is retired locally" + ); + assert!(pending + .reserve_or_share( + &handler.registry.session_layout(DATA).unwrap(), + &handler.registry, + ServicePeerDirection::Inbound + ) + .is_err()); + tokio::time::sleep(limits.prelude_timeout).await; + } + } + assert!( + !cancel.is_cancelled(), + "different pair identifiers retire only the incomplete session" + ); + assert!(pending + .reserve_or_share( + &handler.registry.session_layout(DATA).unwrap(), + &handler.registry, + ServicePeerDirection::Inbound + ) + .is_err()); + assert_eq!(permits.available_permits(), 2); + assert!(pending.deadline().is_none()); + assert!( + admission.workers.is_empty(), + "incomplete and mismatched pairs never activate workers" + ); + drop(offers); + connection.close(0u32.into(), b"done"); + timeout(TEST_TIMEOUT, client.close()).await?; + timeout(TEST_TIMEOUT, router.shutdown()).await??; + Ok(()) +} + +#[tokio::test] +async fn ineligible_pair_opener_is_rejected_before_service_reservation() -> Result<(), BoxError> { + let _guard = zakura_test::init(); + let (router, client, connection, remote) = raw_connection().await?; + let (sessions, _sessions_rx) = mpsc::channel(1); + let service = Arc::new(SessionService { + streams: &[DATA, REQUESTS], + capacity: None, + opening: SessionOpening::InitiatorOnly, + fail_first_reservation: std::sync::atomic::AtomicBool::new(true), + sessions, + }); + let local = ZakuraLocalLimits::from_config(&Config::default()); + let handler = ZakuraProtocolHandler::new_with_registry( + ZakuraSupervisorHandle::new(16), + Network::Mainnet, + ZakuraHandshakeConfig::for_network(&Network::Mainnet), + local.clone(), + Arc::new(ServiceRegistry::new(vec![service.clone()])?), + ); + let slots = Arc::new(Semaphore::new(2)); + let context = raw_worker_context(&client, Arc::new(Semaphore::new(1))); + let mut workers = JoinSet::new(); + let mut open_limiter = TokenBucket::new(100); + let mut buckets = MessageRateBuckets::new(); + let mut pending = PendingSessions::default(); + let (exits, _exit_rx) = mpsc::unbounded_channel(); + let mut admission = StreamAdmission { + is_initiator: true, + direction: ServicePeerDirection::Outbound, + conn: context.conn, + peer_id: &context.peer_id, + stream_sem: &slots, + open_limiter: &mut open_limiter, + message_buckets: &mut buckets, + workers: &mut workers, + limits: context.limits, + accepted_capabilities: DATA.capability, + connection_token: context.connection_token.clone(), + close_cause: context.close_cause, + freshness_tx: context.freshness_tx, + }; + let (mut send, _recv) = connection.open_bi().await?; + // Withhold the pair ID: even the first role cannot reserve service capacity. + let prelude = StreamPrelude { + magic: STREAM_PRELUDE_MAGIC, + stream_kind: REQUESTS.kind, + stream_version: REQUESTS.version, + request_id: None, + max_frame_bytes: DATA.frame_cap, + }; + send.write_all(&prelude.encode()?).await?; + let (send, recv) = timeout(TEST_TIMEOUT, remote.accept_bi()).await??; + assert!(timeout( + Duration::from_secs(1), + handler.admit_bi_stream(send, recv, &mut admission, 2, exits, &mut pending,) + ) + .await? + .is_none()); + assert!( + service.fail_first_reservation.load(Ordering::SeqCst), + "an ineligible opener must not call service reservation" + ); + assert!(context.connection_token.is_cancelled()); + assert_eq!(slots.available_permits(), 2); + assert!(pending.deadline().is_none()); + assert!(workers.is_empty()); + connection.close(0u32.into(), b"done"); + timeout(TEST_TIMEOUT, client.close()).await?; + timeout(TEST_TIMEOUT, router.shutdown()).await??; + Ok(()) +} + +#[tokio::test] +async fn initial_pair_capacity_race_preserves_the_connection() -> Result<(), BoxError> { + let mut fixture = Fixture::start_with_reservation_race(true).await?; + let mut client_sibling = timeout(TEST_TIMEOUT, fixture.client_siblings.recv()) + .await? + .ok_or("missing client sibling")?; + let mut server_sibling = timeout(TEST_TIMEOUT, fixture.server_siblings.recv()) + .await? + .ok_or("missing server sibling")?; + let (mut recv, _) = client_sibling.take_stream(SIBLING.kind).unwrap(); + let (_, send) = server_sibling.take_stream(SIBLING.kind).unwrap(); + let ping = frame(1, 19, 64); + timeout(TEST_TIMEOUT, send.send(ping.clone())).await??; + assert_eq!(timeout(TEST_TIMEOUT, recv.recv()).await?, Some(ping)); + let (mut client, mut server) = fixture.sessions().await?; + exchange(&mut client, &mut server).await?; + assert!(fixture.connection.close_reason().is_none()); + fixture.close().await +} + +#[tokio::test] +async fn session_retirement_preserves_remote_cause_and_local_neutrality() -> Result<(), BoxError> { + for streams in [&[DATA][..], &[DATA, REQUESTS][..]] { + for local in [false, true] { + let mut fixture = + RawFixture::start_with_streams(1, Duration::from_secs(3), streams).await?; + let (mut peer_send, _peer_recv) = fixture + .offer(DATA, (streams.len() > 1).then_some(73)) + .await?; + let _requests = if streams.len() > 1 { + Some(fixture.offer(REQUESTS, Some(73)).await?) + } else { + None + }; + let mut peer = timeout(TEST_TIMEOUT, fixture.sessions.recv()) + .await? + .ok_or("missing session")?; + let cancel = peer.service_cancel_token(); + let connection_cancel = peer.cancel_token(); + let (mut recv, send) = peer.take_stream(DATA.kind).unwrap(); + // In a pair, the unused request member waits for the data member. + drop(peer); + assert!(timeout(Duration::from_millis(100), cancel.cancelled()) + .await + .is_err()); + if local { + cancel.cancel(); + } else { + peer_send.finish()?; + } + timeout(TEST_TIMEOUT, cancel.cancelled()).await?; + assert!(timeout(TEST_TIMEOUT, recv.recv()).await?.is_none()); + assert_eq!( + recv.failure(), + (!local).then_some(OrderedStreamFailure::RemoteClose) + ); + assert!(!connection_cancel.is_cancelled()); + drop(recv); + drop(send); + fixture.wait_for_slots(1, TEST_TIMEOUT).await?; + fixture.close().await?; + } + } + Ok(()) +} + +/// Model the request owner's partial-write abort contract at the worker boundary. +#[derive(Debug)] +struct CancellingWriteClaim { + cancel: CancellationToken, + failure: OrderedStreamFailureCause, + expected: Option, +} + +impl crate::zakura::FrameWriteClaim for CancellingWriteClaim { + fn try_start(&self) -> bool { + true + } + + fn written(&self) { + panic!("the peer cannot accept this frame within its receive window"); + } +} + +impl Drop for CancellingWriteClaim { + fn drop(&mut self) { + assert_eq!(self.failure.get(), self.expected); + self.cancel.cancel(); + } +} + +#[tokio::test] +async fn failed_write_records_cause_before_claim_cancels_session() -> Result<(), BoxError> { + let _guard = zakura_test::init(); + let (router, client, connection, remote) = raw_connection().await?; + for expected in [ + Some(OrderedStreamFailure::RemoteClose), + Some(OrderedStreamFailure::WriteTimeout), + None, + ] { + let (mut peer_send, mut peer_recv) = connection.open_bi().await?; + peer_send + .write_all(&frame(1, 0, 0).encode(DATA.frame_cap)?) + .await?; + let (send, recv) = timeout(TEST_TIMEOUT, remote.accept_bi()).await??; + let slots = Arc::new(Semaphore::new(1)); + let mut context = raw_worker_context(&client, slots.clone()); + if expected == Some(OrderedStreamFailure::WriteTimeout) { + context.write_policy = StreamWritePolicy::Timeout(Duration::from_millis(250)); + } + let connection_cancel = context.connection_token.clone(); + let cancel = context.stream_token.clone(); + let failure = OrderedStreamFailureCause::default(); + let prelude = StreamPrelude { + magic: STREAM_PRELUDE_MAGIC, + stream_kind: REQUESTS.kind, + stream_version: REQUESTS.version, + request_id: None, + max_frame_bytes: DATA.frame_cap, + }; + let (inbound_tx, mut inbound_rx) = mpsc::channel(1); + let (outbound_tx, outbound_rx) = worker_framed_channel(1); + let worker = AbortOnDropHandle::new(tokio::spawn(persistent_stream_worker_with_policy( + send, + recv, + prelude, + context, + inbound_tx, + outbound_rx, + 1, + Some(failure.clone()), + None, + ))); + assert_eq!( + timeout(TEST_TIMEOUT, inbound_rx.recv()).await?, + Some(frame(1, 0, 0)) + ); + assert!(outbound_tx.try_reserve_guarded().unwrap().send_request( + frame(1, 17, 1024 * 1024), + Arc::new(CancellingWriteClaim { + cancel: cancel.clone(), + failure: failure.clone(), + expected + }), + )); + // Partial bytes prove the claim started. Withhold the remaining QUIC + // credit so every case exercises a failed or cancelled partial write. + timeout(TEST_TIMEOUT, peer_recv.read_exact(&mut [0; 1])).await??; + match expected { + Some(OrderedStreamFailure::RemoteClose) => peer_recv.stop(0u32.into())?, + Some(OrderedStreamFailure::WriteTimeout) => {} + None => cancel.cancel(), + } + timeout(TEST_TIMEOUT, cancel.cancelled()).await?; + assert_eq!(failure.get(), expected); + timeout(TEST_TIMEOUT, worker).await??; + assert_eq!(slots.available_permits(), 1); + assert!(!connection_cancel.is_cancelled()); + } + connection.close(0u32.into(), b"done"); + timeout(TEST_TIMEOUT, client.close()).await?; + timeout(TEST_TIMEOUT, router.shutdown()).await??; + Ok(()) +} + +#[tokio::test] +async fn abandoned_application_session_releases_capacity_and_preserves_sibling( +) -> Result<(), BoxError> { + for streams in [&[DATA][..], &[DATA, REQUESTS][..]] { + let mut fixture = + RawFixture::start_with_streams(1, Duration::from_secs(3), streams).await?; + let (mut sibling_send, _sibling_recv) = fixture.offer(SIBLING, None).await?; + let mut sibling = timeout(TEST_TIMEOUT, fixture.siblings.recv()) + .await? + .ok_or("no sibling")?; + let (mut sibling_recv, _sibling_send) = sibling.take_stream(SIBLING.kind).unwrap(); + let mut remote_handles = Vec::new(); + for stream in streams { + remote_handles.push( + fixture + .offer(*stream, (streams.len() > 1).then_some(71)) + .await?, + ); + } + let peer = timeout(TEST_TIMEOUT, fixture.sessions.recv()) + .await? + .ok_or("no session")?; + let cancel = peer.service_cancel_token(); + drop(peer); + timeout(TEST_TIMEOUT, cancel.cancelled()).await?; + fixture.wait_for_slots(1, TEST_TIMEOUT).await?; + // A late duplicate cancellation cannot release the reservation twice. + cancel.cancel(); + let ping = frame(1, 19, 8); + sibling_send + .write_all(&ping.encode(SIBLING.frame_cap)?) + .await?; + assert_eq!( + timeout(TEST_TIMEOUT, sibling_recv.recv()).await?, + Some(ping) + ); + assert_eq!(fixture.capacity().available_permits(), 1); + assert!(!sibling.cancel_token().is_cancelled()); + fixture.close().await?; + } + Ok(()) +} + +#[tokio::test] +async fn application_receive_half_or_sender_clone_keeps_session_alive() -> Result<(), BoxError> { + for streams in [&[DATA][..], &[DATA, REQUESTS][..]] { + for keep_receiver in [true, false] { + let mut fixture = + RawFixture::start_with_streams(1, Duration::from_secs(3), streams).await?; + let (mut peer_send, mut peer_recv) = fixture + .offer(DATA, (streams.len() > 1).then_some(74)) + .await?; + let _requests = if streams.len() > 1 { + Some(fixture.offer(REQUESTS, Some(74)).await?) + } else { + None + }; + let mut peer = timeout(TEST_TIMEOUT, fixture.sessions.recv()) + .await? + .ok_or("no session")?; + let cancel = peer.service_cancel_token(); + let (mut recv, send) = peer.take_stream(DATA.kind).unwrap(); + drop(peer); + let ping = frame(1, 19, 8); + if keep_receiver { + drop(send); + assert!(timeout(Duration::from_millis(100), cancel.cancelled()) + .await + .is_err()); + peer_send.write_all(&ping.encode(DATA.frame_cap)?).await?; + assert_eq!(timeout(TEST_TIMEOUT, recv.recv()).await?, Some(ping)); + assert!(recv.failure().is_none()); + assert_eq!(fixture.capacity().available_permits(), 0); + drop(recv); + } else { + let clone = send.clone(); + drop(send); + drop(recv); + assert!(timeout(Duration::from_millis(100), cancel.cancelled()) + .await + .is_err()); + // Incoming traffic cannot retire a session with a retained sender. + peer_send.write_all(&ping.encode(DATA.frame_cap)?).await?; + assert!(timeout(Duration::from_millis(100), cancel.cancelled()) + .await + .is_err()); + timeout(TEST_TIMEOUT, clone.send(ping.clone())).await??; + assert_eq!( + read_frame( + &mut peer_recv, + DATA.frame_cap, + &[], + None, + TEST_TIMEOUT, + Some(TEST_TIMEOUT) + ) + .await?, + ping + ); + assert_eq!(fixture.capacity().available_permits(), 0); + drop(clone); + } + timeout(TEST_TIMEOUT, cancel.cancelled()).await?; + fixture.wait_for_slots(1, TEST_TIMEOUT).await?; + fixture.close().await?; + } + } + Ok(()) +} + +#[tokio::test] +async fn abandoned_application_drains_queued_writes_before_retirement() -> Result<(), BoxError> { + for streams in [&[DATA][..], &[DATA, REQUESTS][..]] { + let mut fixture = + RawFixture::start_with_streams(1, Duration::from_secs(3), streams).await?; + let (mut peer_send, mut peer_recv) = fixture + .offer(DATA, (streams.len() > 1).then_some(72)) + .await?; + let mut requests = if streams.len() > 1 { + Some(fixture.offer(REQUESTS, Some(72)).await?) + } else { + None + }; + let mut peer = timeout(TEST_TIMEOUT, fixture.sessions.recv()) + .await? + .ok_or("no session")?; + let cancel = peer.service_cancel_token(); + let (recv, send) = peer.take_stream(DATA.kind).unwrap(); + drop(recv); + let first = frame(1, 19, 1024 * 1024); + let second = frame(1, 20, 1024 * 1024); + timeout(TEST_TIMEOUT, send.send(first.clone())).await??; + // A second enqueue proves the first frame left the bounded application queue. + timeout(TEST_TIMEOUT, send.send(second.clone())).await??; + drop(send); + // The empty request stream must let the data stream drain before the + // whole session retires, even though all application handles are gone. + drop(peer); + assert!(timeout(Duration::from_millis(100), cancel.cancelled()) + .await + .is_err()); + // The 64 KB QUIC window keeps data writes blocked while this frame reaches + // the dropped receiver, including on the otherwise idle request member. + let incoming = frame(1, 23, 8).encode(DATA.frame_cap)?; + match requests.as_mut() { + Some((request_send, _)) => request_send.write_all(&incoming).await?, + None => peer_send.write_all(&incoming).await?, + } + assert!(timeout(Duration::from_millis(100), cancel.cancelled()) + .await + .is_err()); + assert_eq!(fixture.capacity().available_permits(), 0); + for expected in [first, second] { + assert_eq!( + read_frame( + &mut peer_recv, + DATA.frame_cap, + &[], + None, + TEST_TIMEOUT, + Some(TEST_TIMEOUT) + ) + .await?, + expected + ); + } + timeout(TEST_TIMEOUT, cancel.cancelled()).await?; + fixture.wait_for_slots(1, TEST_TIMEOUT).await?; + assert_eq!( + timeout(TEST_TIMEOUT, peer_recv.read(&mut [0; 1])).await??, + None + ); + assert!(fixture.connection.close_reason().is_none()); + fixture.close().await?; + } + Ok(()) +} diff --git a/crates/zakura-network/src/zakura/handler/tests/connection.rs b/crates/zakura-network/src/zakura/handler/tests/connection.rs new file mode 100644 index 0000000000..1208f3e242 --- /dev/null +++ b/crates/zakura-network/src/zakura/handler/tests/connection.rs @@ -0,0 +1,79 @@ +//! Native connection setup shared by paired transport fixtures. + +use super::*; +use tokio_util::task::AbortOnDropHandle; + +impl ZakuraProtocolHandler { + /// Drive both setup phases directly in tests that isolate admission rules. + pub(in crate::zakura::handler) async fn admit_bi_stream( + &self, + send: SendStream, + recv: RecvStream, + admission: &mut StreamAdmission<'_>, + queue_depth: usize, + exits: mpsc::UnboundedSender, + pending: &mut PendingSessions, + ) -> Option { + let incoming = self.begin_bi_stream_setup(send, recv, admission)?.await?; + self.finish_bi_stream_setup(incoming, admission, queue_depth, exits, pending) + } +} + +pub(in crate::zakura::handler) async fn connect_and_serve( + client: &Endpoint, + address: EndpointAddr, + handler: ZakuraProtocolHandler, + limits: ZakuraLocalLimits, + alpn: &[u8], + deadline: Duration, +) -> Result< + ( + Connection, + AbortOnDropHandle>, + ), + BoxError, +> { + let remote_id = address.id; + let local_id = client.id(); + let connection = timeout(deadline, client.connect(address, alpn)).await??; + let local_peer = ZakuraPeerId::new(local_id.as_bytes().to_vec())?; + let remote_peer = ZakuraPeerId::new(remote_id.as_bytes().to_vec())?; + let conn = ZakuraConnTrace::without_peer(1); + let negotiated = timeout( + deadline, + run_native_initiator_handshake( + &connection, + &limits, + &handler.current_handshake_config(), + &local_peer, + &ZakuraTrace::noop(), + &conn, + ), + ) + .await??; + let remote_ip = confirmed_remote_ip(&connection); + let serving_connection = connection.clone(); + let transport = AbortOnDropHandle::new(tokio::spawn(async move { + handler + .register_and_serve( + serving_connection, + remote_peer, + remote_ip, + ConnectionServeContext { + limits: limits.clamp(&negotiated.limits), + accepted_capabilities: negotiated.accepted_capabilities, + role: "initiator", + direction: ServicePeerDirection::Outbound, + transcript_hash: native_connection_transcript_hash( + ServicePeerDirection::Outbound, + &local_id, + &remote_id, + ), + i_open_collision_winner: i_open_collision_winner(&local_id, &remote_id), + conn, + }, + ) + .await + })); + Ok((connection, transport)) +} diff --git a/crates/zakura-network/src/zakura/header_sync/service.rs b/crates/zakura-network/src/zakura/header_sync/service.rs index 3af949d2c2..5fb49055e2 100644 --- a/crates/zakura-network/src/zakura/header_sync/service.rs +++ b/crates/zakura-network/src/zakura/header_sync/service.rs @@ -18,8 +18,8 @@ use super::{events::*, pipe::run_peer, wire::*, FRAME_HEADER_BYTES}; use crate::zakura::ZakuraSupervisorHandle; use crate::zakura::{ handle_pipe_exit, spawn_supervised_pipe, BoxRunFuture, CloseCause, Frame, FramedRecv, - FramedSend, OrderedSendError, OrderedSessionDemand, OrderedStreamOpening, OrderedStreamPolicy, - Peer, PeerStreamSession, Service, ServicePeerDirection, Sink, SinkReject, Stream, StreamMode, + FramedSend, OrderedSendError, Peer, PeerStreamSession, Service, ServicePeerDirection, + SessionDemand, SessionOpening, SessionPolicy, Sink, SinkReject, Stream, StreamMode, ZakuraConnId, ZakuraPeerId, ZAKURA_CAP_HEADER_SYNC, }; @@ -35,7 +35,7 @@ const HEADER_SYNC_SERVICE_STREAMS: [Stream; 1] = [Stream { version: ZAKURA_HEADER_SYNC_STREAM_VERSION, frame_cap: HEADER_SYNC_FRAME_CAP, capability: ZAKURA_CAP_HEADER_SYNC, - mode: StreamMode::Ordered, + mode: StreamMode::Persistent, }]; /// The sole stream declaration for native header sync. @@ -58,7 +58,7 @@ mod stream_tests { version: 8, frame_cap: HEADER_SYNC_FRAME_CAP, capability: 1 << 5, - mode: StreamMode::Ordered, + mode: StreamMode::Persistent, } ); } @@ -482,12 +482,12 @@ impl HeaderSyncService { self } - fn coordinator_demand(&self) -> Option { + fn coordinator_demand(&self) -> Option { let mut service_demand = self.service_demand.clone()?; if service_demand.borrow().header.is_enabled() { return None; } - Some(OrderedSessionDemand::WaitForChange(Box::pin(async move { + Some(SessionDemand::WaitForChange(Box::pin(async move { loop { if service_demand.changed().await.is_err() { std::future::pending::<()>().await; @@ -509,20 +509,20 @@ impl Service for HeaderSyncService { header_sync_streams() } - fn ordered_stream_policy(&self, _kind: u16) -> OrderedStreamPolicy { - OrderedStreamPolicy { - opening: OrderedStreamOpening::InitiatorOnly, + fn session_policy(&self) -> SessionPolicy { + SessionPolicy { + opening: SessionOpening::InitiatorOnly, reopen: true, } } - fn ordered_session_demand( + fn session_demand( &self, _conn_id: ZakuraConnId, peer: &ZakuraPeerId, _negotiated: u64, direction: ServicePeerDirection, - ) -> OrderedSessionDemand { + ) -> SessionDemand { if let Some(demand) = self.coordinator_demand() { return demand; } @@ -533,7 +533,7 @@ impl Service for HeaderSyncService { ServicePeerDirection::Outbound => snapshot.outbound_slots_free, }; if slots_free == 0 { - return OrderedSessionDemand::WaitForChange(Box::pin(async move { + return SessionDemand::WaitForChange(Box::pin(async move { if peers.changed().await.is_err() { std::future::pending::<()>().await; } @@ -541,7 +541,7 @@ impl Service for HeaderSyncService { } let Some(node_id) = header_peer_node_id(peer) else { - return OrderedSessionDemand::Retire; + return SessionDemand::Retire; }; if self .header_sync @@ -551,12 +551,12 @@ impl Service for HeaderSyncService { { // The reactor publishes the backed-off set without per-node deadlines. // The service re-offers each target after one conservative backoff window. - return OrderedSessionDemand::RetryAt( + return SessionDemand::RetryAt( std::time::Instant::now() + HEADER_SYNC_ADVISORY_BACKOFF, ); } - OrderedSessionDemand::OpenNow + SessionDemand::OpenNow } fn wants_peer( @@ -802,19 +802,19 @@ impl Service for HeaderSyncPassthroughService { header_sync_streams() } - fn ordered_stream_policy(&self, kind: u16) -> OrderedStreamPolicy { - self.inner.ordered_stream_policy(kind) + fn session_policy(&self) -> SessionPolicy { + self.inner.session_policy() } - fn ordered_session_demand( + fn session_demand( &self, conn_id: ZakuraConnId, peer: &ZakuraPeerId, negotiated: u64, direction: ServicePeerDirection, - ) -> OrderedSessionDemand { + ) -> SessionDemand { self.inner - .ordered_session_demand(conn_id, peer, negotiated, direction) + .session_demand(conn_id, peer, negotiated, direction) } fn wants_peer( diff --git a/crates/zakura-network/src/zakura/legacy_gossip.rs b/crates/zakura-network/src/zakura/legacy_gossip.rs index c4d91a3319..799501a2b2 100644 --- a/crates/zakura-network/src/zakura/legacy_gossip.rs +++ b/crates/zakura-network/src/zakura/legacy_gossip.rs @@ -40,10 +40,10 @@ use crate::{ use super::trace::BlockBodySource; use super::{ - spawn_supervised_peer_task, BoxRunFuture, Frame, FramedSend, OrderedSendError, - OrderedSessionDemand, OrderedStreamOpening, OrderedStreamPolicy, Peer, RequestResponseService, - Service as ZakuraService, ServicePeerDirection, SinkReject, Stream, StreamMode, ZakuraConnId, - ZakuraPeerHandle, ZakuraPeerId, ZakuraSupervisorHandle, ZakuraTrace, FRAME_HEADER_BYTES, + spawn_supervised_peer_task, BoxRunFuture, Frame, FramedSend, OrderedSendError, Peer, + RequestResponseService, Service as ZakuraService, ServicePeerDirection, SessionDemand, + SessionOpening, SessionPolicy, SinkReject, Stream, StreamMode, ZakuraConnId, ZakuraPeerHandle, + ZakuraPeerId, ZakuraSupervisorHandle, ZakuraTrace, FRAME_HEADER_BYTES, LOCAL_MAX_CONTROL_FRAME_BYTES, ZAKURA_CAP_LEGACY_GOSSIP, }; @@ -145,7 +145,7 @@ const LEGACY_GOSSIP_SERVICE_STREAMS: [Stream; 2] = [ version: LEGACY_GOSSIP_VERSION, frame_cap: LOCAL_MAX_CONTROL_FRAME_BYTES, capability: ZAKURA_CAP_LEGACY_GOSSIP, - mode: StreamMode::Ordered, + mode: StreamMode::Persistent, }, Stream { kind: ZAKURA_STREAM_LEGACY_REQUESTS, @@ -2544,9 +2544,9 @@ impl ZakuraService for LegacyGossipSink { legacy_gossip_streams() } - fn ordered_stream_policy(&self, _kind: u16) -> OrderedStreamPolicy { - OrderedStreamPolicy { - opening: OrderedStreamOpening::InitiatorOnly, + fn session_policy(&self) -> SessionPolicy { + SessionPolicy { + opening: SessionOpening::InitiatorOnly, reopen: true, } } @@ -2560,17 +2560,17 @@ impl ZakuraService for LegacyGossipSink { self.outbound.owns_connection(peer, conn_id) } - fn ordered_session_demand( + fn session_demand( &self, conn_id: ZakuraConnId, peer: &ZakuraPeerId, _negotiated: u64, _direction: ServicePeerDirection, - ) -> OrderedSessionDemand { + ) -> SessionDemand { if self.outbound.is_retired(peer, conn_id) { - return OrderedSessionDemand::Retire; + return SessionDemand::Retire; } - OrderedSessionDemand::OpenNow + SessionDemand::OpenNow } fn add_peer(&self, mut peer: Peer) { @@ -4704,8 +4704,8 @@ mod tests { "a retired gossip stream must not hold a reopen-gap claim" ); assert!(matches!( - sink.ordered_session_demand(conn_id, &peer_id, 0, ServicePeerDirection::Outbound), - OrderedSessionDemand::Retire, + sink.session_demand(conn_id, &peer_id, 0, ServicePeerDirection::Outbound), + SessionDemand::Retire, )); let (send, _rx) = framed_channel(1); assert!( @@ -4767,8 +4767,8 @@ mod tests { "a reset churn counter must keep the reopen-gap claim" ); assert!(matches!( - sink.ordered_session_demand(conn_id, &peer_id, 0, ServicePeerDirection::Outbound), - OrderedSessionDemand::OpenNow, + sink.session_demand(conn_id, &peer_id, 0, ServicePeerDirection::Outbound), + SessionDemand::OpenNow, )); } diff --git a/crates/zakura-network/src/zakura/regulation/mod.rs b/crates/zakura-network/src/zakura/regulation/mod.rs index bdfbd3b3bc..2299766917 100644 --- a/crates/zakura-network/src/zakura/regulation/mod.rs +++ b/crates/zakura-network/src/zakura/regulation/mod.rs @@ -1,27 +1,18 @@ -//! Reusable resource accounting for native Zakura services. -//! -//! This facade provides the ownership mechanics shared by message-specific -//! policies. It does not decide what a message costs or what should happen -//! when capacity is unavailable. Those decisions stay with each service. +#![allow(dead_code, unused_imports)] // activated by the serving migration -#[allow(dead_code)] // used by the first message policy in the stacked PR -mod outstanding_bytes; -#[allow(dead_code)] // used by the first message policy in the stacked PR -mod rate; -#[allow(dead_code)] // used by the first message policy in the stacked PR -mod slots; +//! Shared admission and ownership for native Zakura message policies. +//! +//! Finite request policies supply their codec and response bound. The shared +//! admission path owns concurrency, rollback, execution, and response lifetimes. +//! Peer routines retain protocol dispatch and scheduling decisions. -#[allow(unused_imports)] // used by the first message policy in the stacked PR -pub(crate) use outstanding_bytes::{ - FrameLease, OutstandingByteBudget, OutstandingByteReservation, OutstandingCapacityError, +mod request; +pub(crate) use request::{ + RequestAdmission, RequestPolicy, RequestSession, ResponsePermit, WorkAttempt, WorkLease, }; -#[allow(unused_imports)] // used by the first message policy in the stacked PR -pub(crate) use rate::{ - CommittedRateReservation, RateBudget, RateBudgetConfigError, RateReservation, - RateReservationError, RateReservationSpendError, -}; -#[allow(unused_imports)] // used by the first message policy in the stacked PR -pub(crate) use slots::{SlotBudget, SlotBudgetCapacityError, SlotPermit}; + +mod slots; +pub(crate) use slots::{SlotBudget, SlotPermit}; #[cfg(test)] mod tests; diff --git a/crates/zakura-network/src/zakura/regulation/outstanding_bytes.rs b/crates/zakura-network/src/zakura/regulation/outstanding_bytes.rs deleted file mode 100644 index 8f50868591..0000000000 --- a/crates/zakura-network/src/zakura/regulation/outstanding_bytes.rs +++ /dev/null @@ -1,255 +0,0 @@ -//! Linear accounting for response bytes until their transport owner drops them. - -use std::sync::{ - atomic::{AtomicU64, Ordering}, - Arc, -}; - -use thiserror::Error; -use tokio::sync::{futures::Notified, Notify}; - -/// A reservation exceeded an outstanding-byte budget's total capacity. -#[derive(Copy, Clone, Debug, Eq, Error, PartialEq)] -#[error("requested {requested} outstanding bytes exceeds capacity {capacity}")] -pub(crate) struct OutstandingCapacityError { - /// Requested byte count. - pub(crate) requested: u64, - /// Configured byte capacity. - pub(crate) capacity: u64, -} - -/// Shared capacity returned only when outstanding byte ownership ends. -/// -/// Unlike [`RateBudget`](super::RateBudget), this balance does not refill with -/// time. Clones reserve and release against the same atomic counter, allowing a -/// node budget and a peer budget to be held by different tasks. -#[derive(Clone, Debug)] -pub(crate) struct OutstandingByteBudget { - inner: Arc, -} - -#[derive(Debug)] -struct OutstandingByteBudgetInner { - capacity: u64, - reserved: AtomicU64, - capacity_released: Notify, -} - -impl OutstandingByteBudget { - /// Create an empty budget with `capacity` bytes. - pub(crate) fn new(capacity: u64) -> Self { - Self { - inner: Arc::new(OutstandingByteBudgetInner { - capacity, - reserved: AtomicU64::new(0), - capacity_released: Notify::new(), - }), - } - } - - /// Return the configured byte capacity. - pub(crate) fn capacity(&self) -> u64 { - self.inner.capacity - } - - /// Return the bytes currently available for reservation. - pub(crate) fn available(&self) -> u64 { - self.capacity().saturating_sub(self.reserved()) - } - - /// Return the bytes currently owned by reservations and frame leases. - pub(crate) fn reserved(&self) -> u64 { - self.inner.reserved.load(Ordering::Acquire) - } - - /// Reserve bytes and return their linear owner. - /// - /// `Err` means the request can never fit. `Ok(None)` means it can fit after - /// another owner releases capacity. - pub(crate) fn try_reserve( - &self, - bytes: u64, - ) -> Result, OutstandingCapacityError> { - self.ensure_request_fits(bytes)?; - - if !self.reserve_bytes(bytes) { - return Ok(None); - } - - Ok(Some(OutstandingByteReservation { - budget: self.clone(), - remaining: bytes, - })) - } - - /// Wait until `bytes` could fit, without reserving them. - /// - /// The notification is registered before capacity is rechecked, so a - /// concurrent release cannot be missed. - pub(crate) async fn wait_for(&self, bytes: u64) -> Result<(), OutstandingCapacityError> { - self.ensure_request_fits(bytes)?; - - loop { - let released = self.inner.capacity_released.notified(); - tokio::pin!(released); - Notified::enable(released.as_mut()); - - if self.available() >= bytes { - return Ok(()); - } - - released.await; - } - } - - fn ensure_request_fits(&self, bytes: u64) -> Result<(), OutstandingCapacityError> { - if bytes > self.capacity() { - return Err(OutstandingCapacityError { - requested: bytes, - capacity: self.capacity(), - }); - } - - Ok(()) - } - - fn reserve_bytes(&self, bytes: u64) -> bool { - if bytes == 0 { - return true; - } - - let mut reserved = self.reserved(); - loop { - if bytes > self.capacity().saturating_sub(reserved) { - return false; - } - - let next = reserved.saturating_add(bytes); - match self.inner.reserved.compare_exchange_weak( - reserved, - next, - Ordering::AcqRel, - Ordering::Acquire, - ) { - Ok(_) => return true, - Err(observed) => reserved = observed, - } - } - } - - fn release_bytes(&self, bytes: u64) { - if bytes == 0 { - return; - } - - let mut reserved = self.reserved(); - loop { - let next = reserved.saturating_sub(bytes); - match self.inner.reserved.compare_exchange_weak( - reserved, - next, - Ordering::AcqRel, - Ordering::Acquire, - ) { - Ok(_) => break, - Err(observed) => reserved = observed, - } - } - self.inner.capacity_released.notify_waiters(); - } -} - -/// Linear ownership of bytes reserved but not yet handed to a frame. -#[derive(Debug)] -#[must_use = "dropping an outstanding-byte reservation releases it"] -pub(crate) struct OutstandingByteReservation { - budget: OutstandingByteBudget, - remaining: u64, -} - -impl OutstandingByteReservation { - /// Return the bytes this reservation still owns directly. - pub(crate) fn remaining(&self) -> u64 { - self.remaining - } - - /// Transfer equal bytes from every reservation into one frame lease. - /// - /// This operation validates every reservation before changing any of them. - /// It is intended for the peer and node response-byte reservations attached - /// to the same outbound frame. - pub(crate) fn transfer_to_frame( - reservations: [&mut OutstandingByteReservation; N], - bytes: u64, - ) -> Option { - if reservations - .iter() - .any(|reservation| reservation.remaining < bytes) - { - return None; - } - - let mut releases = Vec::with_capacity(N); - for reservation in reservations { - reservation.remaining -= bytes; - releases.push(FrameLeaseRelease { - budget: reservation.budget.clone(), - bytes, - }); - } - - Some(FrameLease { - accounted_bytes: bytes, - releases, - }) - } - - /// Release all remaining bytes now. - pub(crate) fn release(self) { - drop(self); - } -} - -impl Drop for OutstandingByteReservation { - fn drop(&mut self) { - self.budget.release_bytes(self.remaining); - self.remaining = 0; - } -} - -/// Linear ownership of one queued frame's outstanding-byte charges. -#[derive(Debug)] -#[must_use = "the transport must retain a frame lease until write completion or drop"] -pub(crate) struct FrameLease { - accounted_bytes: u64, - releases: Vec, -} - -#[derive(Debug)] -struct FrameLeaseRelease { - budget: OutstandingByteBudget, - bytes: u64, -} - -impl FrameLease { - /// Return the response bytes represented by this lease. - pub(crate) fn accounted_bytes(&self) -> u64 { - self.accounted_bytes - } - - #[cfg(test)] - pub(crate) fn empty_for_test() -> Self { - Self { - accounted_bytes: 0, - releases: Vec::new(), - } - } -} - -impl Drop for FrameLease { - fn drop(&mut self) { - for release in self.releases.drain(..) { - release.budget.release_bytes(release.bytes); - } - } -} diff --git a/crates/zakura-network/src/zakura/regulation/rate.rs b/crates/zakura-network/src/zakura/regulation/rate.rs deleted file mode 100644 index 103dcb7ab7..0000000000 --- a/crates/zakura-network/src/zakura/regulation/rate.rs +++ /dev/null @@ -1,356 +0,0 @@ -//! Refundable reservations from a monotonic rate budget. - -use std::{ - sync::{Arc, Mutex, MutexGuard}, - time::Duration, -}; - -use thiserror::Error; -use tokio::{ - sync::{futures::Notified, Notify}, - time::sleep, -}; - -use crate::zakura::transport::{Clock, RealClock}; - -const NANOS_PER_SECOND: u128 = 1_000_000_000; - -/// Invalid local configuration for a rate budget. -#[derive(Copy, Clone, Debug, Eq, Error, PartialEq)] -pub(crate) enum RateBudgetConfigError { - /// A zero capacity cannot admit any positive work. - #[error("rate budget capacity must be greater than zero")] - ZeroCapacity, - /// A zero refill would make spent capacity unavailable forever. - #[error("rate budget refill must be greater than zero")] - ZeroRefill, -} - -/// A rate reservation could not be admitted. -#[derive(Copy, Clone, Debug, Eq, Error, PartialEq)] -pub(crate) enum RateReservationError { - /// The requested amount can never fit in this budget. - #[error("requested rate reservation {requested} exceeds capacity {capacity}")] - ExceedsCapacity { - /// Requested units. - requested: u64, - /// Configured burst capacity. - capacity: u64, - }, - /// The request can fit after refill or an earlier reservation is returned. - #[error("rate reservation is temporarily unavailable for {retry_after:?}")] - TemporarilyUnavailable { - /// Refill time assuming no earlier return. - retry_after: Duration, - }, -} - -impl RateReservationError { - /// Return the refill delay for a temporary rejection. - pub(crate) fn retry_after(self) -> Option { - match self { - Self::TemporarilyUnavailable { retry_after } => Some(retry_after), - Self::ExceedsCapacity { .. } => None, - } - } -} - -/// A caller tried to spend more than its refundable reservation. -#[derive(Copy, Clone, Debug, Eq, Error, PartialEq)] -#[error("rate spend {spent} exceeds refundable reservation {refundable}")] -pub(crate) struct RateReservationSpendError { - /// Units the caller tried to spend. - pub(crate) spent: u64, - /// Units that were still refundable. - pub(crate) refundable: u64, -} - -/// Shared tokens that bound a burst and sustained work rate. -/// -/// Each instance has one caller-defined unit, such as response bytes or a -/// measured work unit. Incomparable resources use separate budgets. Time -/// replenishes spent tokens; dropping an uncommitted reservation returns its -/// tokens immediately. -#[derive(Clone, Debug)] -pub(crate) struct RateBudget { - inner: Arc>, -} - -#[derive(Debug)] -struct RateBudgetInner { - capacity: u64, - refill_per_second: u64, - clock: C, - state: Mutex, - tokens_returned: Notify, -} - -#[derive(Debug)] -struct RateState { - available: u64, - /// Fractional unit numerator in unit-nanoseconds. - refill_remainder: u128, - last_refill: tokio::time::Instant, -} - -impl RateBudget { - /// Create a production budget initialized at full capacity. - pub(crate) fn new( - capacity: u64, - refill_per_second: u64, - ) -> Result { - Self::with_clock(capacity, refill_per_second, RealClock) - } - - /// Wait until `units` could be reserved, without consuming them. - /// - /// The notification is registered before the balance is rechecked, so a - /// concurrent return cannot be missed. - pub(crate) async fn wait_for(&self, units: u64) -> Result<(), RateReservationError> { - self.ensure_reservation_fits(units)?; - - loop { - let returned = self.inner.tokens_returned.notified(); - tokio::pin!(returned); - Notified::enable(returned.as_mut()); - - let Some(retry_after) = self.time_until_available(units) else { - return Ok(()); - }; - - tokio::select! { - _ = &mut returned => {} - _ = sleep(retry_after) => {} - } - } - } -} - -impl RateBudget { - /// Create a budget with an injected monotonic clock. - pub(crate) fn with_clock( - capacity: u64, - refill_per_second: u64, - clock: C, - ) -> Result { - if capacity == 0 { - return Err(RateBudgetConfigError::ZeroCapacity); - } - if refill_per_second == 0 { - return Err(RateBudgetConfigError::ZeroRefill); - } - - let now = clock.now(); - Ok(Self { - inner: Arc::new(RateBudgetInner { - capacity, - refill_per_second, - clock, - state: Mutex::new(RateState { - available: capacity, - refill_remainder: 0, - last_refill: now, - }), - tokens_returned: Notify::new(), - }), - }) - } - - /// Return the configured burst capacity. - pub(crate) fn capacity(&self) -> u64 { - self.inner.capacity - } - - /// Return the configured refill rate in units per second. - pub(crate) fn refill_per_second(&self) -> u64 { - self.inner.refill_per_second - } - - /// Return the currently available units after applying elapsed refill. - pub(crate) fn available(&self) -> u64 { - let mut state = self.lock_state(); - self.refill(&mut state); - state.available - } - - /// Reserve units now or return why they are unavailable. - pub(crate) fn try_reserve( - &self, - units: u64, - ) -> Result, RateReservationError> { - self.ensure_reservation_fits(units)?; - - let mut state = self.lock_state(); - self.refill(&mut state); - if state.available < units { - return Err(RateReservationError::TemporarilyUnavailable { - retry_after: self.retry_after(&state, units), - }); - } - - state.available -= units; - Ok(RateReservation { - budget: self.clone(), - refundable: units, - }) - } - - fn ensure_reservation_fits(&self, units: u64) -> Result<(), RateReservationError> { - if units > self.capacity() { - return Err(RateReservationError::ExceedsCapacity { - requested: units, - capacity: self.capacity(), - }); - } - - Ok(()) - } - - fn time_until_available(&self, units: u64) -> Option { - let mut state = self.lock_state(); - self.refill(&mut state); - (state.available < units).then(|| self.retry_after(&state, units)) - } - - fn retry_after(&self, state: &RateState, units: u64) -> Duration { - let deficit = units.saturating_sub(state.available); - let scaled_deficit = u128::from(deficit) - .saturating_mul(NANOS_PER_SECOND) - .saturating_sub(state.refill_remainder); - let refill = u128::from(self.refill_per_second().max(1)); - let nanos = scaled_deficit.saturating_add(refill - 1) / refill; - Duration::from_nanos(u64::try_from(nanos).unwrap_or(u64::MAX)) - } - - fn refill(&self, state: &mut RateState) { - let now = self.inner.clock.now(); - let elapsed = now.saturating_duration_since(state.last_refill); - if elapsed.is_zero() || state.available == self.capacity() { - state.last_refill = now; - if state.available == self.capacity() { - state.refill_remainder = 0; - } - return; - } - - let produced = elapsed - .as_nanos() - .saturating_mul(u128::from(self.refill_per_second())) - .saturating_add(state.refill_remainder); - let whole_units = u64::try_from(produced / NANOS_PER_SECOND).unwrap_or(u64::MAX); - state.available = self - .capacity() - .min(state.available.saturating_add(whole_units)); - state.refill_remainder = if state.available == self.capacity() { - 0 - } else { - produced % NANOS_PER_SECOND - }; - state.last_refill = now; - } - - fn return_units(&self, units: u64) { - if units == 0 { - return; - } - - let mut state = self.lock_state(); - self.refill(&mut state); - state.available = self.capacity().min(state.available.saturating_add(units)); - if state.available == self.capacity() { - state.refill_remainder = 0; - } - drop(state); - self.inner.tokens_returned.notify_waiters(); - } - - fn lock_state(&self) -> MutexGuard<'_, RateState> { - self.inner - .state - .lock() - .expect("rate budget mutex should not be poisoned") - } -} - -/// Provisional ownership of a fully refundable rate reservation. -#[derive(Debug)] -#[must_use = "dropping a provisional rate reservation returns it"] -pub(crate) struct RateReservation { - budget: RateBudget, - refundable: u64, -} - -impl RateReservation { - /// Commit the reservation and permanently spend `initial_spend` units. - pub(crate) fn commit( - mut self, - initial_spend: u64, - ) -> Result, RateReservationSpendError> { - if initial_spend > self.refundable { - return Err(RateReservationSpendError { - spent: initial_spend, - refundable: self.refundable, - }); - } - - self.refundable -= initial_spend; - let refundable = self.refundable; - self.refundable = 0; - Ok(CommittedRateReservation { - budget: self.budget.clone(), - refundable, - }) - } - - /// Return the units held by this provisional reservation. - pub(crate) fn reserved(&self) -> u64 { - self.refundable - } -} - -impl Drop for RateReservation { - fn drop(&mut self) { - self.budget.return_units(self.refundable); - self.refundable = 0; - } -} - -/// A committed rate reservation whose unused units remain refundable. -#[derive(Debug)] -#[must_use = "dropping a committed rate reservation returns its unused units"] -pub(crate) struct CommittedRateReservation { - budget: RateBudget, - refundable: u64, -} - -impl CommittedRateReservation { - /// Permanently spend units consumed by completed work. - pub(crate) fn spend(&mut self, units: u64) -> Result<(), RateReservationSpendError> { - if units > self.refundable { - return Err(RateReservationSpendError { - spent: units, - refundable: self.refundable, - }); - } - - self.refundable -= units; - Ok(()) - } - - /// Return the units that will be returned when this owner is dropped. - pub(crate) fn refundable(&self) -> u64 { - self.refundable - } - - /// Finish this reservation and return its unused units now. - pub(crate) fn finish(self) { - drop(self); - } -} - -impl Drop for CommittedRateReservation { - fn drop(&mut self) { - self.budget.return_units(self.refundable); - self.refundable = 0; - } -} diff --git a/crates/zakura-network/src/zakura/regulation/request.rs b/crates/zakura-network/src/zakura/regulation/request.rs new file mode 100644 index 0000000000..d9abb9bd6c --- /dev/null +++ b/crates/zakura-network/src/zakura/regulation/request.rs @@ -0,0 +1,273 @@ +//! Admission and ownership for requests that produce a finite response. +//! +//! Policies supply the codec and response bound. Sequential serving tasks wait +//! for admission before dispatching work; this layer owns capacity and lifetimes. + +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, +}; + +use tokio_util::sync::CancellationToken; + +use super::{slots::WeakSlotBudget, SlotBudget, SlotPermit}; +use crate::zakura::transport::{Frame, FrameGuard}; +use crate::zakura::ZakuraPeerId; + +/// Message-specific rules used by request admission. +/// +/// Decode must use the production codec and validate the request before returning +/// it. The response bound must cover every frame the handler may produce. This +/// interface is for finite requests, not announcements or subscription lifetimes. +pub(crate) trait RequestPolicy { + type Request; + type Error; + + fn decode(&self, frame: Frame) -> Result; + fn response_cap(&self, request: &Self::Request) -> u64; +} + +/// Node work capacity shared by the sessions of one configured request policy. +#[derive(Clone, Debug)] +pub(crate) struct RequestAdmission

{ + policy: P, + node: SlotBudget, + peer_capacity: usize, + peers: Arc>>, +} + +impl RequestAdmission

{ + pub(crate) fn new(policy: P, node: SlotBudget, peer_capacity: usize) -> Self { + // Validate before a session is created rather than panicking on ingress. + SlotBudget::new(peer_capacity).expect("request peer capacity is validated"); + Self { + policy, + node, + peer_capacity, + peers: Arc::new(Mutex::new(HashMap::new())), + } + } + + /// Reconnects share the authenticated peer's capacity with its old work. + pub(crate) fn session(&self, peer: &ZakuraPeerId) -> RequestSession

{ + let mut peers = self + .peers + .lock() + .expect("peer budget registry is not poisoned"); + // Permits keep their semaphore alive even after the old session closes. + // Remove expired identities on each connection so churn cannot grow the map. + peers.retain(|_, budget| budget.is_alive()); + let budget = peers + .get(peer) + .and_then(WeakSlotBudget::upgrade) + .unwrap_or_else(|| { + let budget = SlotBudget::new(self.peer_capacity) + .expect("request peer capacity was validated at construction"); + peers.insert(peer.clone(), budget.downgrade()); + budget + }); + RequestSession { + policy: self.policy.clone(), + node: self.node.clone(), + peer: budget, + } + } + + #[cfg(test)] + pub(crate) fn reserved_by_peers(&self) -> usize { + self.peers + .lock() + .unwrap() + .values() + .filter_map(WeakSlotBudget::upgrade) + .map(|budget| budget.reserved()) + .sum() + } +} + +/// One session's request policy, sharing capacity with the peer's other sessions. +#[derive(Clone, Debug)] +pub(crate) struct RequestSession

{ + policy: P, + node: SlotBudget, + peer: SlotBudget, +} + +impl RequestSession

{ + pub(crate) fn decode(&self, frame: Frame) -> Result { + self.policy.decode(frame) + } + + /// One sequential serving task waits for its peer before entering the node + /// queue. A previous response cannot make this peer hold extra node slots. + /// Dropping this future removes its FIFO waiter and releases a partial claim. + pub(crate) async fn admit(&self, request: &P::Request) -> WorkAttempt { + let peer = reserve_response_slot(&self.peer, WorkBound::Peer).await; + let node = reserve_response_slot(&self.node, WorkBound::Node).await; + WorkAttempt { + resources: Arc::new(WorkResources { + _peer: peer, + _node: node, + }), + response_cap: self.policy.response_cap(request), + } + } + + #[cfg(test)] + pub(crate) fn peer_budget(&self) -> &SlotBudget { + &self.peer + } +} + +/// Scope of the work capacity that delayed a request. +#[derive(Copy, Clone, Debug, Eq, PartialEq)] +pub(crate) enum WorkBound { + Peer, + Node, +} + +impl WorkBound { + /// Stable resource names used by delay metrics and traces. + pub(crate) fn label(self) -> &'static str { + match self { + Self::Peer => "peer_active", + Self::Node => "node_active", + } + } +} + +async fn reserve_response_slot(budget: &SlotBudget, bound: WorkBound) -> SlotPermit { + if let Some(permit) = budget.try_reserve() { + return permit; + } + metrics::counter!("sync.block.serving.delayed", "bound" => bound.label()).increment(1); + budget.reserve().await +} + +/// Provisional work ownership. Dropping it rolls back admission. +#[derive(Debug)] +#[must_use = "dropping an admission attempt returns its capacity"] +pub(crate) struct WorkAttempt { + resources: Arc, + response_cap: u64, +} + +impl WorkAttempt { + pub(crate) fn commit(self) -> ResponsePermit { + ResponsePermit { + resources: self.resources, + execution: Arc::new(Execution::default()), + response_cap: self.response_cap, + queued_bytes: 0, + } + } +} + +#[derive(Debug)] +pub(crate) struct WorkResources { + _peer: SlotPermit, + _node: SlotPermit, +} + +/// The handler's response ownership. +/// +/// Dropping it prevents unclaimed work from starting and signals cancellation +/// to already running work. This does not interrupt an operation in progress or +/// release capacity held by work leases and queued frame guards. +#[derive(Debug)] +#[must_use = "retain the response permit until settlement or cancellation"] +pub(crate) struct ResponsePermit { + resources: Arc, + execution: Arc, + response_cap: u64, + queued_bytes: u64, +} + +impl ResponsePermit { + pub(crate) fn can_queue_frame(&self, bytes: u64) -> bool { + bytes <= self.response_cap.saturating_sub(self.queued_bytes) + } + + /// Call only once queue capacity is reserved. The guard lives through writing. + pub(crate) fn frame_guard(&mut self, bytes: u64) -> FrameGuard { + assert!( + self.can_queue_frame(bytes), + "encoded response fits its declared cap" + ); + self.queued_bytes += bytes; + FrameGuard::new(self.resources.clone()) + } + + pub(crate) fn work_lease(&self) -> WorkLease { + WorkLease { + _resources: self.resources.clone(), + execution: self.execution.clone(), + } + } +} + +impl Drop for ResponsePermit { + fn drop(&mut self) { + self.execution.cancel(); + } +} + +#[derive(Debug, Default, PartialEq, Eq)] +enum ExecutionState { + #[default] + Queued, + Started, + Cancelled, +} + +#[derive(Debug, Default)] +struct Execution { + state: Mutex, + cancelled: CancellationToken, +} + +impl Execution { + fn try_start(&self) -> bool { + let mut state = self + .state + .lock() + .expect("request execution state is not poisoned"); + if *state != ExecutionState::Queued { + return false; + } + *state = ExecutionState::Started; + true + } + + fn cancel(&self) { + *self + .state + .lock() + .expect("request execution state is not poisoned") = ExecutionState::Cancelled; + self.cancelled.cancel(); + } +} + +/// Capacity shared by an execution and its returned result. +/// +/// Clones share one execution claim. Cancellation prevents an unclaimed start +/// and lets running work skip further steps at its next cancellation check. +/// An operation and its returned result must retain this lease until they are +/// released, even after cancellation. +#[derive(Clone, Debug)] +pub(crate) struct WorkLease { + _resources: Arc, + execution: Arc, +} + +impl WorkLease { + pub(crate) fn try_start(&self) -> bool { + self.execution.try_start() + } + pub(crate) fn is_cancelled(&self) -> bool { + self.execution.cancelled.is_cancelled() + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/zakura-network/src/zakura/regulation/request/tests.rs b/crates/zakura-network/src/zakura/regulation/request/tests.rs new file mode 100644 index 0000000000..48ef76c724 --- /dev/null +++ b/crates/zakura-network/src/zakura/regulation/request/tests.rs @@ -0,0 +1,307 @@ +//! A second finite-request adapter, using the production discovery codec. +//! This exercises reuse without enabling discovery regulation in production. + +use std::time::Duration; + +use futures::FutureExt; + +use super::*; +use crate::zakura::discovery::{DiscoveryMessage, DiscoveryWireError, MAX_DISCOVERY_MESSAGE_BYTES}; + +impl ResponsePermit { + pub(crate) fn weak_resources(&self) -> std::sync::Weak { + Arc::downgrade(&self.resources) + } +} + +#[derive(Clone, Debug)] +pub(super) struct GetPeersPolicy; + +#[derive(Debug)] +pub(super) struct GetPeersRequest { + limit: u16, +} + +impl RequestPolicy for GetPeersPolicy { + type Request = GetPeersRequest; + type Error = String; + + fn decode(&self, frame: Frame) -> Result { + if frame.message_type != 1 || frame.flags != 0 { + return Err("invalid discovery envelope".into()); + } + match DiscoveryMessage::decode(&frame.payload) + .map_err(|e: DiscoveryWireError| e.to_string())? + { + DiscoveryMessage::GetPeers { limit, .. } => Ok(GetPeersRequest { limit }), + _ => Err("expected GetPeers".into()), + } + } + + fn response_cap(&self, request: &GetPeersRequest) -> u64 { + assert!(request.limit > 0); + // Today's codec caps the complete discovery response, independently of + // the smaller per-record limits proposed in the regulation draft. + u64::try_from(MAX_DISCOVERY_MESSAGE_BYTES).unwrap() + } +} + +pub(super) fn frame(limit: u16) -> Frame { + Frame { + message_type: 1, + flags: 0, + payload: DiscoveryMessage::GetPeers { + limit, + wanted_services: vec![], + exclude_node_ids: vec![], + } + .encode() + .unwrap(), + } +} + +fn peer(byte: u8) -> ZakuraPeerId { + ZakuraPeerId::new(vec![byte; 32]).unwrap() +} + +fn admission(node: &SlotBudget) -> RequestAdmission { + RequestAdmission::new(GetPeersPolicy, node.clone(), 1) +} + +#[test] +fn discovery_codec_runs_before_work_admission() { + let node = SlotBudget::new(1).unwrap(); + let session = admission(&node).session(&peer(1)); + let mut malformed = frame(1); + malformed.payload.push(0); + assert!(session.decode(malformed).is_err()); + assert_eq!(node.reserved(), 0); + let request = session.decode(frame(1)).unwrap(); + assert_eq!(request.limit, 1); + let attempt = session.admit(&request).now_or_never().unwrap(); + assert_eq!(node.reserved(), 1); + drop(attempt); + assert_eq!(node.reserved(), 0); +} + +#[tokio::test] +async fn finite_discovery_response_retains_work_through_its_write() { + use crate::zakura::transport::worker_framed_channel; + let node = SlotBudget::new(1).unwrap(); + let session = admission(&node).session(&peer(1)); + let request = session.decode(frame(1)).unwrap(); + let mut response = session.admit(&request).now_or_never().unwrap().commit(); + let lifetime = response.weak_resources(); + let work = response.work_lease(); + assert!(work.try_start()); + assert!(!work.clone().try_start()); + let (send, mut recv) = worker_framed_channel(1); + let payload = DiscoveryMessage::Peers { records: vec![] } + .encode() + .unwrap(); + let payload_bytes = u64::try_from(payload.len()).unwrap(); + send.try_reserve_guarded().unwrap().send( + Frame { + message_type: 1, + flags: 0, + payload, + }, + response.frame_guard(payload_bytes), + ); + let queued = recv.recv().await.unwrap(); + drop(response); + drop(work); + assert_eq!(node.reserved(), 1); + assert!(lifetime.upgrade().is_some()); + assert!(session.admit(&request).now_or_never().is_none()); + queued + .write_with(|_| async { Ok::<_, std::io::Error>(()) }) + .await + .unwrap(); + assert_eq!(node.reserved(), 0); + assert!(lifetime.upgrade().is_none()); +} + +#[tokio::test] +async fn admission_waiters_retain_peer_capacity_and_receive_node_slots_in_order() { + let node = SlotBudget::new(1).unwrap(); + let admission = admission(&node); + let first = admission.session(&peer(1)); + let second = admission.session(&peer(2)); + let third = admission.session(&peer(3)); + let request = first.decode(frame(1)).unwrap(); + let held = first.admit(&request).await; + let mut second_wait = Box::pin(second.admit(&request)); + let mut third_wait = Box::pin(third.admit(&request)); + assert!(futures::poll!(&mut second_wait).is_pending()); + assert!(futures::poll!(&mut third_wait).is_pending()); + assert_eq!(second.peer_budget().reserved(), 1); + assert_eq!(third.peer_budget().reserved(), 1); + assert_eq!(node.reserved(), 1); + drop(held); + assert!( + node.try_reserve().is_none(), + "the first waiter owns the released slot" + ); + let second_owner = tokio::time::timeout(Duration::from_secs(1), second_wait) + .await + .unwrap(); + assert!(futures::poll!(&mut third_wait).is_pending()); + drop(second_owner); + let third_owner = tokio::time::timeout(Duration::from_secs(1), third_wait) + .await + .unwrap(); + drop(third_owner); + assert_eq!(node.reserved(), 0); + assert_eq!(admission.reserved_by_peers(), 0); +} + +#[test] +fn closing_response_prevents_an_unclaimed_execution() { + let node = SlotBudget::new(1).unwrap(); + let session = admission(&node).session(&peer(1)); + let request = session.decode(frame(1)).unwrap(); + let response = session.admit(&request).now_or_never().unwrap().commit(); + let work = response.work_lease(); + drop(response); + assert!(work.is_cancelled()); + assert!(!work.try_start()); + assert_eq!(node.reserved(), 1); + drop(work); + assert_eq!(node.reserved(), 0); +} + +#[tokio::test] +async fn cancelling_a_node_waiter_releases_its_partial_peer_claim() { + let node = SlotBudget::new(1).unwrap(); + let admission = admission(&node); + let first = admission.session(&peer(1)); + let second = admission.session(&peer(2)); + let request = first.decode(frame(1)).unwrap(); + let held = first.admit(&request).await; + let mut wait = Box::pin(second.admit(&request)); + assert!(futures::poll!(&mut wait).is_pending()); + assert_eq!(second.peer_budget().reserved(), 1); + drop(wait); + assert_eq!(second.peer_budget().reserved(), 0); + assert_eq!(node.reserved(), 1); + drop(held); + assert!(second.admit(&request).now_or_never().is_some()); + assert_eq!(node.reserved(), 0); +} + +#[test] +fn reconnect_waits_for_all_old_read_and_frame_owners() { + for read_finishes_first in [false, true] { + let node = SlotBudget::new(2).unwrap(); + let admission = admission(&node); + let original = admission.session(&peer(1)); + let request = original.decode(frame(1)).unwrap(); + let mut response = original.admit(&request).now_or_never().unwrap().commit(); + let mut work = Some(response.work_lease()); + assert!(work.as_ref().unwrap().try_start()); + let writing = response.frame_guard(1); + let ending = response.frame_guard(1); + drop(response); + assert!( + work.as_ref().unwrap().is_cancelled(), + "producer closure signals cancellation even after execution starts" + ); + assert!(!work.as_ref().unwrap().try_start()); + drop(original); + + for _ in 0..64 { + let replacement = admission.session(&peer(1)); + assert!(replacement.admit(&request).now_or_never().is_none()); + assert_eq!(node.reserved(), 1); + } + let replacement = admission.session(&peer(1)); + assert!(admission + .session(&peer(2)) + .admit(&request) + .now_or_never() + .is_some()); + if read_finishes_first { + drop(work.take()); + } + drop(writing); + assert!(replacement.admit(&request).now_or_never().is_none()); + drop(ending); + if !read_finishes_first { + assert!(replacement.admit(&request).now_or_never().is_none()); + drop(work); + } + assert!(replacement.admit(&request).now_or_never().is_some()); + assert_eq!(node.reserved(), 0); + } +} + +#[test] +fn peer_registry_prunes_churn_but_retains_outstanding_work() { + let node = SlotBudget::new(2).unwrap(); + let admission = admission(&node); + let original = admission.session(&peer(0)); + let request = original.decode(frame(1)).unwrap(); + let held = original.admit(&request).now_or_never().unwrap(); + drop(original); + for id in 1..=255 { + drop(admission.session(&peer(id))); + assert_eq!(admission.peers.lock().unwrap().len(), 2); + } + assert!(admission + .session(&peer(0)) + .admit(&request) + .now_or_never() + .is_none()); + drop(held); + let replacement = admission.session(&peer(0)); + assert_eq!(admission.peers.lock().unwrap().len(), 1); + assert!(replacement.admit(&request).now_or_never().is_some()); +} + +#[test] +fn concurrent_sessions_for_one_identity_share_capacity() { + let node = SlotBudget::new(8).unwrap(); + let admission = admission(&node); + let barrier = std::sync::Barrier::new(8); + let sessions = std::thread::scope(|scope| { + let tasks: Vec<_> = (0..8) + .map(|_| { + scope.spawn(|| { + barrier.wait(); + admission.session(&peer(1)) + }) + }) + .collect(); + tasks + .into_iter() + .map(|task| task.join().unwrap()) + .collect::>() + }); + let request = sessions[0].decode(frame(1)).unwrap(); + let held = sessions[0].admit(&request).now_or_never().unwrap(); + for session in &sessions[1..] { + assert!(session.admit(&request).now_or_never().is_none()); + } + drop(held); + assert_eq!(node.reserved(), 0); +} + +#[tokio::test] +async fn a_waiters_peer_slot_survives_session_replacement_and_cancellation() { + let node = SlotBudget::new(1).unwrap(); + let admission = admission(&node); + let original = admission.session(&peer(1)); + let request = original.decode(frame(1)).unwrap(); + let held = original.admit(&request).await; + let mut wait = Box::pin(async move { original.admit(&request).await }); + assert!(futures::poll!(&mut wait).is_pending()); + drop(held); + let replacement = admission.session(&peer(1)); + let request = replacement.decode(frame(1)).unwrap(); + assert!(replacement.admit(&request).now_or_never().is_none()); + // Cancel after capacity is assigned but before the waiting future resumes. + drop(wait); + assert!(replacement.admit(&request).now_or_never().is_some()); + assert_eq!(node.reserved(), 0); +} diff --git a/crates/zakura-network/src/zakura/regulation/slots.rs b/crates/zakura-network/src/zakura/regulation/slots.rs index 19349a5a32..e10ef238fb 100644 --- a/crates/zakura-network/src/zakura/regulation/slots.rs +++ b/crates/zakura-network/src/zakura/regulation/slots.rs @@ -1,6 +1,6 @@ //! Owned permits for bounded collections of retained or active work. -use std::sync::Arc; +use std::sync::{Arc, Weak}; use thiserror::Error; use tokio::sync::{OwnedSemaphorePermit, Semaphore}; @@ -21,6 +21,7 @@ pub(crate) struct SlotBudgetCapacityError { /// admitted item makes ordinary drop and cancellation release capacity. #[derive(Clone, Debug)] pub(crate) struct SlotBudget { + #[cfg(test)] capacity: usize, permits: Arc, } @@ -36,22 +37,35 @@ impl SlotBudget { } Ok(Self { + #[cfg(test)] capacity, permits: Arc::new(Semaphore::new(capacity)), }) } /// Return the maximum number of owned slots. + #[cfg(test)] pub(crate) fn capacity(&self) -> usize { self.capacity } /// Return the number of currently owned slots. + #[cfg(test)] pub(crate) fn reserved(&self) -> usize { self.capacity .saturating_sub(self.permits.available_permits()) } + /// Create a non-owning handle to this same pool for registry bookkeeping. + /// This does not change capacity or release any reserved slots. + pub(super) fn downgrade(&self) -> WeakSlotBudget { + WeakSlotBudget { + #[cfg(test)] + capacity: self.capacity, + permits: Arc::downgrade(&self.permits), + } + } + /// Reserve one slot without waiting. pub(crate) fn try_reserve(&self) -> Option { self.permits @@ -63,8 +77,7 @@ impl SlotBudget { /// Wait for a slot and return its ownership in semaphore queue order. /// - /// Keep the returned permit through admission. If another resource cannot - /// be reserved, drop it before waiting for that resource. Cancelling this + /// Keep the returned permit while owning the resource. Cancelling this /// future removes its waiter without consuming a slot. pub(crate) async fn reserve(&self) -> SlotPermit { let permit = self @@ -77,9 +90,43 @@ impl SlotBudget { } } -/// Linear ownership of one slot. +/// One reserved slot. The permit keeps its pool alive even if the session ends. +/// Dropping the permit returns the slot to that same pool. #[derive(Debug)] #[must_use = "dropping a slot permit releases its capacity"] pub(crate) struct SlotPermit { _permit: OwnedSemaphorePermit, } + +/// A registry handle that can find a slot pool without keeping it alive. +/// +/// Sessions and outstanding permits keep the pool alive. If a peer reconnects +/// while its old work remains, upgrading this handle reuses the same pool, +/// so reconnecting cannot bypass the peer's serving limit. +/// +/// Once all strong owners are gone, upgrading returns `None`. The registry +/// can then remove the stale entry. +#[derive(Debug)] +pub(super) struct WeakSlotBudget { + #[cfg(test)] + capacity: usize, + permits: Weak, +} + +impl WeakSlotBudget { + /// Check whether the pool currently has any strong owners, for registry cleanup. + /// This is a snapshot; use `upgrade` to obtain a handle that keeps it alive. + pub(super) fn is_alive(&self) -> bool { + self.permits.strong_count() > 0 + } + + /// Obtain shared ownership of the existing pool, or `None` if it is gone. + /// This never creates a new pool or resets its available capacity. + pub(super) fn upgrade(&self) -> Option { + Some(SlotBudget { + #[cfg(test)] + capacity: self.capacity, + permits: self.permits.upgrade()?, + }) + } +} diff --git a/crates/zakura-network/src/zakura/regulation/tests.rs b/crates/zakura-network/src/zakura/regulation/tests.rs index fc163edbaf..1dc9301787 100644 --- a/crates/zakura-network/src/zakura/regulation/tests.rs +++ b/crates/zakura-network/src/zakura/regulation/tests.rs @@ -1,225 +1,6 @@ use std::time::Duration; use super::*; -use crate::zakura::testkit::TestClock; - -#[test] -fn rate_budget_rejects_configuration_that_cannot_refill() { - assert_eq!( - RateBudget::with_clock(0, 1, TestClock::new()).unwrap_err(), - RateBudgetConfigError::ZeroCapacity, - ); - assert_eq!( - RateBudget::with_clock(1, 0, TestClock::new()).unwrap_err(), - RateBudgetConfigError::ZeroRefill, - ); -} - -#[test] -fn provisional_rate_reservation_returns_all_units_on_drop() { - let budget = RateBudget::with_clock(100, 10, TestClock::new()) - .expect("the rate budget configuration is valid"); - let reservation = budget - .try_reserve(80) - .expect("the full budget covers eighty units"); - - assert_eq!(budget.available(), 20); - drop(reservation); - assert_eq!(budget.available(), 100); -} - -#[test] -fn committed_rate_reservation_keeps_consumed_units_spent() { - let budget = RateBudget::with_clock(100, 10, TestClock::new()) - .expect("the rate budget configuration is valid"); - let reservation = budget - .try_reserve(80) - .expect("the full budget covers eighty units"); - let mut committed = reservation - .commit(20) - .expect("the initial spend fits the reservation"); - committed - .spend(30) - .expect("completed work fits the refundable remainder"); - - committed.finish(); - - assert_eq!(budget.available(), 50); -} - -#[test] -fn rate_refill_preserves_fractional_time() { - let clock = TestClock::new(); - let budget = RateBudget::with_clock(10, 3, clock.clone()) - .expect("the rate budget configuration is valid"); - budget - .try_reserve(10) - .expect("the full budget covers its capacity") - .commit(10) - .expect("the full reservation can be spent") - .finish(); - - clock.advance(Duration::from_millis(333)); - assert_eq!(budget.available(), 0); - clock.advance(Duration::from_millis(1)); - assert_eq!(budget.available(), 1); -} - -#[tokio::test(start_paused = true)] -async fn rate_wait_wakes_when_a_reservation_is_returned() { - let budget = RateBudget::new(100, 1).expect("the rate budget configuration is valid"); - let reservation = budget - .try_reserve(100) - .expect("the full budget covers its capacity"); - let waiting_budget = budget.clone(); - let waiter = tokio::spawn(async move { waiting_budget.wait_for(50).await }); - tokio::task::yield_now().await; - assert!(!waiter.is_finished()); - - drop(reservation); - - tokio::task::yield_now().await; - waiter - .await - .expect("the waiter task should not panic") - .expect("the returned reservation fits the budget"); -} - -#[tokio::test(start_paused = true)] -async fn rate_wait_wakes_at_the_refill_deadline() { - let budget = RateBudget::new(100, 10).expect("the rate budget configuration is valid"); - budget - .try_reserve(100) - .expect("the full budget covers its capacity") - .commit(100) - .expect("the full reservation can be spent") - .finish(); - let waiting_budget = budget.clone(); - let waiter = tokio::spawn(async move { waiting_budget.wait_for(50).await }); - tokio::task::yield_now().await; - assert!(!waiter.is_finished()); - - tokio::time::advance(Duration::from_secs(4)).await; - tokio::task::yield_now().await; - assert!(!waiter.is_finished()); - - tokio::time::advance(Duration::from_secs(1)).await; - tokio::task::yield_now().await; - waiter - .await - .expect("the waiter task should not panic") - .expect("five seconds refill fifty units"); -} - -#[tokio::test] -async fn impossible_waits_return_capacity_errors() { - let outstanding = OutstandingByteBudget::new(10); - assert_eq!( - outstanding.wait_for(11).await, - Err(OutstandingCapacityError { - requested: 11, - capacity: 10, - }) - ); - - let rate = RateBudget::new(10, 1).expect("the rate budget configuration is valid"); - assert_eq!( - rate.wait_for(11).await, - Err(RateReservationError::ExceedsCapacity { - requested: 11, - capacity: 10, - }) - ); -} - -#[tokio::test] -async fn outstanding_wait_does_not_miss_a_concurrent_release() { - let budget = OutstandingByteBudget::new(10); - let reservation = budget - .try_reserve(10) - .expect("the request fits the budget") - .expect("the empty budget admits the reservation"); - let waiting_budget = budget.clone(); - let waiter = tokio::spawn(async move { waiting_budget.wait_for(1).await }); - tokio::task::yield_now().await; - - drop(reservation); - - tokio::time::timeout(Duration::from_secs(1), waiter) - .await - .expect("the release must wake the waiter") - .expect("the waiter task should not panic") - .expect("one byte fits after release"); -} - -#[tokio::test(start_paused = true)] -async fn outstanding_budget_does_not_refill_with_time() { - let budget = OutstandingByteBudget::new(10); - let reservation = budget - .try_reserve(10) - .expect("the request fits the budget") - .expect("the empty budget admits the reservation"); - - tokio::time::advance(Duration::from_secs(24 * 60 * 60)).await; - - assert_eq!(budget.available(), 0); - reservation.release(); - assert_eq!(budget.available(), 10); -} - -#[test] -fn frame_lease_holds_peer_and_node_bytes() { - let node = OutstandingByteBudget::new(100); - let peer = OutstandingByteBudget::new(100); - let mut node_reservation = node - .try_reserve(100) - .expect("the request fits the node budget") - .expect("the node budget is empty"); - let mut peer_reservation = peer - .try_reserve(100) - .expect("the request fits the peer budget") - .expect("the peer budget is empty"); - - let lease = OutstandingByteReservation::transfer_to_frame( - [&mut node_reservation, &mut peer_reservation], - 60, - ) - .expect("both reservations cover the frame"); - assert_eq!(lease.accounted_bytes(), 60); - - drop(node_reservation); - drop(peer_reservation); - assert_eq!(node.reserved(), 60); - assert_eq!(peer.reserved(), 60); - - drop(lease); - assert_eq!(node.reserved(), 0); - assert_eq!(peer.reserved(), 0); -} - -#[test] -fn failed_multi_budget_transfer_changes_nothing() { - let first = OutstandingByteBudget::new(100); - let second = OutstandingByteBudget::new(50); - let mut first_reservation = first - .try_reserve(100) - .expect("the request fits the first budget") - .expect("the first budget is empty"); - let mut second_reservation = second - .try_reserve(50) - .expect("the request fits the second budget") - .expect("the second budget is empty"); - - assert!(OutstandingByteReservation::transfer_to_frame( - [&mut first_reservation, &mut second_reservation], - 60, - ) - .is_none()); - assert_eq!(first_reservation.remaining(), 100); - assert_eq!(second_reservation.remaining(), 50); - assert_eq!(first.reserved(), 100); - assert_eq!(second.reserved(), 50); -} #[test] fn slot_permits_bound_and_release_owned_items() { diff --git a/crates/zakura-network/src/zakura/testkit/cluster.rs b/crates/zakura-network/src/zakura/testkit/cluster.rs index f6c59625f1..dd1d49f059 100644 --- a/crates/zakura-network/src/zakura/testkit/cluster.rs +++ b/crates/zakura-network/src/zakura/testkit/cluster.rs @@ -495,7 +495,7 @@ mod tests { version: 1, frame_cap: CUSTOM_FRAME_CAP_BYTES, capability: CUSTOM_FRAME_CAP_CAPABILITY, - mode: StreamMode::Ordered, + mode: StreamMode::Persistent, }]; #[derive(Debug, Default)] diff --git a/crates/zakura-network/src/zakura/testkit/recorder.rs b/crates/zakura-network/src/zakura/testkit/recorder.rs index 1c018ee76e..bb0ce6fc88 100644 --- a/crates/zakura-network/src/zakura/testkit/recorder.rs +++ b/crates/zakura-network/src/zakura/testkit/recorder.rs @@ -141,7 +141,7 @@ impl Service for InboundRecorder { for stream in self .streams() .iter() - .filter(|stream| matches!(stream.mode, crate::zakura::StreamMode::Ordered)) + .filter(|stream| matches!(stream.mode, crate::zakura::StreamMode::Persistent)) { let Some((session_id, mut recv, _send)) = peer.take_stream_with_session_id(stream.kind) else { diff --git a/crates/zakura-network/src/zakura/transport/io.rs b/crates/zakura-network/src/zakura/transport/io.rs index fba2f75966..32eeff2c8f 100644 --- a/crates/zakura-network/src/zakura/transport/io.rs +++ b/crates/zakura-network/src/zakura/transport/io.rs @@ -4,16 +4,46 @@ //! stream frames. The transport applies each stream's declared `Stream::frame_cap`, //! per-kind message-rate buckets, and idle freshness updates in its stream workers //! before frames reach these handles. +//! +//! A persistent service session closes after every application stream handle is +//! dropped and each member's queued writes finish. A retained receiver or sender +//! clone keeps the session alive. Transport failures still cancel all members. +//! After an application receiver closes, incoming frames still pass transport +//! checks and are discarded until the session retires. + +#![allow(dead_code)] // guarded producers are activated by the next chunks use tokio::sync::mpsc; use super::Frame; -use crate::zakura::regulation::FrameLease; +use std::sync::{Arc, OnceLock}; + +/// Why a persistent stream ended before local cancellation. +#[derive(Copy, Clone, Debug, Eq, PartialEq)] +pub(crate) enum OrderedStreamFailure { + RemoteClose, + WriteTimeout, +} + +/// Preserve the first transport failure before cancelling the service session. +#[derive(Clone, Debug, Default)] +pub(crate) struct OrderedStreamFailureCause(Arc>); + +impl OrderedStreamFailureCause { + pub(crate) fn record(&self, failure: OrderedStreamFailure) { + let _ = self.0.set(failure); + } + + pub(crate) fn get(&self) -> Option { + self.0.get().copied() + } +} /// Receive half for bounded, rate-admitted Zakura frames. #[derive(Debug)] pub struct FramedRecv { receiver: FramedReceiver, + failure_cause: Option, } #[derive(Debug)] @@ -27,12 +57,44 @@ impl FramedRecv { pub fn new(receiver: mpsc::Receiver) -> Self { Self { receiver: FramedReceiver::Plain(receiver), + failure_cause: None, } } fn queued(receiver: mpsc::Receiver) -> Self { Self { receiver: FramedReceiver::Queued(receiver), + failure_cause: None, + } + } + + pub(crate) fn with_failure_cause(mut self, failure_cause: OrderedStreamFailureCause) -> Self { + self.failure_cause = Some(failure_cause); + self + } + + /// Failure of any member, retained through session cancellation for service policy. + pub(crate) fn failure(&self) -> Option { + self.failure_cause.as_ref().and_then(|cause| cause.get()) + } + + /// Stop new ingress while preserving frames already in the bounded queue. + pub(crate) fn close(&mut self) { + match &mut self.receiver { + FramedReceiver::Plain(receiver) => receiver.close(), + FramedReceiver::Queued(receiver) => receiver.close(), + } + } + + /// Receive an already queued frame without waiting for transport progress. + pub(crate) fn try_recv(&mut self) -> Result { + match &mut self.receiver { + FramedReceiver::Plain(receiver) => receiver.try_recv(), + FramedReceiver::Queued(receiver) => loop { + if let Some(frame) = receiver.try_recv()?.receive() { + return Ok(frame); + } + }, } } @@ -41,7 +103,12 @@ impl FramedRecv { match &mut self.receiver { FramedReceiver::Plain(receiver) => receiver.recv().await, FramedReceiver::Queued(receiver) => { - receiver.recv().await.map(|queued| queued.into_parts().0) + while let Some(queued) = receiver.recv().await { + if let Some(frame) = queued.receive() { + return Some(frame); + } + } + None } } } @@ -51,6 +118,7 @@ impl FramedRecv { #[derive(Clone, Debug)] pub struct FramedSend { sender: FramedSender, + session_resources: Option>, } #[derive(Clone, Debug)] @@ -64,15 +132,26 @@ impl FramedSend { pub fn new(sender: mpsc::Sender) -> Self { Self { sender: FramedSender::Plain(sender), + session_resources: None, } } fn queued(sender: mpsc::Sender) -> Self { Self { sender: FramedSender::Queued(sender), + session_resources: None, } } + /// Keep service admission charged while application senders still own the session. + pub(crate) fn with_session_resources( + mut self, + resources: Option>, + ) -> Self { + self.session_resources = resources; + self + } + /// Queue a frame for transport-owned encoding and writing. pub async fn send(&self, frame: Frame) -> Result<(), mpsc::error::SendError> { match &self.sender { @@ -94,29 +173,32 @@ impl FramedSend { } } - /// Reserve a queue slot before attaching an outstanding-byte lease. - /// - /// `make_lease` is called only after the transport owns a queue slot. This - /// prevents accounting from moving to the transport when the queue is full - /// or closed. - #[allow(dead_code)] // consumed by the GetBlocks policy in the stacked PR - pub(crate) fn try_send_leased( - &self, - frame: Frame, - make_lease: impl FnOnce() -> FrameLease, - ) -> Result<(), LeasedSendError> { + /// Reserve queue space before encoding a response or sharing its ownership. + pub(crate) fn try_reserve_guarded(&self) -> Result, GuardedReserveError> { let FramedSender::Queued(sender) = &self.sender else { - return Err(LeasedSendError::Unsupported(frame)); + return Err(GuardedReserveError::Unsupported); }; + sender + .try_reserve() + .map(|permit| GuardedFrameSlot { permit, sender }) + .map_err(|error| match error { + mpsc::error::TrySendError::Full(()) => GuardedReserveError::Full, + mpsc::error::TrySendError::Closed(()) => GuardedReserveError::Closed, + }) + } - match sender.try_reserve() { - Ok(slot) => { - slot.send(QueuedFrame::leased(frame, make_lease())); - Ok(()) - } - Err(mpsc::error::TrySendError::Full(())) => Err(LeasedSendError::Full(frame)), - Err(mpsc::error::TrySendError::Closed(())) => Err(LeasedSendError::Closed(frame)), - } + /// Wait for queue space. Cancellation leaves response ownership with the caller. + pub(crate) async fn reserve_guarded( + &self, + ) -> Result, GuardedReserveError> { + let FramedSender::Queued(sender) = &self.sender else { + return Err(GuardedReserveError::Unsupported); + }; + sender + .reserve() + .await + .map(|permit| GuardedFrameSlot { permit, sender }) + .map_err(|_| GuardedReserveError::Closed) } /// Current free slots in the bounded transport queue. @@ -136,71 +218,131 @@ impl FramedSend { } } -/// Failure to queue a leased frame. +/// One reserved queue slot. Dropping it returns capacity without sending a frame. #[derive(Debug)] -#[allow(dead_code)] // consumed by the GetBlocks policy in the stacked PR -pub(crate) enum LeasedSendError { - /// The bounded transport queue has no free slot. - Full(Frame), - /// The transport worker has closed its receive half. - Closed(Frame), - /// This handle wraps a compatibility channel without lease support. - Unsupported(Frame), +pub(crate) struct GuardedFrameSlot<'a> { + permit: mpsc::Permit<'a, QueuedFrame>, + sender: &'a mpsc::Sender, } -#[allow(dead_code)] // consumed by the GetBlocks policy in the stacked PR -impl LeasedSendError { - /// Recover the frame that was not queued. - pub(crate) fn into_frame(self) -> Frame { - match self { - Self::Full(frame) | Self::Closed(frame) | Self::Unsupported(frame) => frame, - } +impl GuardedFrameSlot<'_> { + /// Transfer a validated frame and its ownership to the reserved queue slot. + pub(crate) fn send(self, frame: Frame, guard: FrameGuard) { + self.permit.send(QueuedFrame::guarded(frame, guard)); } - /// Return whether the queue was temporarily full. - pub(crate) fn is_full(&self) -> bool { - matches!(self, Self::Full(_)) + /// Publish a request whose ownership must be claimed before its first byte. + /// A false result requires explicit settlement after publication unlocks: + /// Tokio can retain a send made through a permit after its receiver drops. + pub(crate) fn send_request(self, frame: Frame, claim: Arc) -> bool { + if self.sender.is_closed() { + return false; + } + self.permit.send(QueuedFrame { + frame, + guard: None, + claim: Some(claim), + }); + !self.sender.is_closed() } +} + +/// Arbitrates an unwritten request against expiry and reset. Dropping a started +/// but unfinished claim must retire the stream session before another write. +pub(crate) trait FrameWriteClaim: std::fmt::Debug + Send + Sync { + /// Atomically claim current ownership, or skip this obsolete frame. + fn try_start(&self) -> bool; + /// Mark the complete frame accepted by the transport write. + fn written(&self); +} + +/// Failure to reserve space for a guarded response. +#[derive(Debug)] +pub(crate) enum GuardedReserveError { + /// The bounded transport queue has no free slot. + Full, + /// The transport worker has closed its receive half. + Closed, + /// This handle wraps a compatibility channel without guard support. + Unsupported, +} - /// Return whether the worker permanently closed the queue. - pub(crate) fn is_closed(&self) -> bool { - matches!(self, Self::Closed(_)) +/// Shared service ownership held until a frame finishes its application write. +/// +/// This is a completion guard, not a byte budget or acknowledgement of delivery. +/// QUIC owns its own bounded send buffers after the write accepts the frame. +#[derive(Clone, Debug)] +pub(crate) struct FrameGuard { + _owner: Arc, +} + +impl FrameGuard { + /// Share an existing work owner without acquiring more capacity. + pub(crate) fn new(owner: Arc) -> Self { + Self { _owner: owner } } } -/// Frame plus optional byte ownership retained through its transport write. +/// Frame plus optional response ownership retained through its transport write. #[derive(Debug)] pub(crate) struct QueuedFrame { frame: Frame, - lease: Option, + guard: Option, + claim: Option>, } impl QueuedFrame { + /// Complete an in-process channel delivery if its claim can still start. + fn receive(self) -> Option { + if let Some(claim) = &self.claim { + if !claim.try_start() { + return None; + } + claim.written(); + } + Some(self.frame) + } + fn plain(frame: Frame) -> Self { - Self { frame, lease: None } + Self { + frame, + guard: None, + claim: None, + } } - #[allow(dead_code)] // consumed through `try_send_leased` in the stacked PR - fn leased(frame: Frame, lease: FrameLease) -> Self { + fn guarded(frame: Frame, guard: FrameGuard) -> Self { Self { frame, - lease: Some(lease), + guard: Some(guard), + claim: None, } } - /// Split the frame from its lease while retaining both in the caller. - pub(crate) fn into_parts(self) -> (Frame, Option) { - (self.frame, self.lease) + /// Split the frame from its guard while retaining both in the caller. + pub(crate) fn into_parts(self) -> (Frame, Option) { + (self.frame, self.guard) } - /// Run the transport write while retaining this frame's lease. - pub(crate) async fn write_with(self, write: F) -> T + /// Run the transport write while retaining this frame's guard. + pub(crate) async fn write_with(self, write: F) -> Result<(), E> where F: FnOnce(Frame) -> Fut, - Fut: std::future::Future, + Fut: std::future::Future>, { - let (frame, _lease) = self.into_parts(); - write(frame).await + let Self { + frame, + guard: _guard, + claim, + } = self; + if claim.as_ref().is_some_and(|claim| !claim.try_start()) { + return Ok(()); + } + write(frame).await?; + if let Some(claim) = &claim { + claim.written(); + } + Ok(()) } } @@ -244,22 +386,30 @@ fn map_queued_try_send_error( #[cfg(test)] mod tests { - use std::sync::{ - atomic::{AtomicBool, Ordering}, - Arc, - }; + use std::sync::Arc; use super::*; - use crate::zakura::regulation::{OutstandingByteBudget, OutstandingByteReservation}; + use crate::zakura::regulation::SlotBudget; fn frame(message_type: u16) -> Frame { Frame { message_type, flags: 0, - payload: vec![u8::try_from(message_type).unwrap_or(u8::MAX)], + payload: vec![u8::try_from(message_type).unwrap_or(u8::MAX); 10], } } + fn guarded_queue() -> (FramedSend, FramedWorkerRecv, SlotBudget) { + let (sender, receiver) = worker_framed_channel(1); + let budget = SlotBudget::new(1).unwrap(); + let reservation = Arc::new(budget.try_reserve().expect("the producer is free")); + sender + .try_reserve_guarded() + .expect("the worker queue has a slot") + .send(frame(1), FrameGuard::new(reservation)); + (sender, receiver, budget) + } + #[tokio::test] async fn public_channel_preserves_order_capacity_and_errors() { let (sender, mut receiver) = framed_channel(2); @@ -305,47 +455,22 @@ mod tests { } #[tokio::test] - async fn queued_frame_holds_lease_until_transport_consumes_it() { - let (sender, mut receiver) = worker_framed_channel(1); - let budget = OutstandingByteBudget::new(10); - let mut reservation = budget - .try_reserve(10) - .expect("the frame fits the budget") - .expect("the budget has capacity"); - - sender - .try_send_leased(frame(1), || { - OutstandingByteReservation::transfer_to_frame([&mut reservation], 10) - .expect("the reservation covers the frame") - }) - .expect("the worker queue has a slot"); - drop(reservation); - assert_eq!(budget.reserved(), 10); + async fn queued_frame_holds_guard_until_transport_consumes_it() { + let (_sender, mut receiver, budget) = guarded_queue(); + assert_eq!(budget.reserved(), 1); let queued = receiver.recv().await.expect("worker receives the frame"); - let (received, lease) = queued.into_parts(); + let (received, guard) = queued.into_parts(); assert_eq!(received, frame(1)); - assert_eq!(budget.reserved(), 10); + assert_eq!(budget.reserved(), 1); - drop(lease); + drop(guard); assert_eq!(budget.reserved(), 0); } #[tokio::test] - async fn queued_frame_holds_lease_while_write_is_pending() { - let (sender, mut receiver) = worker_framed_channel(1); - let budget = OutstandingByteBudget::new(10); - let mut reservation = budget - .try_reserve(10) - .expect("the frame fits the budget") - .expect("the budget has capacity"); - sender - .try_send_leased(frame(1), || { - OutstandingByteReservation::transfer_to_frame([&mut reservation], 10) - .expect("the reservation covers the frame") - }) - .expect("the worker queue has a slot"); - drop(reservation); + async fn queued_frame_holds_guard_while_write_is_pending() { + let (_sender, mut receiver, budget) = guarded_queue(); let queued = receiver.recv().await.expect("worker receives the frame"); let (started_tx, started_rx) = tokio::sync::oneshot::channel(); let (finish_tx, finish_rx) = tokio::sync::oneshot::channel(); @@ -353,39 +478,31 @@ mod tests { let write = tokio::spawn(queued.write_with(move |_frame| async move { let _ = started_tx.send(()); let _ = finish_rx.await; + Ok::<_, std::convert::Infallible>(()) })); started_rx.await.expect("the write reaches its wait point"); - assert_eq!(budget.reserved(), 10); + assert_eq!(budget.reserved(), 1); let _ = finish_tx.send(()); - write.await.expect("the write task should not panic"); + write + .await + .expect("the write task should not panic") + .unwrap(); assert_eq!(budget.reserved(), 0); } #[tokio::test] - async fn cancelling_pending_write_releases_lease() { - let (sender, mut receiver) = worker_framed_channel(1); - let budget = OutstandingByteBudget::new(10); - let mut reservation = budget - .try_reserve(10) - .expect("the frame fits the budget") - .expect("the budget has capacity"); - sender - .try_send_leased(frame(1), || { - OutstandingByteReservation::transfer_to_frame([&mut reservation], 10) - .expect("the reservation covers the frame") - }) - .expect("the worker queue has a slot"); - drop(reservation); + async fn cancelling_pending_write_releases_guard() { + let (_sender, mut receiver, budget) = guarded_queue(); let queued = receiver.recv().await.expect("worker receives the frame"); let (started_tx, started_rx) = tokio::sync::oneshot::channel(); let write = tokio::spawn(queued.write_with(move |_frame| async move { let _ = started_tx.send(()); - std::future::pending::<()>().await; + std::future::pending::>().await })); started_rx.await.expect("the write reaches its wait point"); - assert_eq!(budget.reserved(), 10); + assert_eq!(budget.reserved(), 1); write.abort(); assert!(write @@ -395,63 +512,197 @@ mod tests { assert_eq!(budget.reserved(), 0); } - #[test] - fn failed_leased_send_does_not_create_a_lease() { - let (sender, _receiver) = worker_framed_channel(1); - sender.try_send(frame(1)).expect("the queue slot is free"); - let full_called = Arc::new(AtomicBool::new(false)); - let called_by_factory = full_called.clone(); - - let result = sender.try_send_leased(frame(2), move || { - called_by_factory.store(true, Ordering::SeqCst); - FrameLease::empty_for_test() - }); + #[tokio::test] + async fn waiting_guarded_send_transfers_only_after_capacity_and_holds_through_write() { + let (sender, mut receiver) = worker_framed_channel(1); + sender.try_send(frame(1)).expect("filler fits"); + let budget = SlotBudget::new(1).unwrap(); + let reservation = Arc::new(budget.try_reserve().expect("the producer is free")); + let mut pending = Box::pin(sender.reserve_guarded()); + assert!(futures::poll!(&mut pending).is_pending()); + assert_eq!(Arc::strong_count(&reservation), 1); + assert_eq!(budget.reserved(), 1); + drop(receiver.recv().await.expect("filler queued")); + pending + .await + .expect("the freed slot admits the frame") + .send(frame(2), FrameGuard::new(reservation.clone())); + drop(reservation); + assert_eq!(budget.reserved(), 1); + let queued = receiver.recv().await.expect("guarded frame queued"); + let mut write = Box::pin(queued.write_with(|received| async move { + assert_eq!(received, frame(2)); + std::future::pending::>().await + })); + assert!(futures::poll!(&mut write).is_pending()); + assert_eq!(budget.reserved(), 1); + drop(write); + assert_eq!(budget.reserved(), 0); + } - assert!(matches!(result, Err(LeasedSendError::Full(_)))); - assert!(!full_called.load(Ordering::SeqCst)); + #[tokio::test] + async fn cancelling_or_closing_a_guarded_queue_wait_keeps_the_callers_reservation() { + let (sender, mut receiver) = worker_framed_channel(1); + sender.try_send(frame(1)).expect("filler fits"); + let budget = SlotBudget::new(1).unwrap(); + let reservation = Arc::new(budget.try_reserve().expect("the producer is free")); + let mut pending = Box::pin(sender.reserve_guarded()); + assert!(futures::poll!(&mut pending).is_pending()); + drop(pending); + drop(receiver.recv().await.expect("filler queued")); + assert_eq!(Arc::strong_count(&reservation), 1); + assert_eq!(budget.reserved(), 1); + drop(receiver); + assert!(matches!( + sender.reserve_guarded().await, + Err(GuardedReserveError::Closed) + )); + assert_eq!(Arc::strong_count(&reservation), 1); + assert_eq!(budget.reserved(), 1); + drop(reservation); + assert_eq!(budget.reserved(), 0); + } + #[test] + fn guarded_reservation_reports_full_closed_and_unsupported_queues() { let (sender, receiver) = worker_framed_channel(1); + sender.try_send(frame(1)).unwrap(); + assert!(matches!( + sender.try_reserve_guarded(), + Err(GuardedReserveError::Full) + )); drop(receiver); - let closed_called = Arc::new(AtomicBool::new(false)); - let called_by_factory = closed_called.clone(); - let result = sender.try_send_leased(frame(3), move || { - called_by_factory.store(true, Ordering::SeqCst); - FrameLease::empty_for_test() - }); - assert!(matches!(result, Err(LeasedSendError::Closed(_)))); - assert!(!closed_called.load(Ordering::SeqCst)); + assert!(matches!( + sender.try_reserve_guarded(), + Err(GuardedReserveError::Closed) + )); let (raw_sender, _raw_receiver) = mpsc::channel(1); let sender = FramedSend::new(raw_sender); - let unsupported_called = Arc::new(AtomicBool::new(false)); - let called_by_factory = unsupported_called.clone(); - let result = sender.try_send_leased(frame(4), move || { - called_by_factory.store(true, Ordering::SeqCst); - FrameLease::empty_for_test() - }); - assert!(matches!(result, Err(LeasedSendError::Unsupported(_)))); - assert!(!unsupported_called.load(Ordering::SeqCst)); + assert!(matches!( + sender.try_reserve_guarded(), + Err(GuardedReserveError::Unsupported) + )); } #[test] - fn dropping_worker_queue_releases_queued_lease() { - let (sender, receiver) = worker_framed_channel(1); - let budget = OutstandingByteBudget::new(10); - let mut reservation = budget - .try_reserve(10) - .expect("the frame fits the budget") - .expect("the budget has capacity"); - sender - .try_send_leased(frame(1), || { - OutstandingByteReservation::transfer_to_frame([&mut reservation], 10) - .expect("the reservation covers the frame") - }) - .expect("the worker queue has a slot"); - drop(reservation); - assert_eq!(budget.reserved(), 10); + fn dropping_reserved_slot_returns_queue_capacity() { + let (sender, _receiver) = worker_framed_channel(1); + let slot = sender.try_reserve_guarded().unwrap(); + assert_eq!(sender.capacity(), 0); + drop(slot); + assert_eq!(sender.capacity(), 1); + } + + #[test] + fn dropping_worker_queue_releases_queued_guard() { + let (_sender, receiver, budget) = guarded_queue(); + assert_eq!(budget.reserved(), 1); drop(receiver); assert_eq!(budget.reserved(), 0); } + #[tokio::test] + async fn quic_backpressure_holds_producer_and_preserves_another_stream() { + use crate::zakura::testkit::LocalEndpointFactory; + use iroh::{ + endpoint::{Connection, QuicTransportConfig, VarInt}, + protocol::{AcceptError, ProtocolHandler, Router}, + }; + use std::time::Duration; + + #[derive(Debug)] + struct AcceptConnection(mpsc::Sender); + impl ProtocolHandler for AcceptConnection { + async fn accept(&self, connection: Connection) -> Result<(), AcceptError> { + let _ = self.0.send(connection).await; + Ok(()) + } + } + + const ALPN: &[u8] = b"/zakura/test/producer-backpressure"; + // Scale down the windows so a single bounded frame reaches flow control. + let transport_config = || { + QuicTransportConfig::builder() + .max_remote_nat_traversal_addresses(0) + .stream_receive_window(VarInt::from_u32(16 * 1024)) + .receive_window(VarInt::from_u32(128 * 1024)) + .send_window(128 * 1024) + .build() + }; + let server = LocalEndpointFactory::with_transport_config(transport_config()) + .endpoint(92_001) + .await + .unwrap(); + let client = LocalEndpointFactory::with_transport_config(transport_config()) + .endpoint(92_002) + .await + .unwrap(); + let (accepted, mut incoming) = mpsc::channel(1); + let router = Router::builder(server) + .accept(ALPN, AcceptConnection(accepted)) + .spawn(); + let address = LocalEndpointFactory::node_addr(router.endpoint()).await; + let connection = client.connect(address, ALPN).await.unwrap(); + let remote = tokio::time::timeout(Duration::from_secs(5), incoming.recv()) + .await + .unwrap() + .unwrap(); + let (mut send, _recv) = connection.open_bi().await.unwrap(); + let producer = SlotBudget::new(1).unwrap(); + let owner = Arc::new(producer.try_reserve().unwrap()); + let (queue, mut writer) = worker_framed_channel(1); + queue.try_reserve_guarded().unwrap().send( + Frame { + message_type: 1, + flags: 0, + payload: vec![0; 2_000_001], + }, + FrameGuard::new(owner.clone()), + ); + drop(owner); + let queued = writer.recv().await.unwrap(); + let mut write = tokio::spawn(queued.write_with(move |frame| async move { + send.write_all(&frame.payload).await.unwrap(); + send.finish().unwrap(); + Ok::<_, std::convert::Infallible>(()) + })); + let (_remote_send, mut slow_read) = + tokio::time::timeout(Duration::from_secs(5), remote.accept_bi()) + .await + .unwrap() + .unwrap(); + assert!(tokio::time::timeout(Duration::from_millis(100), &mut write) + .await + .is_err()); + assert!( + producer.try_reserve().is_none(), + "a pending QUIC write retains the producer" + ); + + let (mut other_send, _other_recv) = connection.open_bi().await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), async { + other_send.write_all(b"progress").await.unwrap(); + other_send.finish().unwrap(); + let (_send, mut recv) = remote.accept_bi().await.unwrap(); + assert_eq!(recv.read_to_end(8).await.unwrap(), b"progress"); + }) + .await + .expect("the blocked stream does not consume all connection credit"); + assert!(producer.try_reserve().is_none()); + tokio::time::timeout(Duration::from_secs(5), async { + assert_eq!( + slow_read.read_to_end(2_000_001).await.unwrap().len(), + 2_000_001 + ); + write.await.unwrap().unwrap(); + }) + .await + .expect("draining the peer resumes the write"); + assert!(producer.try_reserve().is_some()); + connection.close(0u32.into(), b"done"); + client.close().await; + router.shutdown().await.unwrap(); + } } diff --git a/crates/zakura-network/src/zakura/transport/mod.rs b/crates/zakura-network/src/zakura/transport/mod.rs index bd787c4ea7..6fffc0cc8d 100644 --- a/crates/zakura-network/src/zakura/transport/mod.rs +++ b/crates/zakura-network/src/zakura/transport/mod.rs @@ -20,16 +20,21 @@ pub use frame::{Frame, StreamPrelude, ZakuraTrace}; #[allow(unused_imports)] pub(crate) use guard::{Admit, ByteBudget, PeerMeters, SessionGuard}; pub use io::{framed_channel, FramedRecv, FramedSend}; -#[allow(unused_imports)] // used by the first message policy in the stacked PR -pub(crate) use io::{worker_framed_channel, FramedWorkerRecv, LeasedSendError, QueuedFrame}; +#[allow(unused_imports)] // guarded producers are activated by the next chunks +pub(crate) use io::{ + worker_framed_channel, FrameGuard, FrameWriteClaim, FramedWorkerRecv, GuardedReserveError, + OrderedStreamFailure, OrderedStreamFailureCause, QueuedFrame, +}; pub(crate) use pipe::{ handle_pipe_exit, spawn_supervised_peer_task, spawn_supervised_pipe, CloseCause, Edge, Flow, Node, NodeKind, Pipe, PipeCx, PipeShape, }; +pub(crate) use registry::SessionLayout; pub use registry::{RegistryError, ServiceRegistry}; pub(crate) use service::ServiceStream; pub use service::{ - BoxRunFuture, OrderedSessionDemand, OrderedStreamOpening, OrderedStreamPolicy, Peer, - RequestResponseService, Service, Sink, SinkReject, Source, Stream, StreamMode, + BoxRunFuture, Peer, RequestResponseService, Service, SessionDemand, SessionFull, + SessionOpening, SessionPolicy, SessionResources, Sink, SinkReject, Source, Stream, StreamMode, + StreamWritePolicy, }; pub use session::{OrderedSendError, PeerStreamSession}; diff --git a/crates/zakura-network/src/zakura/transport/registry.rs b/crates/zakura-network/src/zakura/transport/registry.rs index 48cb07ff63..7c6e4c6b6f 100644 --- a/crates/zakura-network/src/zakura/transport/registry.rs +++ b/crates/zakura-network/src/zakura/transport/registry.rs @@ -8,13 +8,22 @@ use std::{ use thiserror::Error; use super::{ - Frame, OrderedSessionDemand, OrderedStreamPolicy, Peer, Service, SinkReject, Stream, StreamMode, + Frame, Peer, Service, SessionDemand, SessionPolicy, SinkReject, Stream, StreamMode, + StreamWritePolicy, }; use crate::zakura::{ServicePeerDirection, ZakuraConnId, ZakuraPeerId}; /// Errors returned while building a [`ServiceRegistry`]. #[derive(Debug, Error)] pub enum RegistryError { + /// A service declared incompatible persistent session layouts. + #[error("service {service} declared an invalid service session for kind {kind}")] + InvalidSessionLayout { + /// Service declaring the session. + service: &'static str, + /// Stream with an inconsistent session declaration. + kind: u16, + }, /// Two services declared the same stream kind. #[error( "duplicate Zakura stream kind {kind} declared by {first_service} and {second_service}" @@ -60,6 +69,23 @@ pub enum RegistryError { }, } +/// A validated set of persistent streams admitted as one service session. +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) struct SessionLayout { + pub(crate) streams: Arc<[Stream]>, +} + +impl SessionLayout { + /// The lowest stream kind supplies the session identity and layout version. + pub(crate) fn primary(&self) -> Stream { + self.streams[0] + } + + pub(crate) fn is_multi_stream(&self) -> bool { + self.streams.len() > 1 + } +} + /// Registry of Zakura protocol services. #[derive(Clone, Debug, Default)] pub struct ServiceRegistry { @@ -67,6 +93,7 @@ pub struct ServiceRegistry { by_kind: HashMap, by_capability: HashMap>, supported_capabilities: u64, + session_layouts: HashMap<(u16, u16), SessionLayout>, } impl ServiceRegistry { @@ -75,6 +102,7 @@ impl ServiceRegistry { let mut by_kind: HashMap = HashMap::new(); let mut by_capability: HashMap> = HashMap::new(); let mut supported_capabilities = 0; + let mut session_layouts = HashMap::new(); for (index, service) in services.iter().enumerate() { let mut service_capabilities = HashSet::new(); @@ -116,6 +144,37 @@ impl ServiceRegistry { service_capabilities.insert(stream.capability); } + let mut layouts: HashMap> = HashMap::new(); + for stream in service + .streams() + .iter() + .filter(|s| s.mode == StreamMode::Persistent) + { + layouts.entry(stream.capability).or_default().push(*stream); + } + let mut primary_kind = None; + for mut streams in layouts.into_values() { + streams.sort_unstable_by_key(|stream| stream.kind); + let primary = streams[0]; + let invalid = primary_kind.is_some_and(|kind| kind != primary.kind) + || streams + .windows(2) + .any(|roles| roles[0].kind == roles[1].kind); + if invalid { + return Err(RegistryError::InvalidSessionLayout { + service: service.name(), + kind: primary.kind, + }); + } + primary_kind = Some(primary.kind); + let layout = SessionLayout { + streams: streams.into(), + }; + for stream in layout.streams.iter() { + session_layouts.insert((stream.kind, stream.version), layout.clone()); + } + } + for capability in service_capabilities { by_capability.entry(capability).or_default().push(index); } @@ -126,6 +185,7 @@ impl ServiceRegistry { by_kind, by_capability, supported_capabilities, + session_layouts, }) } @@ -141,6 +201,22 @@ impl ServiceRegistry { .map(|index| Arc::clone(&self.services[*index])) } + /// Local message limits supplied by the service that owns this stream. + pub(crate) fn message_payload_limits(&self, stream: Stream) -> &'static [(u16, usize)] { + self.service_for_kind(stream.kind) + .map(|service| service.message_payload_limits(stream)) + .unwrap_or(&[]) + } + + pub(crate) fn stream_queue_depths(&self, stream: Stream) -> Option<(usize, usize)> { + self.service_for_kind(stream.kind)? + .stream_queue_depths(stream) + } + + pub(crate) fn message_types(&self, stream: Stream) -> Option<&'static [u16]> { + self.service_for_kind(stream.kind)?.message_types(stream) + } + /// Lookup the declared stream for `kind`. pub fn stream_for_kind(&self, kind: u16) -> Option { let service = self.service_for_kind(kind)?; @@ -215,28 +291,50 @@ impl ServiceRegistry { self.supported_capabilities } - /// Ordered streams negotiated with a peer, in registry service order. - pub fn ordered_streams_for_negotiated(&self, negotiated: u64) -> Vec { + /// Return the complete persistent layout containing this exact stream. + pub(crate) fn session_layout(&self, stream: Stream) -> Option { + self.session_layouts + .get(&(stream.kind, stream.version)) + .filter(|layout| layout.streams.contains(&stream)) + .cloned() + } + + pub(crate) fn stream_write_policy(&self, stream: Stream) -> StreamWritePolicy { + self.service_for_kind(stream.kind) + .expect("a registered stream has an owning service") + .stream_write_policy(stream) + } + + fn selected_session_streams(&self, service: &dyn Service, negotiated: u64) -> Vec { + service + .streams() + .iter() + .filter(|stream| { + stream.mode == StreamMode::Persistent && negotiated & stream.capability != 0 + }) + .filter_map(|stream| self.session_layout(*stream)) + .max_by_key(|layout| layout.primary().version) + .map_or_else(Vec::new, |layout| layout.streams.to_vec()) + } + + /// Persistent streams negotiated with a peer, in registry service order. + pub fn persistent_streams_for_negotiated(&self, negotiated: u64) -> Vec { let mut streams = Vec::new(); for service in self.services_for_negotiated(negotiated) { - streams.extend(selected_streams( - service.streams(), - negotiated, - StreamMode::Ordered, - )); + streams.extend(self.selected_session_streams(service.as_ref(), negotiated)); } streams } - /// Ordered streams that should be lazily escalated for this peer now. + /// Persistent streams that should be lazily escalated for this peer now. /// /// The connection loop applies its per-kind opening policy to each returned /// stream. This demand check narrows the negotiated capabilities to services /// that currently have local interest and room; the owning reactor still /// makes the final admission decision after the typed session arrives. - pub fn ordered_streams_for_escalation( + pub fn persistent_streams_for_escalation( &self, negotiated: u64, peer_id: &ZakuraPeerId, @@ -249,18 +347,14 @@ impl ServiceRegistry { continue; } - streams.extend(selected_streams( - service.streams(), - negotiated, - StreamMode::Ordered, - )); + streams.extend(self.selected_session_streams(service.as_ref(), negotiated)); } streams } /// Return true when the service owning `kind` still wants this peer. - pub fn wants_ordered_stream( + pub fn wants_session( &self, kind: u16, negotiated: u64, @@ -275,26 +369,41 @@ impl ServiceRegistry { } /// Return the owning service's static ordered-stream policy. - pub fn ordered_stream_policy(&self, kind: u16) -> OrderedStreamPolicy { + pub fn session_policy(&self, kind: u16) -> SessionPolicy { self.service_for_kind(kind) - .map(|service| service.ordered_stream_policy(kind)) + .map(|service| service.session_policy()) .unwrap_or_default() } /// Return the owning service's current demand for an absent ordered session. - pub fn ordered_session_demand( + pub fn session_demand( &self, kind: u16, conn_id: ZakuraConnId, negotiated: u64, peer_id: &ZakuraPeerId, direction: ServicePeerDirection, - ) -> OrderedSessionDemand { + ) -> SessionDemand { let Some(service) = self.service_for_kind(kind) else { - return OrderedSessionDemand::Retire; + return SessionDemand::Retire; }; - service.ordered_session_demand(conn_id, peer_id, negotiated, direction) + service.session_demand(conn_id, peer_id, negotiated, direction) + } + + /// Recheck demand after a complete session has reserved its service capacity. + pub(crate) fn reserved_session_demand( + &self, + kind: u16, + conn_id: ZakuraConnId, + negotiated: u64, + peer_id: &ZakuraPeerId, + direction: ServicePeerDirection, + ) -> SessionDemand { + let Some(service) = self.service_for_kind(kind) else { + return SessionDemand::Retire; + }; + service.reserved_session_demand(conn_id, peer_id, negotiated, direction) } /// Request/response streams negotiated with a peer, in registry service order. @@ -303,7 +412,7 @@ impl ServiceRegistry { for service in self.services_for_negotiated(negotiated) { streams.extend(selected_streams( - service.streams(), + service.as_ref(), negotiated, StreamMode::RequestResponse, )); @@ -470,9 +579,10 @@ impl ServiceRegistry { /// Each capability bit declares one version alternative. /// Select the highest matching version before opening the prelude. /// The stream selection scopes decoding and preserves existing streams for older peers. -fn selected_streams(streams: &[Stream], negotiated: u64, mode: StreamMode) -> Vec { +fn selected_streams(service: &dyn Service, negotiated: u64, mode: StreamMode) -> Vec { let mut selected = Vec::::new(); - for stream in streams + for stream in service + .streams() .iter() .copied() .filter(|stream| stream.mode == mode && negotiated & stream.capability != 0) @@ -601,7 +711,7 @@ mod tests { version: 1, frame_cap: 1024, capability, - mode: StreamMode::Ordered, + mode: StreamMode::Persistent, } } @@ -627,15 +737,48 @@ mod tests { assert!(registry.is_supported_stream(5, 7)); assert!(registry.is_supported_stream(5, 8)); assert_eq!( - registry.ordered_streams_for_negotiated(0b0001), + registry.persistent_streams_for_negotiated(0b0001), vec![versioned_stream(5, 7, 0b0001)] ); assert_eq!( - registry.ordered_streams_for_negotiated(0b0011), + registry.persistent_streams_for_negotiated(0b0011), vec![versioned_stream(5, 8, 0b0010)] ); } + #[test] + fn session_versions_are_selected_as_complete_layouts() { + let legacy = versioned_stream(6, 1, 1); + let older = [versioned_stream(6, 2, 2), versioned_stream(7, 4, 2)]; + let newer = [ + versioned_stream(6, 3, 4), + versioned_stream(7, 1, 4), + versioned_stream(8, 1, 4), + ]; + let service = TestService::new( + "session", + [vec![legacy], older.to_vec(), newer.to_vec()].concat(), + ); + let registry = ServiceRegistry::new(vec![service]).unwrap(); + assert_eq!(registry.persistent_streams_for_negotiated(1), vec![legacy]); + assert_eq!(registry.persistent_streams_for_negotiated(3), older); + // Per-kind selection would incorrectly choose stream 7 version 4. + assert_eq!(registry.persistent_streams_for_negotiated(7), newer); + } + + #[test] + fn session_layouts_require_one_version_per_kind_and_a_stable_primary_kind() { + for streams in [ + vec![versioned_stream(6, 1, 1), versioned_stream(6, 2, 1)], + vec![versioned_stream(6, 1, 1), versioned_stream(7, 1, 2)], + ] { + assert!(matches!( + ServiceRegistry::new(vec![TestService::new("invalid", streams)]), + Err(RegistryError::InvalidSessionLayout { .. }) + )); + } + } + #[test] fn duplicate_kind_and_version_in_one_service_is_rejected() { let header = TestService::new( @@ -658,7 +801,10 @@ mod tests { #[test] fn registry_builds_kind_and_capability_lookups() { - let header = TestService::new("header", vec![stream(5, 0b0001), stream(6, 0b0010)]); + let header = TestService::new( + "header", + vec![stream(5, 0b0001), versioned_stream(5, 2, 0b0010)], + ); let gossip = TestService::new("gossip", vec![stream(2, 0b0100)]); let registry = ServiceRegistry::new(vec![header.clone(), gossip.clone()]) @@ -739,7 +885,10 @@ mod tests { #[test] fn supported_capabilities_are_or_of_declared_streams() { - let header = TestService::new("header", vec![stream(5, 0b0001), stream(6, 0b0010)]); + let header = TestService::new( + "header", + vec![stream(5, 0b0001), versioned_stream(5, 2, 0b0010)], + ); let gossip = TestService::new("gossip", vec![stream(2, 0b0100)]); let registry = ServiceRegistry::new(vec![header, gossip]) @@ -750,7 +899,10 @@ mod tests { #[test] fn services_for_negotiated_matches_any_bit_once_in_registration_order() { - let header = TestService::new("header", vec![stream(5, 0b0001), stream(6, 0b0010)]); + let header = TestService::new( + "header", + vec![stream(5, 0b0001), versioned_stream(5, 2, 0b0010)], + ); let gossip = TestService::new("gossip", vec![stream(2, 0b0100)]); let discovery = TestService::new("discovery", vec![stream(4, 0b1000)]); @@ -773,7 +925,7 @@ mod tests { "multi-capability", vec![ stream(5, 0b0001), - stream(6, 0b0010), + versioned_stream(5, 2, 0b0010), request_response_one, request_response_two, ], @@ -783,12 +935,12 @@ mod tests { let peer = ZakuraPeerId::new(vec![8; 32]).expect("32-byte test peer id is valid"); let ordered_kinds: Vec<_> = registry - .ordered_streams_for_negotiated(0b0001) + .persistent_streams_for_negotiated(0b0001) .iter() .map(|stream| stream.kind) .collect(); let escalated_kinds: Vec<_> = registry - .ordered_streams_for_escalation(0b0001, &peer, ServicePeerDirection::Outbound) + .persistent_streams_for_escalation(0b0001, &peer, ServicePeerDirection::Outbound) .iter() .map(|stream| stream.kind) .collect(); @@ -911,7 +1063,7 @@ mod tests { } #[test] - fn ordered_streams_for_escalation_filters_services_without_demand() { + fn persistent_streams_for_escalation_filters_services_without_demand() { let header = TestService::new("header", vec![stream(5, 0b0001)]); let discovery = TestService::new("discovery", vec![stream(4, 0b0010)]); let registry = ServiceRegistry::new(vec![header.clone(), discovery.clone()]) @@ -920,8 +1072,11 @@ mod tests { header.set_wants(false); - let streams = - registry.ordered_streams_for_escalation(0b0011, &peer, ServicePeerDirection::Outbound); + let streams = registry.persistent_streams_for_escalation( + 0b0011, + &peer, + ServicePeerDirection::Outbound, + ); let stream_kinds: Vec<_> = streams.iter().map(|stream| stream.kind).collect(); assert_eq!(stream_kinds, [4]); diff --git a/crates/zakura-network/src/zakura/transport/service.rs b/crates/zakura-network/src/zakura/transport/service.rs index db7dfd7935..f828a642a7 100644 --- a/crates/zakura-network/src/zakura/transport/service.rs +++ b/crates/zakura-network/src/zakura/transport/service.rs @@ -1,6 +1,13 @@ //! Zakura protocol service trait surface. -use std::{collections::HashMap, fmt, future::Future, net::IpAddr, pin::Pin, time::Instant}; +use std::{ + collections::HashMap, + fmt, + future::Future, + net::IpAddr, + pin::Pin, + time::{Duration, Instant}, +}; use thiserror::Error; use tokio_util::sync::CancellationToken; @@ -20,41 +27,41 @@ pub type BoxRunFuture<'a, T> = Pin + Send + 'a>>; #[derive(Copy, Clone, Debug, Eq, PartialEq)] pub enum StreamMode { /// A long-lived ordered stream between connected peers. - Ordered, + Persistent, /// A short-lived request/response stream opened per request. RequestResponse, } -/// Which endpoint may proactively open an ordered service stream. +/// Which endpoint may proactively open a service session. #[derive(Copy, Clone, Debug, Eq, PartialEq)] -pub enum OrderedStreamOpening { +pub enum SessionOpening { /// Only the endpoint that initiated the authenticated connection opens the stream. InitiatorOnly, /// Either endpoint may open the stream; simultaneous opens use the transport tiebreak. EitherSide, } -/// Static transport policy for one ordered service stream. +/// Static transport policy for one persistent service session. #[derive(Copy, Clone, Debug, Eq, PartialEq)] -pub struct OrderedStreamPolicy { +pub struct SessionPolicy { /// Which endpoint may proactively open the stream. - pub opening: OrderedStreamOpening, + pub opening: SessionOpening, /// Whether a locally ended session may be re-admitted on the same connection. pub reopen: bool, } -impl Default for OrderedStreamPolicy { +impl Default for SessionPolicy { fn default() -> Self { Self { - opening: OrderedStreamOpening::InitiatorOnly, + opening: SessionOpening::InitiatorOnly, reopen: false, } } } -/// A service's current decision for an absent ordered session. -pub enum OrderedSessionDemand { - /// Open and admit the ordered stream now. +/// A service's current decision for an absent service session. +pub enum SessionDemand { + /// Open and admit the complete session now. OpenNow, /// Re-check demand at this instant. RetryAt(Instant), @@ -64,7 +71,7 @@ pub enum OrderedSessionDemand { Retire, } -impl fmt::Debug for OrderedSessionDemand { +impl fmt::Debug for SessionDemand { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::OpenNow => formatter.write_str("OpenNow"), @@ -90,7 +97,29 @@ pub struct Stream { pub mode: StreamMode, } -/// Transport state for one ordered service stream. +/// A persistent stream's write deadline within its service session. +#[derive(Copy, Clone, Debug, Eq, PartialEq)] +pub enum StreamWritePolicy { + /// Retire the session if a complete frame cannot be written within this time. + Timeout(Duration), + /// Let the service's progress policy or session cancellation end the wait. + UntilCancelled, +} + +/// A service slot held from session setup through the last transport and +/// application sender owner. The service can release its setup allowance once +/// all members are ready, while retaining its session allowance through teardown. +pub trait SessionResources: fmt::Debug + Send + Sync { + /// All required streams have completed setup. Called once per session. + fn admitted(&self); +} + +/// The service has no capacity for another establishing or retiring session. +#[derive(Debug, Error)] +#[error("service session capacity is full")] +pub struct SessionFull; + +/// Transport state for one persistent stream within a service session. #[derive(Debug)] pub(crate) struct ServiceStream { pub(crate) session_id: u64, @@ -285,14 +314,14 @@ impl Peer { } } - /// Take ownership of a stream pair for `kind`. + /// Take ownership of a stream's receive and send handles for `kind`. pub fn take_stream(&mut self, kind: u16) -> Option<(FramedRecv, FramedSend)> { self.streams .remove(&kind) .map(|stream| (stream.recv, stream.send)) } - /// Take ownership of a stream pair and its owning ordered-stream generation. + /// Take ownership of a stream's receive and send handles and its owning session identity. pub fn take_stream_with_session_id( &mut self, kind: u16, @@ -302,7 +331,7 @@ impl Peer { .map(|stream| (stream.session_id, stream.recv, stream.send)) } - /// Take ownership of a stream pair, its version, and ordered-stream generation. + /// Take ownership of a stream's receive and send handles, its version, and session identity. pub fn take_versioned_stream_with_session_id( &mut self, kind: u16, @@ -364,36 +393,98 @@ pub trait Service: fmt::Debug + Send + Sync + 'static { /// Stable service name for logs and diagnostics. fn name(&self) -> &'static str; - /// Streams this service owns. + /// Stream types this service owns. + /// + /// Persistent streams with the same capability form one complete session + /// layout. The transport admits every member together and retires the + /// session when any member ends. Request/response streams open per request + /// and do not participate in session setup. + /// + /// Alternative layouts use distinct capabilities and the same lowest stream + /// kind. That stream's version ranks complete layouts during negotiation. + /// Advance its version whenever the session layout changes. fn streams(&self) -> &[Stream]; - /// Return the transport-owned opening and re-admission policy for `kind`. + /// Payload size limits for this stream, as `(message_type, maximum_bytes)` pairs. + /// + /// The reader checks these limits before allocating a payload. Limits exclude + /// the frame header and may only tighten the stream's existing cap. Unlisted + /// message types keep that cap; message validity is checked by the codec. + fn message_payload_limits(&self, _stream: Stream) -> &'static [(u16, usize)] { + &[] + } + + /// Optional message types accepted on this role. The transport rejects an + /// unlisted type from its header, before allocating or reading its payload. + fn message_types(&self, _stream: Stream) -> Option<&'static [u16]> { + None + } + + /// Optional per-stream inbound and outbound application queue limits. + /// The transport also applies its connection-wide inbound queue allowance. + fn stream_queue_depths(&self, _stream: Stream) -> Option<(usize, usize)> { + None + } + + /// Reserve service capacity before starting a persistent session. The returned + /// owner lives through incomplete setup and all workers' eventual teardown. + fn reserve_session( + &self, + _direction: ServicePeerDirection, + ) -> Result>, SessionFull> { + Ok(None) + } + + /// Control writes independently for each persistent stream. + /// + /// Services that allow indefinite writes must enforce their own bounded + /// progress policy. Cancelling the session always interrupts a pending write. + fn stream_write_policy(&self, _stream: Stream) -> StreamWritePolicy { + StreamWritePolicy::Timeout(Duration::from_secs(10)) + } + + /// Return the opening and re-admission policy for the whole service session. /// /// The default preserves the legacy one-shot initiator-opens behavior. - fn ordered_stream_policy(&self, _kind: u16) -> OrderedStreamPolicy { - OrderedStreamPolicy::default() + fn session_policy(&self) -> SessionPolicy { + SessionPolicy::default() } - /// Return this service's current demand for an absent ordered session. + /// Return this service's current demand for an absent service session. /// - /// Services that opt into [`OrderedStreamPolicy::reopen`] should override + /// Services that opt into [`SessionPolicy::reopen`] should override /// this method so local cooldowns, capacity, and usefulness remain - /// reactor-owned. [`OrderedSessionDemand::WaitForChange`] avoids periodic + /// reactor-owned. [`SessionDemand::WaitForChange`] avoids periodic /// transport polling while the service is full or has no useful work. - fn ordered_session_demand( + fn session_demand( &self, _conn_id: ZakuraConnId, peer: &ZakuraPeerId, negotiated: u64, direction: ServicePeerDirection, - ) -> OrderedSessionDemand { + ) -> SessionDemand { if self.wants_peer(peer, negotiated, direction) { - OrderedSessionDemand::OpenNow + SessionDemand::OpenNow } else { - OrderedSessionDemand::Retire + SessionDemand::Retire } } + /// Recheck demand for a complete session that already owns its setup reservation. + /// + /// Services with session reservations must retain cooldown and usefulness + /// checks here without requiring capacity for a second reservation. The + /// default preserves ordinary demand checks for services without reservations. + fn reserved_session_demand( + &self, + conn_id: ZakuraConnId, + peer: &ZakuraPeerId, + negotiated: u64, + direction: ServicePeerDirection, + ) -> SessionDemand { + self.session_demand(conn_id, peer, negotiated, direction) + } + /// Return whether this service currently wants a new session for `peer`. /// /// This is a cheap, advisory demand check used by the transport before diff --git a/crates/zakurad/tests/common/configs/v1.4.0.toml b/crates/zakurad/tests/common/configs/v1.4.0.toml deleted file mode 100644 index 2a6a188d1f..0000000000 --- a/crates/zakurad/tests/common/configs/v1.4.0.toml +++ /dev/null @@ -1,151 +0,0 @@ -# Default configuration for zakurad. -# -# This file can be used as a skeleton for custom configs. -# -# Unspecified fields use default values. Optional fields are Some(field) if the -# field is present and None if it is absent. -# -# This file is generated as an example using zakurad's current defaults. -# You should set only the config options you want to keep, and delete the rest. -# Only a subset of fields are present in the skeleton, since optional values -# whose default is None are omitted. -# -# The config format (including a complete list of sections and fields) is -# documented here: -# https://docs.rs/zakura/latest/zakurad/config/struct.ZakuradConfig.html -# -# CONFIGURATION SOURCES (in order of precedence, highest to lowest): -# -# 1. Environment variables with ZAKURA_ prefix (highest precedence) -# - Format: ZAKURA_SECTION__KEY (double underscore for nested keys) -# - Examples: -# - ZAKURA_NETWORK__NETWORK=Testnet -# - ZAKURA_RPC__LISTEN_ADDR=127.0.0.1:8232 -# - ZAKURA_STATE__CACHE_DIR=/path/to/cache -# - ZAKURA_TRACING__FILTER=debug -# - ZAKURA_METRICS__ENDPOINT_ADDR=0.0.0.0:9999 -# -# 2. Environment variables with deprecated ZEBRA_ prefix -# -# 3. Configuration file (TOML format) -# - At the path specified via -c flag, e.g. `zakurad -c myconfig.toml start`, or -# - At the default path in the user's preference directory (platform-dependent, see below) -# -# 4. Hard-coded defaults (lowest precedence) -# -# The user's preference directory and the default path to the `zakurad` config are platform dependent, -# based on `dirs::preference_dir`, see https://docs.rs/dirs/latest/dirs/fn.preference_dir.html : -# -# | Platform | Value | Example | -# | -------- | ------------------------------------- | ---------------------------------------------- | -# | Linux | `$XDG_CONFIG_HOME` or `$HOME/.config` | `/home/alice/.config/zakura.toml` | -# | macOS | `$HOME/Library/Preferences` | `/Users/Alice/Library/Preferences/zakura.toml` | -# | Windows | `{FOLDERID_RoamingAppData}` | `C:\Users\Alice\AppData\Local\zakura.toml` | - -[consensus] -checkpoint_sync = true - -[health] -enforce_on_test_networks = false -min_connected_peers = 1 -ready_max_blocks_behind = 2 -ready_max_tip_age = "5m" - -[mempool] -eviction_memory_time = "1h" -max_datacarrier_bytes = 83 -max_transaction_bytes = 250000 -tx_cost_limit = 80000000 - -[metrics] - -[mining] -internal_miner = false -optimistic_block_inventory = true - -[network] -cache_dir = true -crawl_new_peer_interval = "1m 1s" -expose_peer_addresses = false -identity_dir = "identity_dir" -initial_mainnet_peers = [ - "dnsseed.str4d.xyz:8233", - "dnsseed.z.cash:8233", - "mainnet.seeder.shieldedinfra.net:8233", - "mainnet.seeder.zfnd.org:8233", -] -initial_testnet_peers = [ - "dnsseed.testnet.z.cash:18233", - "testnet.seeder.zfnd.org:18233", -] -listen_addr = "[::]:8233" -max_connections_per_ip = 1 -network = "Mainnet" -# The peer-to-peer stack to run: -# - "legacy": the legacy TCP Zcash P2P stack only. -# - "zakura": the experimental native Zakura P2P v2 stack only. -# - "dual": both stacks, enabling experimental v2 with legacy fallback. -# - "default": Zakura's default for this network, which can change between -# releases. Currently "legacy" on Mainnet, and "dual" everywhere else. -p2p_stack = "default" -peerset_initial_target_size = 100 - -[network.zakura] -bootstrap_peers = [ - "1398f62c6d1a457c51ba6a4b5f3dbd2f69fca93216218dc8997e416bd17d93ca@165.22.54.66:8234", - "fd1724385aa0c75b64fb78cd602fa1d991fdebf76b13c58ed702eac835e9f618@104.131.184.123:8234", - "9ec67ad6834bc2ca0d659c240e042d3446c37cabcc092b527d459c87d938b4a4@159.65.183.89:8234", - "bd3dc5d2a3d44c6bf90e364bf446231dbf9737e38a562ccf9e91ea631ea59b22@143.244.184.176:8234", - "14ab98fa0c4b07d40119e1dbc9f3c36d20c8f226ae5ba4216218a2034f148e57@159.203.38.10:8234", - "681d21b18644cd82ec13256a97f92bec1fff815683ef6f65dc7c993f098a4fe5@64.227.44.93:8234", - "058b3f20dc9bef7bb447f94d7663d793cfbc036720f97e52d7f13661b21818e1@161.35.156.226:8234", - "291323d78eb7186c3fa225ef5e305e95363e0ef06d42dca91bd4ef0254aed1ae@139.59.64.115:8234", - "85e425233a68697d4be91dd5d542305a8a327cd06d992d53c0913cef2fa75084@168.144.173.250:8234", -] -listen_addr = "0.0.0.0:8234" -max_connections = 256 -max_connections_per_ip = 16 -max_pending_handshakes = 32 -message_rate_per_second = 2048 -stream_open_rate_per_second = 32 - -[rpc] -cookie_dir = "cache_dir" -cookie_file_name = ".cookie" -debug_force_finished_sync = false -enable_cookie_auth = true -max_response_body_size = 52428800 -parallel_cpu_threads = 0 - -[state] -cache_dir = "cache_dir" -debug_skip_non_finalized_state_backup_task = false -delete_old_database = true -ephemeral = false -should_backup_non_finalized_state = true -storage_mode = "archive" - -[sync] -checkpoint_verify_concurrency_limit = 1000 -download_concurrency_limit = 100 -full_verify_concurrency_limit = 20 -parallel_cpu_threads = 0 -zakura_block_apply_concurrency_limit = 32 - -[tracing] -buffer_limit = 128000 -force_use_color = false -use_color = true -use_journald = false - -[zcashd_compat] -block_gossip_peer_ips = [] -enabled = false -manage_zcashd = false -restart_backoff = "2s" -restart_backoff_max = "5m" -restart_reset_after = "1h" -shutdown_grace_period = "5m" -startup_delay = "1s" -zcashd_extra_args = [] -zcashd_source = "path" diff --git a/docs/changelog/params.md b/docs/changelog/params.md index 835164d39f..de7df7d957 100644 --- a/docs/changelog/params.md +++ b/docs/changelog/params.md @@ -32,6 +32,7 @@ Keep entries **newest-first**. Each row records: | Parameter | Location | Old → New | PR | Why | | --- | --- | --- | --- | --- | +| `DEFAULT_ZAKURA_STREAM_RECEIVE_WINDOW` | `crates/zakura-network/src/zakura/handler.rs` | `32 MiB` → `16 MiB` | [#943](https://github.com/zakura-core/zakura/pull/943) | Leave connection receive credit for another service while one stream's application reads are paused. | | `MAX_CONCURRENT_UTXO_LOOKUPS` | `crates/zakura-consensus/src/transaction.rs` | serial (`1`) → `64` per block transaction | [#918](https://github.com/zakura-core/zakura/pull/918) | Overlap external UTXO waits while bounding pending lookups per transaction. Concurrent lookups start their six-minute timeout clocks together. | | `MAX_MINED_SUBMISSIONS` | `crates/zakura-rpc/src/methods/types/submit_block.rs` | unbounded → `16` submissions | [#748](https://github.com/zakura-core/zakura/pull/748) | Bound detached verification work across RPC cancellation. | | `non_finalized_write_slots` | `crates/zakura-state/src/service.rs` | unbounded → `1,000` contextual writes | [#748](https://github.com/zakura-core/zakura/pull/748) | Bound writer block bodies using the existing orphan queue capacity. | diff --git a/docs/changelog/unreleased/943.md b/docs/changelog/unreleased/943.md new file mode 100644 index 0000000000..8bb7827087 --- /dev/null +++ b/docs/changelog/unreleased/943.md @@ -0,0 +1,15 @@ +## Fixed + +- Keep native P2P session setup and recovery responsive under congestion, and + release stream capacity safely when sessions end + ([#943](https://github.com/zakura-core/zakura/pull/943)). +- Reject disallowed messages and malformed buffered responses even after a + session fails, while preserving block download stall penalties without + penalizing local delays + ([#943](https://github.com/zakura-core/zakura/pull/943)). + +## Changed + +- Keep the connection and unrelated services running when a persistent native + P2P stream write times out + ([#943](https://github.com/zakura-core/zakura/pull/943)). diff --git a/docs/design/service-sessions.md b/docs/design/service-sessions.md new file mode 100644 index 0000000000..e24b408c96 --- /dev/null +++ b/docs/design/service-sessions.md @@ -0,0 +1,147 @@ +# Service sessions + +A service declares its stream types through `Service::streams()`. The transport +groups persistent streams with the same capability into one session layout. +The transport admits the complete layout through one `Service::add_peer()` call. +Each stream has independent readers, writers, and queues. + +Request/response stream types remain in the declaration, but the transport opens +their streams per request. They do not participate in persistent session setup. +The service controls their application lifecycle. + +## Declaring a session + +For example, a service can declare data, requests, and events as three persistent +streams with the same capability: + +```rust +const DATA: Stream = Stream { + kind: 64, + version: 1, + frame_cap: 1024 * 1024, + capability: 1 << 16, + mode: StreamMode::Persistent, +}; +const REQUESTS: Stream = Stream { kind: 65, ..DATA }; +const EVENTS: Stream = Stream { kind: 66, ..DATA }; +const LOOKUP: Stream = Stream { + kind: 67, + mode: StreamMode::RequestResponse, + ..DATA +}; + +// Inside impl Service: +fn streams(&self) -> &[Stream] { + &[DATA, REQUESTS, EVENTS, LOOKUP] +} + +fn session_policy(&self) -> SessionPolicy { + SessionPolicy { + opening: SessionOpening::EitherSide, + reopen: true, + } +} +``` + +These identifiers illustrate the API. A production protocol must allocate its +own stream kinds and capability bit. + +The service does not declare membership a second time. The transport waits for +data, requests, and events before handing their receive/send handles to the +service. It does not wait for a lookup request. + +The service uses `message_types()`, `message_payload_limits()`, and +`stream_queue_depths()` to specify each stream's traffic and bounds. The protocol +defines message assignments; peers do not negotiate individual message types. +The service routes outgoing messages to the appropriate sender. + +`stream_write_policy()` sets each persistent stream's write deadline. The default +is ten seconds. A service can choose another duration or `UntilCancelled`. +A service that chooses `UntilCancelled` must enforce its own progress deadline. + +## Negotiating complete layouts + +Each capability identifies a complete persistent layout for its service. +Alternative layouts use different capabilities. Each alternative retains the +same lowest stream kind, whose version ranks the layouts. The registry selects +the highest mutually supported version of that primary stream and includes every +member of its layout. It never mixes members from different alternatives. + +For example, primary/request versions `3/1` and `2/4` select `3/1` when both +capabilities are available. The request stream's higher version in the older +layout does not override that choice. + +Adding a required stream changes the protocol layout. Allocate a new capability +and advance the primary stream's version. Keep the primary kind stable. +Changing a message assignment also requires a compatible protocol transition. +The transport cannot make an old peer understand a new layout or message. + +## Setup and retirement + +Multi-stream sessions append the same nonzero eight-byte session identifier to +each ordinary stream prelude. This retains #943's two-stream setup encoding. +Single-stream sessions retain their existing prelude without an extra identifier. + +The transport holds at most one incomplete session per service and connection. +The first complete member starts the setup deadline. Later members cannot extend +that deadline. Duplicate members, mismatched identifiers, and invalid declarations +cannot complete a session. Expiry releases every arrived member and defers another +offer through the existing cooldown. + +A different identifier can be a retry racing an abandoned offer's reset. The +transport discards both offers and applies the same cooldown, keeping unrelated +services connected. Offers during the cooldown cannot reserve service capacity +or extend the cooldown. + +While a remote session is active or its workers are retiring, additional offers +for that service are reset before reserving service capacity. The existing session +and unrelated services stay connected. A new offer can proceed after cleanup. + +`reserve_session()` charges service capacity once during setup. +`SessionResources::admitted()` signals complete setup. +The workers and application senders retain the shared resource owner until they +finish or drop it. Every member also consumes a transport stream slot. + +Every persistent member shares a local session identity, cancellation token, and +message-rate budget. A remote close on any member retires the session. +Cancellation resets unfinished writes before a replacement can send frames. +A write deadline retires the session without closing unrelated services on the +connection. Protocol violations can still close the connection. + +Dropping an application receiver stops delivery to that receiver. The transport +keeps reading within its frame and message-rate limits and discards those frames. +Retained application handles keep the session alive. Normal retirement waits for +every application handle to close and every member's queued writes to finish. + +The transport records the first remote close or write timeout before cancelling +its session. Block sync settles that failure against unanswered download work, +including when cancellation wins over receiving EOF. Local cancellation alone +does not charge the peer for a stall. + +Buffered block responses retain their bounded raw frames until decode capacity +is available, even after a remote session failure. Block sync validates those +responses before settling unanswered work, so malformed payloads cannot bypass +peer rejection during local backpressure. Connection shutdown still cancels the +pending validation and releases the session. + +The transport reports session exit after every worker and reader finishes. +Reopening follows the service's policy and demand. Ephemeral request completion +does not cancel the persistent session. + +Setup readiness does not impose ordering across streams. For example, a request +can arrive before a status message on another stream. The service must handle +that ordering or perform an application handshake before processing requests. + +## Migrating a pair consumer + +Remove `OrderedStreamPair` and `ordered_stream_pair()`. Declare the persistent +members with one capability in `streams()`. Use the `Session*` policy, demand, +and resource APIs. The transport supplies all declared members together. + +Move role-specific queue limits and write deadlines into the service hooks. +For the block-sync activation following #943, the service must declare the +one-slot request queue, the request write policy, and the 32-second data write +deadline. The transport no longer assigns those policies by role name. + +Production block sync in #943 remains a single-stream protocol. This change does +not activate the later block-sync layout.