velo_ext/admission.rs
1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Ordered per-target send admission.
5//!
6//! ## The hazard this exists to fix
7//!
8//! The obvious way to handle a full per-target channel is to hand the caller a
9//! future that owns the frame and completes the enqueue when polled. Velo
10//! shipped exactly that until 0.7 and it reorders frames: nothing enqueues the
11//! parked frame until somebody polls the future, so a *later* send to the same
12//! target can win the race — its `try_send` succeeds while the earlier frame is
13//! still sitting in an unpolled future. Two sends issued in order A, B arrive
14//! at the remote as B, A. Fire-and-forget senders make it worse; they may never
15//! poll at all, so A can sit behind an unbounded number of successors.
16//!
17//! The fix is structural rather than advisory: take the frame at `send` time
18//! and never let its delivery depend on the caller.
19//!
20//! ## The guarantee
21//!
22//! An [`AdmissionGate`] wraps one bounded [`flume::Sender`] and serialises
23//! everything that goes into it:
24//!
25//! > Frames enter the channel in the order their [`AdmissionGate::send`] calls
26//! > returned, regardless of which admissions (if any) are ever polled.
27//!
28//! Two structural choices carry that guarantee:
29//!
30//! 1. **Frames live in the gate, not in the future.** [`SendAdmission`] is a
31//! completion observer and a cancellation handle — never the owner of the
32//! frame. Delivery therefore cannot depend on who polls.
33//! 2. **A lazy driver task drains the queue.** The first queued ticket spawns a
34//! per-gate driver that pushes frames with `send_async` in FIFO order and
35//! resolves each ticket as its frame is enqueued. The driver parks (exits)
36//! when the queue empties and is respawned by the next queued ticket.
37//!
38//! The fast path is preserved: when the queue is empty *and* `try_send`
39//! succeeds, [`AdmissionGate::send`] returns [`SendOutcome::Admitted`] without
40//! allocating a ticket, waking a driver, or touching a waker. Only contended
41//! sends pay. Crucially, a frame the driver has checked out stays in the queue
42//! (with its payload taken) until it has been enqueued or dropped, so the
43//! "queue is empty" test cannot let a newcomer overtake a frame that is
44//! mid-flight.
45//!
46//! ## Dropping an admission does not cancel it
47//!
48//! Dropping a [`SendAdmission`] leaves the frame in the gate, and the gate
49//! still delivers it. Cancellation is explicit, via [`SendAdmission::cancel`].
50//! This is the point of the design, not an oversight: fire-and-forget senders
51//! drop their handle on the spot and must still see their frame delivered,
52//! which is irreconcilable with drop-cancels-the-send.
53//!
54//! Callers that want to *observe* an outcome without holding the future — a
55//! metric to record, a result channel to feed — register a
56//! [`SendAdmission::on_resolved`] hook instead of polling.
57
58use std::collections::VecDeque;
59use std::future::Future;
60use std::pin::Pin;
61use std::sync::{Arc, Mutex, MutexGuard, OnceLock, Weak};
62use std::task::{Context, Poll};
63
64use futures::task::AtomicWaker;
65use tokio_util::sync::CancellationToken;
66
67/// Take a lock, ignoring poisoning.
68///
69/// Every critical section here is a handful of `VecDeque` operations with no
70/// user code in between, so a poisoned lock means a panic elsewhere rather than
71/// torn state. Propagating the panic would strand every outstanding ticket.
72fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
73 mutex
74 .lock()
75 .unwrap_or_else(|poisoned| poisoned.into_inner())
76}
77
78/// Synchronously observable state of a [`SendAdmission`].
79#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
80pub enum AdmissionState {
81 /// The frame is queued in the gate and has not been enqueued yet.
82 Pending,
83 /// The frame has been enqueued on the transport's send channel.
84 Admitted,
85 /// The frame will never be enqueued; see the admission's error.
86 Failed,
87}
88
89/// Why a frame was never admitted to the transport's send channel.
90///
91/// An admission failure is *not* a delivery failure. A frame that is admitted
92/// can still fail on the wire, and those failures continue to flow through the
93/// transport's error handler.
94#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
95pub enum AdmissionError {
96 /// [`SendAdmission::cancel`] withdrew the ticket before it was enqueued.
97 #[error("send admission was cancelled")]
98 Cancelled,
99
100 /// The connection epoch that owned the ticket was replaced. The frame
101 /// belonged to a connection that no longer exists; resending it on the
102 /// successor connection is the caller's decision.
103 #[error("connection was replaced before the frame was admitted")]
104 ConnectionReplaced,
105
106 /// The transport's send channel was closed (receiver dropped) before the
107 /// frame could be enqueued.
108 #[error("transport send channel closed before the frame was admitted")]
109 ChannelClosed,
110
111 /// The epoch died for a transport-specific reason.
112 #[error("send admission failed: {0}")]
113 Failed(String),
114}
115
116/// Outcome of [`AdmissionGate::send`], and of
117/// [`Transport::send_message`](crate::transport::Transport::send_message).
118///
119/// Dropping this is a legitimate fire-and-forget pattern — the frame is already
120/// owned by the gate and will be delivered either way — so it is deliberately
121/// not `#[must_use]`.
122#[derive(Debug)]
123pub enum SendOutcome {
124 /// The frame was enqueued synchronously. No ticket was taken.
125 Admitted,
126 /// The frame was queued behind the gate's FIFO. The contained handle
127 /// observes (and can withdraw) the ticket; it does not drive delivery.
128 Pending(SendAdmission),
129}
130
131impl SendOutcome {
132 /// `true` if the frame took the synchronous fast path.
133 pub fn is_admitted(&self) -> bool {
134 matches!(self, Self::Admitted)
135 }
136
137 /// Take the admission handle, if this send queued a ticket.
138 pub fn into_pending(self) -> Option<SendAdmission> {
139 match self {
140 Self::Admitted => None,
141 Self::Pending(admission) => Some(admission),
142 }
143 }
144}
145
146/// Completion observer for one queued frame.
147///
148/// Resolves `Ok(())` when the frame is enqueued on the transport's send channel
149/// and `Err` when it will never be. **Polling is optional**: the gate's driver
150/// delivers queued frames whether or not anyone awaits, and dropping this
151/// handle does not cancel the send (see the [module docs](self)). Use
152/// [`cancel`](Self::cancel) to withdraw a frame.
153pub struct SendAdmission {
154 ticket: Arc<Ticket>,
155 /// `None` for admissions that were already resolved at construction.
156 gate: Option<Weak<dyn TicketRegistry>>,
157}
158
159impl SendAdmission {
160 fn new(ticket: Arc<Ticket>, gate: Weak<dyn TicketRegistry>) -> Self {
161 Self {
162 ticket,
163 gate: Some(gate),
164 }
165 }
166
167 /// An admission that is already resolved — nothing is queued anywhere.
168 fn resolved(outcome: Result<(), AdmissionError>) -> Self {
169 let ticket = Ticket::new();
170 ticket.resolve(outcome);
171 Self {
172 ticket: Arc::new(ticket),
173 gate: None,
174 }
175 }
176
177 /// Current state of the ticket. Cheap and synchronous; safe to call from a
178 /// non-async context.
179 pub fn state(&self) -> AdmissionState {
180 self.ticket.state()
181 }
182
183 /// Observe the outcome without polling.
184 ///
185 /// `on_resolved` receives exactly what awaiting this admission would have
186 /// produced, and runs exactly once. A hook registered before resolution —
187 /// or while earlier hooks are still being run — runs on the resolving
188 /// task, after every hook registered before it; only a hook registered
189 /// after all of that runs immediately, on the registering thread. That
190 /// makes it the mechanism for callers who cannot await — a fire-and-forget
191 /// send whose handle is about to be dropped, or a metric that must be
192 /// recorded when the frame really lands rather than when it was offered.
193 ///
194 /// The hook runs on whichever task resolves the ticket, normally the gate's
195 /// driver, so keep it short: a slow hook delays the next frame on this
196 /// target. It must not call back into the same gate.
197 ///
198 /// There is no hook for [`SendOutcome::Admitted`] because there is nothing
199 /// to wait for — that variant *is* the synchronous notification, and a
200 /// caller wanting "exactly once per send" handles it on the spot.
201 ///
202 /// Hooks are additive and run in registration order. The runtime installs
203 /// its own bookkeeping hook (outbound-frame metric, error reporting)
204 /// before the admission reaches the caller, so a caller registering its
205 /// own observer must not — and cannot — displace it.
206 pub fn on_resolved(
207 self,
208 on_resolved: impl FnOnce(&Result<(), AdmissionError>) + Send + 'static,
209 ) -> Self {
210 self.ticket.add_hook(Box::new(on_resolved));
211 self
212 }
213
214 /// Withdraw the frame from the gate.
215 ///
216 /// Successors keep their relative order — the ticket is removed from the
217 /// FIFO, not swapped out.
218 ///
219 /// Exactness has two regimes:
220 ///
221 /// - **Still queued** (the common case, including every ticket taken since
222 /// the driver last parked): the frame is removed and dropped under the
223 /// gate lock and the admission resolves [`AdmissionError::Cancelled`].
224 /// The frame is guaranteed never to reach the channel.
225 /// - **Already checked out** by the driver, i.e. parked in `send_async`
226 /// waiting for capacity: cancellation is best-effort. If the channel
227 /// accepts the frame before the driver observes the cancellation, the
228 /// frame is delivered and the admission resolves `Admitted` instead. In
229 /// the reverse race a frame that landed in the channel may still report
230 /// `Cancelled`. Only one frame per gate is ever in this window.
231 pub fn cancel(self) {
232 if let Some(gate) = self.gate.as_ref().and_then(Weak::upgrade) {
233 gate.cancel_ticket(&self.ticket);
234 } else {
235 // Gate is gone: nothing can deliver the frame any more.
236 self.ticket.resolve(Err(AdmissionError::Cancelled));
237 }
238 }
239}
240
241impl Future for SendAdmission {
242 type Output = Result<(), AdmissionError>;
243
244 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
245 let ticket = &self.get_mut().ticket;
246 if let Some(outcome) = ticket.outcome() {
247 return Poll::Ready(outcome);
248 }
249 ticket.waker.register(cx.waker());
250 // Re-check: the ticket may have resolved between the first read and the
251 // waker registration.
252 match ticket.outcome() {
253 Some(outcome) => Poll::Ready(outcome),
254 None => Poll::Pending,
255 }
256 }
257}
258
259impl std::fmt::Debug for SendAdmission {
260 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
261 f.debug_struct("SendAdmission")
262 .field("state", &self.state())
263 .finish_non_exhaustive()
264 }
265}
266
267/// Ticket state, resolved exactly once.
268enum TicketOutcome {
269 Pending,
270 Admitted,
271 Failed(AdmissionError),
272}
273
274fn read_outcome(outcome: &TicketOutcome) -> Option<Result<(), AdmissionError>> {
275 match outcome {
276 TicketOutcome::Pending => None,
277 TicketOutcome::Admitted => Some(Ok(())),
278 TicketOutcome::Failed(error) => Some(Err(error.clone())),
279 }
280}
281
282/// Callback installed by [`SendAdmission::on_resolved`].
283type ResolveHook = Box<dyn FnOnce(&Result<(), AdmissionError>) + Send>;
284
285/// The outcome and the not-yet-fired hooks share one lock.
286///
287/// That is what makes hook registration race-free: the "has it resolved yet?"
288/// test and the install are one critical section, so a hook can neither be
289/// stored on an already-resolved ticket (never to run) nor be missed by a
290/// `resolve` that ran a moment earlier.
291///
292/// Hooks are a `Vec`, not a slot: the runtime installs its own bookkeeping
293/// hook (outbound metric, error handler) before the admission ever reaches the
294/// caller, and the caller's hook must add to that, never replace it.
295///
296/// `hooks_drained` is what makes registration order hold across the
297/// registration-vs-resolution race: the resolver drains the vec in batches
298/// outside the lock, and only marks it drained once a locked re-check finds
299/// the vec empty. A registration landing in that window joins the vec — and
300/// runs on the resolver, in order, behind everything registered before it —
301/// instead of jumping the queue by running on the registering thread.
302struct TicketState {
303 outcome: TicketOutcome,
304 hooks: Vec<ResolveHook>,
305 hooks_drained: bool,
306}
307
308struct Ticket {
309 state: Mutex<TicketState>,
310 waker: AtomicWaker,
311 /// Set by [`SendAdmission::cancel`] when the driver already owns the frame.
312 cancel: CancellationToken,
313}
314
315impl Ticket {
316 fn new() -> Self {
317 Self {
318 state: Mutex::new(TicketState {
319 outcome: TicketOutcome::Pending,
320 hooks: Vec::new(),
321 hooks_drained: false,
322 }),
323 waker: AtomicWaker::new(),
324 cancel: CancellationToken::new(),
325 }
326 }
327
328 fn state(&self) -> AdmissionState {
329 match lock(&self.state).outcome {
330 TicketOutcome::Pending => AdmissionState::Pending,
331 TicketOutcome::Admitted => AdmissionState::Admitted,
332 TicketOutcome::Failed(_) => AdmissionState::Failed,
333 }
334 }
335
336 fn outcome(&self) -> Option<Result<(), AdmissionError>> {
337 read_outcome(&lock(&self.state).outcome)
338 }
339
340 /// Resolve the ticket. First writer wins; later attempts are no-ops.
341 ///
342 /// Never called with the gate lock held for a ticket whose waker could
343 /// re-enter the gate, so the woken task cannot deadlock against us. The
344 /// hooks run last, outside the ticket lock, for the same reason — in
345 /// registration order, so the runtime's bookkeeping hook fires before any
346 /// caller-installed observer.
347 fn resolve(&self, outcome: Result<(), AdmissionError>) {
348 {
349 let mut state = lock(&self.state);
350 if !matches!(state.outcome, TicketOutcome::Pending) {
351 return;
352 }
353 state.outcome = match &outcome {
354 Ok(()) => TicketOutcome::Admitted,
355 Err(error) => TicketOutcome::Failed(error.clone()),
356 };
357 }
358 self.waker.wake();
359 // Drain in batches until a locked re-check finds nothing new, then
360 // mark the drain complete in the same critical section. A hook
361 // registered while a batch runs lands in the vec and is picked up by
362 // the next iteration — still on this task, still in order.
363 loop {
364 let batch = {
365 let mut state = lock(&self.state);
366 if state.hooks.is_empty() {
367 state.hooks_drained = true;
368 return;
369 }
370 std::mem::take(&mut state.hooks)
371 };
372 for hook in batch {
373 hook(&outcome);
374 }
375 }
376 }
377
378 /// Add a completion hook.
379 ///
380 /// Runs on the spot only when the ticket has resolved *and* the resolver
381 /// has finished running every earlier hook; a registration racing the
382 /// resolver's drain joins the queue instead, so hooks always observe the
383 /// outcome in registration order.
384 fn add_hook(&self, hook: ResolveHook) {
385 let resolved = {
386 let mut state = lock(&self.state);
387 if !state.hooks_drained {
388 state.hooks.push(hook);
389 return;
390 }
391 read_outcome(&state.outcome).expect("hooks_drained implies resolved")
392 };
393 hook(&resolved);
394 }
395
396 fn is_live(&self) -> bool {
397 !self.cancel.is_cancelled() && matches!(lock(&self.state).outcome, TicketOutcome::Pending)
398 }
399}
400
401/// One connection lifetime's worth of tickets.
402///
403/// [`AdmissionGate::fail_all`] cancels the current epoch and installs a fresh
404/// one, so the gate itself is never poisoned: a successor connection's sends
405/// use the new epoch and are unaffected by the old one's failure.
406struct Epoch {
407 token: CancellationToken,
408 reason: OnceLock<AdmissionError>,
409}
410
411impl Epoch {
412 fn new() -> Self {
413 Self {
414 token: CancellationToken::new(),
415 reason: OnceLock::new(),
416 }
417 }
418
419 /// Kill the epoch. Only ever called once per epoch (under the gate lock).
420 fn fail(&self, error: AdmissionError) {
421 let _ = self.reason.set(error);
422 self.token.cancel();
423 }
424
425 fn reason(&self) -> AdmissionError {
426 self.reason
427 .get()
428 .cloned()
429 .unwrap_or(AdmissionError::ConnectionReplaced)
430 }
431}
432
433/// A frame waiting its turn.
434///
435/// `item` is taken when the driver checks the frame out for delivery, but the
436/// entry stays at the head of the queue until the send resolves. That keeps
437/// `queue.is_empty()` false for the whole in-flight window, which is what stops
438/// a fast-path send from overtaking a frame the driver is mid-way through.
439struct QueuedFrame<T> {
440 item: Option<T>,
441 ticket: Arc<Ticket>,
442}
443
444struct GateState<T> {
445 queue: VecDeque<QueuedFrame<T>>,
446 /// A driver task exists and owns the queue. Only the driver clears this,
447 /// and only under the lock with an empty queue.
448 driver_live: bool,
449 epoch: Arc<Epoch>,
450}
451
452struct GateInner<T> {
453 tx: flume::Sender<T>,
454 rt: tokio::runtime::Handle,
455 state: Mutex<GateState<T>>,
456}
457
458/// Type-erased handle so [`SendAdmission`] does not have to carry `T`.
459trait TicketRegistry: Send + Sync {
460 fn cancel_ticket(&self, ticket: &Arc<Ticket>);
461}
462
463impl<T: Send + 'static> TicketRegistry for GateInner<T> {
464 fn cancel_ticket(&self, ticket: &Arc<Ticket>) {
465 let removed = {
466 let mut state = lock(&self.state);
467 let position = state
468 .queue
469 .iter()
470 .position(|frame| Arc::ptr_eq(&frame.ticket, ticket));
471 match position {
472 // The frame is still queued: remove it (dropping the payload)
473 // without disturbing its successors.
474 Some(position) if state.queue[position].item.is_some() => {
475 state.queue.remove(position);
476 true
477 }
478 _ => false,
479 }
480 };
481 if removed {
482 ticket.resolve(Err(AdmissionError::Cancelled));
483 } else {
484 // Either the driver already owns the frame — it will observe this
485 // and abort — or the ticket has already resolved, in which case
486 // this is a no-op.
487 ticket.cancel.cancel();
488 }
489 }
490}
491
492/// Ordered admission to one bounded send channel.
493///
494/// A gate is scoped to whatever the transport treats as a target: one per
495/// connection for stream transports, one per peer over a shared writer for
496/// broker transports. Cloning is cheap (the state is shared) so a gate can be
497/// handed to every task that sends to that target.
498///
499/// See the [module docs](self) for the ordering guarantee and for why
500/// dropping an admission does not cancel its frame.
501pub struct AdmissionGate<T: Send + 'static> {
502 inner: Arc<GateInner<T>>,
503}
504
505impl<T: Send + 'static> Clone for AdmissionGate<T> {
506 fn clone(&self) -> Self {
507 Self {
508 inner: Arc::clone(&self.inner),
509 }
510 }
511}
512
513impl<T: Send + 'static> std::fmt::Debug for AdmissionGate<T> {
514 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
515 f.debug_struct("AdmissionGate")
516 .field("queued", &self.queued_len())
517 .finish_non_exhaustive()
518 }
519}
520
521impl<T: Send + 'static> AdmissionGate<T> {
522 /// Build a gate over a bounded channel.
523 ///
524 /// `rt` is used to spawn the driver task that drains queued frames;
525 /// transports already hold a [`Handle`](tokio::runtime::Handle) from
526 /// [`Transport::start`](crate::transport::Transport::start).
527 ///
528 /// The channel *must* be bounded — an unbounded channel never returns
529 /// `Full`, so every send takes the fast path and the gate is inert.
530 pub fn new(tx: flume::Sender<T>, rt: tokio::runtime::Handle) -> Self {
531 Self {
532 inner: Arc::new(GateInner {
533 tx,
534 rt,
535 state: Mutex::new(GateState {
536 queue: VecDeque::new(),
537 driver_live: false,
538 epoch: Arc::new(Epoch::new()),
539 }),
540 }),
541 }
542 }
543
544 /// Offer a frame to the channel, taking a ticket if it cannot go now.
545 ///
546 /// Synchronous and non-blocking. Returns [`SendOutcome::Admitted`] if
547 /// the queue was empty and the channel had room; otherwise the frame joins
548 /// the gate's FIFO and the returned [`SendAdmission`] observes its ticket.
549 ///
550 /// If the channel's receiver has already been dropped the frame is dropped
551 /// and an already-failed [`SendOutcome::Pending`] carrying
552 /// [`AdmissionError::ChannelClosed`] is returned — the two-variant outcome
553 /// has no honest "admitted" answer for a closed channel. A caller that
554 /// spawns a task per `Pending` will spawn one that finishes immediately.
555 pub fn send(&self, item: T) -> SendOutcome {
556 let mut state = lock(&self.inner.state);
557
558 // Fast path. The emptiness check and the try_send are one critical
559 // section, so no concurrent sender can slip between them.
560 let item = if state.queue.is_empty() {
561 match self.inner.tx.try_send(item) {
562 Ok(()) => return SendOutcome::Admitted,
563 Err(flume::TrySendError::Full(item)) => item,
564 Err(flume::TrySendError::Disconnected(_)) => {
565 return SendOutcome::Pending(SendAdmission::resolved(Err(
566 AdmissionError::ChannelClosed,
567 )));
568 }
569 }
570 } else {
571 item
572 };
573
574 let ticket = Arc::new(Ticket::new());
575 state.queue.push_back(QueuedFrame {
576 item: Some(item),
577 ticket: Arc::clone(&ticket),
578 });
579 let spawn_driver = !state.driver_live;
580 state.driver_live = true;
581 drop(state);
582
583 if spawn_driver {
584 let inner = Arc::clone(&self.inner);
585 self.inner.rt.spawn(drive(inner));
586 }
587
588 let weak = Arc::downgrade(&self.inner);
589 let gate: Weak<dyn TicketRegistry> = weak;
590 SendOutcome::Pending(SendAdmission::new(ticket, gate))
591 }
592
593 /// Fail every outstanding ticket and drop the frames behind them.
594 ///
595 /// Called when the connection this gate feeds dies: each queued frame
596 /// belongs to an epoch that no longer exists, so delivering it on the
597 /// successor connection would be wrong. Every pending [`SendAdmission`]
598 /// resolves `Err(error)` and flips to [`AdmissionState::Failed`].
599 ///
600 /// The gate is **not** poisoned — a fresh epoch is installed and later
601 /// sends admit normally, so a transport may either rebuild a gate per
602 /// connection or keep one and call this on each reconnect.
603 ///
604 /// The one frame the driver may already have handed to the channel is
605 /// resolved by the driver rather than here: if the channel accepted it
606 /// before the epoch died it resolves `Admitted`, because it really was
607 /// delivered. Everything not yet enqueued fails.
608 pub fn fail_all(&self, error: AdmissionError) {
609 let failed = {
610 let mut state = lock(&self.inner.state);
611 let dead = std::mem::replace(&mut state.epoch, Arc::new(Epoch::new()));
612 dead.fail(error.clone());
613
614 let mut failed = Vec::new();
615 let mut retained = VecDeque::new();
616 for frame in std::mem::take(&mut state.queue) {
617 if frame.item.is_some() {
618 // Dropping `frame` here drops the payload.
619 failed.push(frame.ticket);
620 } else {
621 // Checked out by the driver; it owns the resolution. Keep
622 // it at the head so successors stay ordered behind it.
623 retained.push_back(frame);
624 }
625 }
626 state.queue = retained;
627 failed
628 };
629
630 for ticket in failed {
631 ticket.resolve(Err(error.clone()));
632 }
633 }
634
635 /// Number of tickets the gate is still holding.
636 ///
637 /// Zero on a gate whose sends are all taking the fast path. Primarily for
638 /// tests, metrics, and saturation debugging.
639 pub fn queued_len(&self) -> usize {
640 lock(&self.inner.state).queue.len()
641 }
642
643 /// Whether a driver task currently owns the queue.
644 #[cfg(test)]
645 fn driver_live(&self) -> bool {
646 lock(&self.inner.state).driver_live
647 }
648
649 /// Whether the driver has checked the head frame out and is parked in
650 /// `send_async` waiting for capacity.
651 #[cfg(test)]
652 fn head_checked_out(&self) -> bool {
653 lock(&self.inner.state)
654 .queue
655 .front()
656 .is_some_and(|frame| frame.item.is_none())
657 }
658}
659
660/// Drain the gate's queue in FIFO order until it empties.
661///
662/// This is the only thing that ever enqueues a queued frame, which is why the
663/// gate's ordering guarantee holds without any caller polling.
664async fn drive<T: Send + 'static>(inner: Arc<GateInner<T>>) {
665 while let Some(checkout) = check_out_head(&inner) {
666 let Checkout {
667 item,
668 ticket,
669 epoch,
670 } = checkout;
671
672 // The send future owns the frame for the duration of this block and is
673 // dropped before the ticket resolves — dropping it before flume accepts
674 // the frame means the frame is never enqueued, so an aborted frame can
675 // never surface behind its successors.
676 let outcome = {
677 let send = inner.tx.send_async(item);
678 tokio::pin!(send);
679 tokio::select! {
680 // Biased so that a frame flume has already accepted reports
681 // `Admitted` rather than being mislabelled by a cancellation
682 // that lost the race. The frame is in the channel either way.
683 biased;
684 result = &mut send => match result {
685 Ok(()) => Ok(()),
686 Err(flume::SendError(_)) => Err(AdmissionError::ChannelClosed),
687 },
688 () = ticket.cancel.cancelled() => Err(AdmissionError::Cancelled),
689 () = epoch.token.cancelled() => Err(epoch.reason()),
690 }
691 };
692
693 {
694 let mut state = lock(&inner.state);
695 if let Some(head) = state.queue.front()
696 && Arc::ptr_eq(&head.ticket, &ticket)
697 {
698 state.queue.pop_front();
699 }
700 }
701 // Resolved outside the gate lock: the frame has already left (or been
702 // dropped), so nothing a woken task does can reorder anything.
703 ticket.resolve(outcome);
704 }
705}
706
707struct Checkout<T> {
708 item: T,
709 ticket: Arc<Ticket>,
710 epoch: Arc<Epoch>,
711}
712
713/// Take the next deliverable frame, skipping tickets that died while queued.
714///
715/// Returns `None` once the queue is empty, clearing `driver_live` under the
716/// same lock so the next queued ticket spawns a fresh driver.
717fn check_out_head<T: Send + 'static>(inner: &Arc<GateInner<T>>) -> Option<Checkout<T>> {
718 let mut state = lock(&inner.state);
719 loop {
720 let Some(head) = state.queue.front() else {
721 state.driver_live = false;
722 return None;
723 };
724 if !head.ticket.is_live() {
725 let frame = state.queue.pop_front().expect("front was just observed");
726 drop(state);
727 frame.ticket.resolve(Err(AdmissionError::Cancelled));
728 state = lock(&inner.state);
729 continue;
730 }
731
732 let epoch = Arc::clone(&state.epoch);
733 let mut checked_out = None;
734 if let Some(head) = state.queue.front_mut()
735 && let Some(item) = head.item.take()
736 {
737 checked_out = Some((item, Arc::clone(&head.ticket)));
738 }
739 match checked_out {
740 Some((item, ticket)) => {
741 return Some(Checkout {
742 item,
743 ticket,
744 epoch,
745 });
746 }
747 // Defensive: a checked-out frame is always removed before the
748 // driver looks again, so this is unreachable in practice.
749 None => {
750 let frame = state.queue.pop_front().expect("front was just observed");
751 drop(state);
752 frame.ticket.resolve(Err(AdmissionError::Cancelled));
753 state = lock(&inner.state);
754 }
755 }
756 }
757}
758
759#[cfg(test)]
760mod tests;