use std::future::Future;
use crate::{
quic::{self, ConnectionError},
util::set_once::SetOnce,
};
#[derive(Debug, Clone, Default)]
pub struct ConnectionErrorLatch {
terminal_error: SetOnce<ConnectionError>,
}
impl ConnectionErrorLatch {
pub fn new() -> Self {
Self::default()
}
pub(crate) fn latch_with(&self, f: impl FnOnce() -> ConnectionError) -> ConnectionError {
match self.terminal_error.peek() {
Some(existing) => existing,
None => {
let error = f();
let _ = self.terminal_error.set(error.clone());
error
}
}
}
pub(crate) fn latch_raw(&self, error: ConnectionError) -> ConnectionError {
self.latch_with(|| error)
}
pub(crate) fn check(&self) -> Result<(), ConnectionError> {
match self.terminal_error.peek() {
Some(error) => Err(error),
None => Ok(()),
}
}
pub(crate) fn peek(&self) -> Option<ConnectionError> {
self.terminal_error.peek()
}
}
pub(crate) mod sealed {
use super::ConnectionErrorLatch;
pub trait HasLatch {
fn latch(&self) -> &ConnectionErrorLatch;
}
}
pub(crate) use sealed::HasLatch;
#[allow(async_fn_in_trait)]
pub trait LifecycleExt: quic::Lifecycle + HasLatch {
fn check_with_probe(
&self,
probe: impl FnOnce() -> Option<ConnectionError>,
) -> Result<(), ConnectionError> {
self.latch().check()?;
match probe() {
None => Ok(()),
Some(error) => Err(self.latch().latch_with(|| error)),
}
}
async fn resolve_closed(&self, wait: impl Future<Output = ConnectionError>) -> ConnectionError {
if let Some(error) = self.latch().peek() {
return error;
}
let error = wait.await;
self.latch().latch_raw(error)
}
async fn guard<T>(
&self,
fut: impl Future<Output = Result<T, ConnectionError>>,
) -> Result<T, ConnectionError> {
quic::Lifecycle::check(self)?;
match fut.await {
Ok(v) => Ok(v),
Err(e) => Err(self.latch().latch_with(|| e)),
}
}
async fn guard_with<T, E, M>(
&self,
fut: impl Future<Output = Result<T, E>>,
map_err: M,
) -> Result<T, ConnectionError>
where
M: FnOnce(E) -> ConnectionError,
{
quic::Lifecycle::check(self)?;
match fut.await {
Ok(v) => Ok(v),
Err(e) => Err(self.latch().latch_with(|| map_err(e))),
}
}
fn guard_sync<T>(
&self,
f: impl FnOnce() -> Result<T, ConnectionError>,
) -> Result<T, ConnectionError> {
quic::Lifecycle::check(self)?;
match f() {
Ok(v) => Ok(v),
Err(e) => Err(self.latch().latch_with(|| e)),
}
}
fn guard_sync_with<T, E, M>(
&self,
f: impl FnOnce() -> Result<T, E>,
map_err: M,
) -> Result<T, ConnectionError>
where
M: FnOnce(E) -> ConnectionError,
{
quic::Lifecycle::check(self)?;
match f() {
Ok(v) => Ok(v),
Err(e) => Err(self.latch().latch_with(|| map_err(e))),
}
}
}
impl<T: quic::Lifecycle + HasLatch + ?Sized> LifecycleExt for T {}
#[cfg(test)]
mod tests {
use std::{
borrow::Cow,
sync::{Arc, Mutex},
};
use super::*;
use crate::{error::Code, varint::VarInt};
fn make_err(tag: u32) -> ConnectionError {
ConnectionError::Transport {
source: quic::TransportError {
kind: VarInt::from_u32(tag),
frame_type: VarInt::from_u32(0),
reason: format!("test-{tag}").into(),
},
}
}
struct TestLifecycle {
latch: ConnectionErrorLatch,
probe: Mutex<Option<ConnectionError>>,
wait: Mutex<Option<ConnectionError>>,
}
impl TestLifecycle {
fn new() -> Self {
Self {
latch: ConnectionErrorLatch::new(),
probe: Mutex::new(None),
wait: Mutex::new(None),
}
}
fn set_probe(&self, error: ConnectionError) {
*self.probe.lock().unwrap() = Some(error);
}
fn set_wait(&self, error: ConnectionError) {
*self.wait.lock().unwrap() = Some(error);
}
}
impl HasLatch for TestLifecycle {
fn latch(&self) -> &ConnectionErrorLatch {
&self.latch
}
}
impl quic::Lifecycle for TestLifecycle {
fn close(&self, _code: Code, _reason: Cow<'static, str>) {}
fn check(&self) -> Result<(), ConnectionError> {
self.check_with_probe(|| self.probe.lock().unwrap().take())
}
async fn closed(&self) -> ConnectionError {
self.resolve_closed(async {
self.wait
.lock()
.unwrap()
.take()
.unwrap_or_else(|| make_err(99))
})
.await
}
}
#[test]
fn latch_with_is_lazy_when_already_latched() {
let latch = ConnectionErrorLatch::new();
let first = latch.latch_with(|| make_err(1));
assert!(matches!(
&first,
ConnectionError::Transport { source } if source.kind == VarInt::from_u32(1)
));
let called = Mutex::new(false);
let again = latch.latch_with(|| {
*called.lock().unwrap() = true;
make_err(2)
});
assert!(
!*called.lock().unwrap(),
"closure must not be invoked once latched"
);
assert!(matches!(
&again,
ConnectionError::Transport { source } if source.kind == VarInt::from_u32(1)
));
}
#[test]
fn check_with_probe_folds_probe_into_latch() {
let lc = TestLifecycle::new();
lc.set_probe(make_err(42));
let e1 = quic::Lifecycle::check(&lc).unwrap_err();
let e2 = quic::Lifecycle::check(&lc).unwrap_err();
match (&e1, &e2) {
(
ConnectionError::Transport { source: s1 },
ConnectionError::Transport { source: s2 },
) => {
assert_eq!(s1.kind, VarInt::from_u32(42));
assert_eq!(s2.kind, VarInt::from_u32(42));
}
_ => panic!("unexpected error shape"),
}
}
#[test]
fn check_with_probe_no_error_when_clean() {
let lc = TestLifecycle::new();
assert!(quic::Lifecycle::check(&lc).is_ok());
}
#[tokio::test]
async fn resolve_closed_returns_latched_without_awaiting() {
let lc = TestLifecycle::new();
lc.latch.latch_with(|| make_err(7));
let never = Mutex::new(Some(make_err(8)));
let got = lc
.resolve_closed(async { never.lock().unwrap().take().unwrap() })
.await;
assert!(matches!(
&got,
ConnectionError::Transport { source } if source.kind == VarInt::from_u32(7)
));
assert!(never.lock().unwrap().is_some(), "wait must not be polled");
}
#[tokio::test]
async fn resolve_closed_latches_wait_result() {
let lc = TestLifecycle::new();
lc.set_wait(make_err(5));
let got = quic::Lifecycle::closed(&lc).await;
assert!(matches!(
&got,
ConnectionError::Transport { source } if source.kind == VarInt::from_u32(5)
));
let again = quic::Lifecycle::closed(&lc).await;
assert!(matches!(
&again,
ConnectionError::Transport { source } if source.kind == VarInt::from_u32(5)
));
}
#[tokio::test]
async fn guard_with_skips_closure_when_latched() {
let lc = TestLifecycle::new();
lc.latch.latch_with(|| make_err(7));
let called = Arc::new(Mutex::new(false));
let called2 = called.clone();
let res: Result<(), ConnectionError> = lc
.guard_with(async { Result::<(), &'static str>::Err("x") }, move |_| {
*called2.lock().unwrap() = true;
make_err(8)
})
.await;
assert!(res.is_err());
assert!(
!*called.lock().unwrap(),
"map_err must stay lazy once latched"
);
}
#[tokio::test]
async fn guard_with_success_is_untouched() {
let lc = TestLifecycle::new();
let out: Result<i32, ConnectionError> = lc
.guard_with(async { Ok::<_, &'static str>(7) }, |_| make_err(1))
.await;
assert_eq!(out.unwrap(), 7);
}
#[test]
fn guard_sync_latches_only_first_error() {
let lc = TestLifecycle::new();
let a = lc.guard_sync(|| Err::<(), _>(make_err(1))).unwrap_err();
let b = lc.guard_sync(|| Err::<(), _>(make_err(2))).unwrap_err();
match (&a, &b) {
(
ConnectionError::Transport { source: sa },
ConnectionError::Transport { source: sb },
) => {
assert_eq!(sa.kind, sb.kind);
assert_eq!(sa.kind, VarInt::from_u32(1));
}
_ => panic!("unexpected error shape"),
}
}
}