use std::collections::VecDeque;
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex, MutexGuard, OnceLock, Weak};
use std::task::{Context, Poll};
use futures::task::AtomicWaker;
use tokio_util::sync::CancellationToken;
fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
mutex
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum AdmissionState {
Pending,
Admitted,
Failed,
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum AdmissionError {
#[error("send admission was cancelled")]
Cancelled,
#[error("connection was replaced before the frame was admitted")]
ConnectionReplaced,
#[error("transport send channel closed before the frame was admitted")]
ChannelClosed,
#[error("send admission failed: {0}")]
Failed(String),
}
#[derive(Debug)]
pub enum SendOutcome {
Admitted,
Pending(SendAdmission),
}
impl SendOutcome {
pub fn is_admitted(&self) -> bool {
matches!(self, Self::Admitted)
}
pub fn into_pending(self) -> Option<SendAdmission> {
match self {
Self::Admitted => None,
Self::Pending(admission) => Some(admission),
}
}
}
pub struct SendAdmission {
ticket: Arc<Ticket>,
gate: Option<Weak<dyn TicketRegistry>>,
}
impl SendAdmission {
fn new(ticket: Arc<Ticket>, gate: Weak<dyn TicketRegistry>) -> Self {
Self {
ticket,
gate: Some(gate),
}
}
fn resolved(outcome: Result<(), AdmissionError>) -> Self {
let ticket = Ticket::new();
ticket.resolve(outcome);
Self {
ticket: Arc::new(ticket),
gate: None,
}
}
pub fn state(&self) -> AdmissionState {
self.ticket.state()
}
pub fn on_resolved(
self,
on_resolved: impl FnOnce(&Result<(), AdmissionError>) + Send + 'static,
) -> Self {
self.ticket.add_hook(Box::new(on_resolved));
self
}
pub fn cancel(self) {
if let Some(gate) = self.gate.as_ref().and_then(Weak::upgrade) {
gate.cancel_ticket(&self.ticket);
} else {
self.ticket.resolve(Err(AdmissionError::Cancelled));
}
}
}
impl Future for SendAdmission {
type Output = Result<(), AdmissionError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let ticket = &self.get_mut().ticket;
if let Some(outcome) = ticket.outcome() {
return Poll::Ready(outcome);
}
ticket.waker.register(cx.waker());
match ticket.outcome() {
Some(outcome) => Poll::Ready(outcome),
None => Poll::Pending,
}
}
}
impl std::fmt::Debug for SendAdmission {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SendAdmission")
.field("state", &self.state())
.finish_non_exhaustive()
}
}
enum TicketOutcome {
Pending,
Admitted,
Failed(AdmissionError),
}
fn read_outcome(outcome: &TicketOutcome) -> Option<Result<(), AdmissionError>> {
match outcome {
TicketOutcome::Pending => None,
TicketOutcome::Admitted => Some(Ok(())),
TicketOutcome::Failed(error) => Some(Err(error.clone())),
}
}
type ResolveHook = Box<dyn FnOnce(&Result<(), AdmissionError>) + Send>;
struct TicketState {
outcome: TicketOutcome,
hooks: Vec<ResolveHook>,
hooks_drained: bool,
}
struct Ticket {
state: Mutex<TicketState>,
waker: AtomicWaker,
cancel: CancellationToken,
}
impl Ticket {
fn new() -> Self {
Self {
state: Mutex::new(TicketState {
outcome: TicketOutcome::Pending,
hooks: Vec::new(),
hooks_drained: false,
}),
waker: AtomicWaker::new(),
cancel: CancellationToken::new(),
}
}
fn state(&self) -> AdmissionState {
match lock(&self.state).outcome {
TicketOutcome::Pending => AdmissionState::Pending,
TicketOutcome::Admitted => AdmissionState::Admitted,
TicketOutcome::Failed(_) => AdmissionState::Failed,
}
}
fn outcome(&self) -> Option<Result<(), AdmissionError>> {
read_outcome(&lock(&self.state).outcome)
}
fn resolve(&self, outcome: Result<(), AdmissionError>) {
{
let mut state = lock(&self.state);
if !matches!(state.outcome, TicketOutcome::Pending) {
return;
}
state.outcome = match &outcome {
Ok(()) => TicketOutcome::Admitted,
Err(error) => TicketOutcome::Failed(error.clone()),
};
}
self.waker.wake();
loop {
let batch = {
let mut state = lock(&self.state);
if state.hooks.is_empty() {
state.hooks_drained = true;
return;
}
std::mem::take(&mut state.hooks)
};
for hook in batch {
hook(&outcome);
}
}
}
fn add_hook(&self, hook: ResolveHook) {
let resolved = {
let mut state = lock(&self.state);
if !state.hooks_drained {
state.hooks.push(hook);
return;
}
read_outcome(&state.outcome).expect("hooks_drained implies resolved")
};
hook(&resolved);
}
fn is_live(&self) -> bool {
!self.cancel.is_cancelled() && matches!(lock(&self.state).outcome, TicketOutcome::Pending)
}
}
struct Epoch {
token: CancellationToken,
reason: OnceLock<AdmissionError>,
}
impl Epoch {
fn new() -> Self {
Self {
token: CancellationToken::new(),
reason: OnceLock::new(),
}
}
fn fail(&self, error: AdmissionError) {
let _ = self.reason.set(error);
self.token.cancel();
}
fn reason(&self) -> AdmissionError {
self.reason
.get()
.cloned()
.unwrap_or(AdmissionError::ConnectionReplaced)
}
}
struct QueuedFrame<T> {
item: Option<T>,
ticket: Arc<Ticket>,
}
struct GateState<T> {
queue: VecDeque<QueuedFrame<T>>,
driver_live: bool,
epoch: Arc<Epoch>,
}
struct GateInner<T> {
tx: flume::Sender<T>,
rt: tokio::runtime::Handle,
state: Mutex<GateState<T>>,
}
trait TicketRegistry: Send + Sync {
fn cancel_ticket(&self, ticket: &Arc<Ticket>);
}
impl<T: Send + 'static> TicketRegistry for GateInner<T> {
fn cancel_ticket(&self, ticket: &Arc<Ticket>) {
let removed = {
let mut state = lock(&self.state);
let position = state
.queue
.iter()
.position(|frame| Arc::ptr_eq(&frame.ticket, ticket));
match position {
Some(position) if state.queue[position].item.is_some() => {
state.queue.remove(position);
true
}
_ => false,
}
};
if removed {
ticket.resolve(Err(AdmissionError::Cancelled));
} else {
ticket.cancel.cancel();
}
}
}
pub struct AdmissionGate<T: Send + 'static> {
inner: Arc<GateInner<T>>,
}
impl<T: Send + 'static> Clone for AdmissionGate<T> {
fn clone(&self) -> Self {
Self {
inner: Arc::clone(&self.inner),
}
}
}
impl<T: Send + 'static> std::fmt::Debug for AdmissionGate<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AdmissionGate")
.field("queued", &self.queued_len())
.finish_non_exhaustive()
}
}
impl<T: Send + 'static> AdmissionGate<T> {
pub fn new(tx: flume::Sender<T>, rt: tokio::runtime::Handle) -> Self {
Self {
inner: Arc::new(GateInner {
tx,
rt,
state: Mutex::new(GateState {
queue: VecDeque::new(),
driver_live: false,
epoch: Arc::new(Epoch::new()),
}),
}),
}
}
pub fn send(&self, item: T) -> SendOutcome {
let mut state = lock(&self.inner.state);
let item = if state.queue.is_empty() {
match self.inner.tx.try_send(item) {
Ok(()) => return SendOutcome::Admitted,
Err(flume::TrySendError::Full(item)) => item,
Err(flume::TrySendError::Disconnected(_)) => {
return SendOutcome::Pending(SendAdmission::resolved(Err(
AdmissionError::ChannelClosed,
)));
}
}
} else {
item
};
let ticket = Arc::new(Ticket::new());
state.queue.push_back(QueuedFrame {
item: Some(item),
ticket: Arc::clone(&ticket),
});
let spawn_driver = !state.driver_live;
state.driver_live = true;
drop(state);
if spawn_driver {
let inner = Arc::clone(&self.inner);
self.inner.rt.spawn(drive(inner));
}
let weak = Arc::downgrade(&self.inner);
let gate: Weak<dyn TicketRegistry> = weak;
SendOutcome::Pending(SendAdmission::new(ticket, gate))
}
pub fn fail_all(&self, error: AdmissionError) {
let failed = {
let mut state = lock(&self.inner.state);
let dead = std::mem::replace(&mut state.epoch, Arc::new(Epoch::new()));
dead.fail(error.clone());
let mut failed = Vec::new();
let mut retained = VecDeque::new();
for frame in std::mem::take(&mut state.queue) {
if frame.item.is_some() {
failed.push(frame.ticket);
} else {
retained.push_back(frame);
}
}
state.queue = retained;
failed
};
for ticket in failed {
ticket.resolve(Err(error.clone()));
}
}
pub fn queued_len(&self) -> usize {
lock(&self.inner.state).queue.len()
}
#[cfg(test)]
fn driver_live(&self) -> bool {
lock(&self.inner.state).driver_live
}
#[cfg(test)]
fn head_checked_out(&self) -> bool {
lock(&self.inner.state)
.queue
.front()
.is_some_and(|frame| frame.item.is_none())
}
}
async fn drive<T: Send + 'static>(inner: Arc<GateInner<T>>) {
while let Some(checkout) = check_out_head(&inner) {
let Checkout {
item,
ticket,
epoch,
} = checkout;
let outcome = {
let send = inner.tx.send_async(item);
tokio::pin!(send);
tokio::select! {
biased;
result = &mut send => match result {
Ok(()) => Ok(()),
Err(flume::SendError(_)) => Err(AdmissionError::ChannelClosed),
},
() = ticket.cancel.cancelled() => Err(AdmissionError::Cancelled),
() = epoch.token.cancelled() => Err(epoch.reason()),
}
};
{
let mut state = lock(&inner.state);
if let Some(head) = state.queue.front()
&& Arc::ptr_eq(&head.ticket, &ticket)
{
state.queue.pop_front();
}
}
ticket.resolve(outcome);
}
}
struct Checkout<T> {
item: T,
ticket: Arc<Ticket>,
epoch: Arc<Epoch>,
}
fn check_out_head<T: Send + 'static>(inner: &Arc<GateInner<T>>) -> Option<Checkout<T>> {
let mut state = lock(&inner.state);
loop {
let Some(head) = state.queue.front() else {
state.driver_live = false;
return None;
};
if !head.ticket.is_live() {
let frame = state.queue.pop_front().expect("front was just observed");
drop(state);
frame.ticket.resolve(Err(AdmissionError::Cancelled));
state = lock(&inner.state);
continue;
}
let epoch = Arc::clone(&state.epoch);
let mut checked_out = None;
if let Some(head) = state.queue.front_mut()
&& let Some(item) = head.item.take()
{
checked_out = Some((item, Arc::clone(&head.ticket)));
}
match checked_out {
Some((item, ticket)) => {
return Some(Checkout {
item,
ticket,
epoch,
});
}
None => {
let frame = state.queue.pop_front().expect("front was just observed");
drop(state);
frame.ticket.resolve(Err(AdmissionError::Cancelled));
state = lock(&inner.state);
}
}
}
}
#[cfg(test)]
mod tests;