use tokio::sync::oneshot;
use tokio::task::JoinHandle;
pub struct TransportHandle {
cancel: Option<oneshot::Sender<()>>,
join: Option<JoinHandle<()>>,
}
impl TransportHandle {
pub fn noop() -> Self {
Self {
cancel: None,
join: None,
}
}
pub fn new(cancel: oneshot::Sender<()>, join: JoinHandle<()>) -> Self {
Self {
cancel: Some(cancel),
join: Some(join),
}
}
pub async fn shutdown(&mut self) {
if let Some(cancel) = self.cancel.take() {
let _ = cancel.send(());
}
if let Some(join) = self.join.take() {
if let Err(err) = join.await {
tracing::warn!(error = ?err, "peer transport task join error");
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn shutdown_signals_the_task_and_joins() {
let (cancel_tx, cancel_rx) = oneshot::channel::<()>();
let join = tokio::spawn(async move {
let _ = cancel_rx.await;
});
let mut handle = TransportHandle::new(cancel_tx, join);
handle.shutdown().await;
handle.shutdown().await;
}
#[tokio::test]
async fn noop_shutdown_is_harmless() {
let mut handle = TransportHandle::noop();
handle.shutdown().await;
}
#[tokio::test]
async fn shutdown_after_task_already_returned() {
let (cancel_tx, _cancel_rx) = oneshot::channel::<()>();
let join = tokio::spawn(async {});
tokio::task::yield_now().await;
let mut handle = TransportHandle::new(cancel_tx, join);
handle.shutdown().await;
}
}