Jędrzej Stuczyński
827c13b69e
dont build netstack in CI additional rust 2024 fixes fixes removed temp.rs first round of cleanup removed duplicated NS types moved gateway probe to the monorepo
702 lines
22 KiB
Rust
702 lines
22 KiB
Rust
// Copyright 2022 - Nym Technologies SA <contact@nymtech.net>
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
use std::{
|
|
error::Error,
|
|
sync::atomic::{AtomicBool, Ordering},
|
|
time::Duration,
|
|
};
|
|
|
|
use futures::{FutureExt, SinkExt, StreamExt, future::pending};
|
|
use log::{Level, log};
|
|
use tokio::sync::{
|
|
mpsc,
|
|
watch::{self, error::SendError},
|
|
};
|
|
|
|
use crate::event::{SentStatus, StatusReceiver, StatusSender, TaskStatus};
|
|
|
|
#[cfg(not(target_arch = "wasm32"))]
|
|
use tokio::time::{sleep, timeout};
|
|
|
|
#[cfg(target_arch = "wasm32")]
|
|
use wasmtimer::tokio::{sleep, timeout};
|
|
|
|
const DEFAULT_SHUTDOWN_TIMER_SECS: u64 = 5;
|
|
|
|
pub(crate) type SentError = Box<dyn Error + Send + Sync>;
|
|
type ErrorSender = mpsc::UnboundedSender<SentError>;
|
|
type ErrorReceiver = mpsc::UnboundedReceiver<SentError>;
|
|
|
|
fn try_recover_name(name: &Option<String>) -> String {
|
|
if let Some(name) = name {
|
|
name.clone()
|
|
} else {
|
|
"unknown".to_string()
|
|
}
|
|
}
|
|
|
|
#[derive(thiserror::Error, Debug)]
|
|
enum TaskError {
|
|
#[error("Task '{}' halted unexpectedly", try_recover_name(.shutdown_name))]
|
|
UnexpectedHalt { shutdown_name: Option<String> },
|
|
}
|
|
|
|
/// Listens to status and error messages from tasks, as well as notifying them to gracefully
|
|
/// shutdown. Keeps track of if task stop unexpectedly, such as in a panic.
|
|
#[deprecated(note = "use ShutdownManager instead")]
|
|
#[derive(Debug)]
|
|
pub struct TaskManager {
|
|
// optional name assigned to the task manager that all subscribed task clients will inherit
|
|
name: Option<String>,
|
|
|
|
// These channels have the dual purpose of signalling it's time to shutdown, but also to keep
|
|
// track of which tasks we are still waiting for.
|
|
notify_tx: watch::Sender<()>,
|
|
notify_rx: Option<watch::Receiver<()>>,
|
|
#[cfg_attr(target_arch = "wasm32", allow(dead_code))]
|
|
shutdown_timer_secs: u64,
|
|
|
|
// If any task failed, it needs to report separately
|
|
task_return_error_tx: ErrorSender,
|
|
task_return_error_rx: Option<ErrorReceiver>,
|
|
|
|
// Also signal when the notifier is dropped, in case the task exits unexpectedly.
|
|
// Why are we not reusing the return error channel? Well, let me tell you kids, it's because I
|
|
// didn't manage to reliably get the explicitly sent error (and not the error sent during drop)
|
|
task_drop_tx: ErrorSender,
|
|
task_drop_rx: Option<ErrorReceiver>,
|
|
|
|
// A task might also send non-fatal errors (effectively, warnings) while running that is not
|
|
// the result of exiting.
|
|
task_status_tx: StatusSender,
|
|
task_status_rx: Option<StatusReceiver>,
|
|
}
|
|
|
|
#[allow(deprecated)]
|
|
impl Default for TaskManager {
|
|
fn default() -> Self {
|
|
let (notify_tx, notify_rx) = watch::channel(());
|
|
let (task_halt_tx, task_halt_rx) = mpsc::unbounded_channel();
|
|
let (task_drop_tx, task_drop_rx) = mpsc::unbounded_channel();
|
|
// The status channel is bounded (unlike the others), since it's not always the case that
|
|
// there is a listener.
|
|
let (task_status_tx, task_status_rx) = futures::channel::mpsc::channel(128);
|
|
Self {
|
|
name: None,
|
|
notify_tx,
|
|
notify_rx: Some(notify_rx),
|
|
shutdown_timer_secs: DEFAULT_SHUTDOWN_TIMER_SECS,
|
|
task_return_error_tx: task_halt_tx,
|
|
task_return_error_rx: Some(task_halt_rx),
|
|
task_drop_tx,
|
|
task_drop_rx: Some(task_drop_rx),
|
|
task_status_tx,
|
|
task_status_rx: Some(task_status_rx),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[allow(deprecated)]
|
|
#[allow(clippy::expect_used)]
|
|
impl TaskManager {
|
|
pub fn new(shutdown_timer_secs: u64) -> Self {
|
|
Self {
|
|
shutdown_timer_secs,
|
|
..Default::default()
|
|
}
|
|
}
|
|
|
|
#[must_use]
|
|
pub fn named<S: Into<String>>(mut self, name: S) -> Self {
|
|
self.name = Some(name.into());
|
|
self
|
|
}
|
|
|
|
#[cfg(not(target_arch = "wasm32"))]
|
|
pub async fn catch_interrupt(&mut self) -> Result<(), SentError> {
|
|
let res = crate::wait_for_signal_and_error(self).await;
|
|
|
|
log::info!("Sending shutdown");
|
|
self.signal_shutdown().ok();
|
|
|
|
log::info!("Waiting for tasks to finish... (Press ctrl-c to force)");
|
|
self.wait_for_shutdown().await;
|
|
|
|
res
|
|
}
|
|
|
|
pub fn subscribe(&self) -> TaskClient {
|
|
let task_client = TaskClient::new(
|
|
self.notify_rx
|
|
.as_ref()
|
|
.expect("Unable to subscribe to shutdown notifier that is already shutdown")
|
|
.clone(),
|
|
self.task_return_error_tx.clone(),
|
|
self.task_drop_tx.clone(),
|
|
self.task_status_tx.clone(),
|
|
);
|
|
|
|
if let Some(name) = &self.name {
|
|
task_client.named(format!("{name}-child"))
|
|
} else {
|
|
task_client
|
|
}
|
|
}
|
|
|
|
pub fn subscribe_named<S: Into<String>>(&self, suffix: S) -> TaskClient {
|
|
let task_client = self.subscribe();
|
|
let suffix = suffix.into();
|
|
let child_name = if let Some(base) = &self.name {
|
|
format!("{base}-{suffix}")
|
|
} else {
|
|
format!("unknown-{suffix}")
|
|
};
|
|
task_client.named(child_name)
|
|
}
|
|
|
|
pub fn signal_shutdown(&self) -> Result<(), SendError<()>> {
|
|
self.notify_tx.send(())
|
|
}
|
|
|
|
pub async fn start_status_listener(
|
|
&mut self,
|
|
mut sender: StatusSender,
|
|
start_status: TaskStatus,
|
|
) {
|
|
// Announce that we are operational. This means that in the application where this is used,
|
|
// everything is up and running and ready to go.
|
|
if let Err(msg) = sender.send(Box::new(start_status)).await {
|
|
log::error!("Error sending status message: {msg}");
|
|
};
|
|
|
|
if let Some(mut task_status_rx) = self.task_status_rx.take() {
|
|
log::info!("Starting status message listener");
|
|
crate::spawn::spawn_future(async move {
|
|
loop {
|
|
if let Some(msg) = task_status_rx.next().await {
|
|
log::trace!("Got msg: {msg}");
|
|
if let Err(msg) = sender.send(msg).await {
|
|
log::error!("Error sending status message: {msg}");
|
|
}
|
|
} else {
|
|
log::trace!("Stopping since channel closed");
|
|
break;
|
|
}
|
|
}
|
|
log::debug!("Status listener: Exiting");
|
|
});
|
|
}
|
|
}
|
|
|
|
// used for compatibility with the ShutdownManager
|
|
#[cfg(not(target_arch = "wasm32"))]
|
|
pub(crate) fn task_return_error_rx(&mut self) -> ErrorReceiver {
|
|
self.task_return_error_rx
|
|
.take()
|
|
.expect("unable to get error channel: attempt to wait twice?")
|
|
}
|
|
|
|
#[cfg(not(target_arch = "wasm32"))]
|
|
pub(crate) fn task_drop_rx(&mut self) -> ErrorReceiver {
|
|
self.task_drop_rx
|
|
.take()
|
|
.expect("unable to get task drop channel: attempt to wait twice?")
|
|
}
|
|
|
|
pub async fn wait_for_error(&mut self) -> Option<SentError> {
|
|
let mut error_rx = self
|
|
.task_return_error_rx
|
|
.take()
|
|
.expect("Unable to wait for error: attempt to wait twice?");
|
|
let mut drop_rx = self
|
|
.task_drop_rx
|
|
.take()
|
|
.expect("Unable to wait for error: attempt to wait twice?");
|
|
|
|
// During an error we are likely like to be swamped with drop notifications as well, this
|
|
// is a crude way to give priority to real errors (if there are any).
|
|
let drop_rx = drop_rx.recv().then(|msg| async move {
|
|
sleep(Duration::from_millis(50)).await;
|
|
msg
|
|
});
|
|
|
|
tokio::select! {
|
|
msg = error_rx.recv() => msg,
|
|
msg = drop_rx => msg
|
|
}
|
|
}
|
|
|
|
pub async fn wait_for_graceful_shutdown(&mut self) {
|
|
if let Some(notify_rx) = self.notify_rx.take() {
|
|
drop(notify_rx);
|
|
}
|
|
self.notify_tx.closed().await
|
|
}
|
|
|
|
pub async fn wait_for_shutdown(&mut self) {
|
|
log::debug!("Waiting for shutdown");
|
|
if let Some(notify_rx) = self.notify_rx.take() {
|
|
drop(notify_rx);
|
|
}
|
|
|
|
#[cfg(not(target_arch = "wasm32"))]
|
|
let interrupt_future = tokio::signal::ctrl_c();
|
|
|
|
#[cfg(target_arch = "wasm32")]
|
|
let interrupt_future = futures::future::pending::<()>();
|
|
|
|
let wait_future = sleep(Duration::from_secs(self.shutdown_timer_secs));
|
|
|
|
tokio::select! {
|
|
_ = self.notify_tx.closed() => {
|
|
log::info!("All registered tasks succesfully shutdown");
|
|
},
|
|
_ = interrupt_future => {
|
|
log::info!("Forcing shutdown");
|
|
}
|
|
_ = wait_future => {
|
|
log::info!("Timeout reached, forcing shutdown");
|
|
},
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Listen for shutdown notifications, and can send error and status messages back to the
|
|
/// `TaskManager`
|
|
#[derive(Debug)]
|
|
#[deprecated(note = "use ShutdownToken instead")]
|
|
pub struct TaskClient {
|
|
// optional name assigned to the shutdown handle
|
|
name: Option<String>,
|
|
|
|
// If a shutdown notification has been registered
|
|
// the reason for having an atomic here is to be able to cheat and modify that value whilst
|
|
// holding an immutable reference to the `TaskClient`.
|
|
// note: using `Relaxed` ordering everywhere is fine since it's not shared between threads
|
|
shutdown: AtomicBool,
|
|
|
|
// Listen for shutdown notifications, as well as a mechanism to report back that we have
|
|
// finished (the receiver is closed).
|
|
notify: watch::Receiver<()>,
|
|
|
|
// Send back error if we stopped
|
|
return_error: ErrorSender,
|
|
|
|
// Also notify if we dropped without shutdown being registered
|
|
drop_error: ErrorSender,
|
|
|
|
// Send non-exit messages
|
|
status_msg: StatusSender,
|
|
|
|
// The current operating mode
|
|
mode: ClientOperatingMode,
|
|
}
|
|
|
|
#[allow(deprecated)]
|
|
impl Clone for TaskClient {
|
|
fn clone(&self) -> Self {
|
|
// make sure to not accidentally overflow the stack if we keep cloning the handle
|
|
let name = if let Some(name) = &self.name {
|
|
if name != Self::OVERFLOW_NAME && name.len() < Self::MAX_NAME_LENGTH {
|
|
Some(format!("{name}-child"))
|
|
} else {
|
|
Some(Self::OVERFLOW_NAME.to_string())
|
|
}
|
|
} else {
|
|
None
|
|
};
|
|
|
|
log::debug!("Cloned task client: {name:?}");
|
|
|
|
TaskClient {
|
|
name,
|
|
shutdown: AtomicBool::new(self.shutdown.load(Ordering::Relaxed)),
|
|
notify: self.notify.clone(),
|
|
return_error: self.return_error.clone(),
|
|
drop_error: self.drop_error.clone(),
|
|
status_msg: self.status_msg.clone(),
|
|
mode: self.mode.clone(),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[allow(deprecated)]
|
|
impl TaskClient {
|
|
const MAX_NAME_LENGTH: usize = 128;
|
|
const OVERFLOW_NAME: &'static str = "reached maximum TaskClient children name depth";
|
|
|
|
const SHUTDOWN_TIMEOUT_WAITING_FOR_SIGNAL_ON_EXIT: Duration = Duration::from_secs(10);
|
|
|
|
fn new(
|
|
notify: watch::Receiver<()>,
|
|
return_error: ErrorSender,
|
|
drop_error: ErrorSender,
|
|
status_msg: StatusSender,
|
|
) -> TaskClient {
|
|
TaskClient {
|
|
name: None,
|
|
shutdown: AtomicBool::new(false),
|
|
notify,
|
|
return_error,
|
|
drop_error,
|
|
status_msg,
|
|
mode: ClientOperatingMode::Listening,
|
|
}
|
|
}
|
|
|
|
// TODO: not convinced about the name...
|
|
pub fn fork<S: Into<String>>(&self, child_suffix: S) -> Self {
|
|
let mut child = self.clone();
|
|
let suffix = child_suffix.into();
|
|
let child_name = if let Some(base) = &self.name {
|
|
format!("{base}-{suffix}")
|
|
} else {
|
|
format!("unknown-{suffix}")
|
|
};
|
|
|
|
log::debug!("Forked task client: {child_name}");
|
|
child.name = Some(child_name);
|
|
child
|
|
}
|
|
|
|
// just a convenience wrapper for including the shutdown name when logging
|
|
// I really didn't want to create macros for that... because that seemed like an overkill.
|
|
// but I guess it would have resolved needing to call `format!` for additional msg arguments
|
|
fn log<S: Into<String>>(&self, level: Level, msg: S) {
|
|
let msg = msg.into();
|
|
|
|
let target = &if let Some(name) = &self.name {
|
|
format!("TaskClient-{name}")
|
|
} else {
|
|
"unnamed-TaskClient".to_string()
|
|
};
|
|
|
|
log!(target: target, level, "{}", format_args!("[{target}] {msg}"))
|
|
}
|
|
|
|
#[must_use]
|
|
pub fn named<S: Into<String>>(mut self, name: S) -> Self {
|
|
self.name = Some(name.into());
|
|
self
|
|
}
|
|
|
|
#[must_use]
|
|
pub fn with_suffix<S: Into<String>>(self, suffix: S) -> Self {
|
|
let suffix = suffix.into();
|
|
let name = if let Some(base) = &self.name {
|
|
format!("{base}-{suffix}")
|
|
} else {
|
|
format!("unknown-{suffix}")
|
|
};
|
|
log::debug!("Renamed task client: {name}");
|
|
self.named(name)
|
|
}
|
|
|
|
// Create a dummy that will never report that we should shutdown.
|
|
pub fn dummy() -> TaskClient {
|
|
let (_notify_tx, notify_rx) = watch::channel(());
|
|
let (task_halt_tx, _task_halt_rx) = mpsc::unbounded_channel();
|
|
let (task_drop_tx, _task_drop_rx) = mpsc::unbounded_channel();
|
|
let (task_status_tx, _task_status_rx) = futures::channel::mpsc::channel(128);
|
|
TaskClient {
|
|
name: None,
|
|
shutdown: AtomicBool::new(false),
|
|
notify: notify_rx,
|
|
return_error: task_halt_tx,
|
|
drop_error: task_drop_tx,
|
|
status_msg: task_status_tx,
|
|
mode: ClientOperatingMode::Dummy,
|
|
}
|
|
}
|
|
|
|
pub fn is_dummy(&self) -> bool {
|
|
self.mode.is_dummy()
|
|
}
|
|
|
|
pub fn is_shutdown(&self) -> bool {
|
|
if self.mode.is_dummy() {
|
|
false
|
|
} else {
|
|
self.shutdown.load(Ordering::Relaxed)
|
|
}
|
|
}
|
|
|
|
pub async fn recv(&mut self) {
|
|
if self.mode.is_dummy() {
|
|
return pending().await;
|
|
}
|
|
if self.shutdown.load(Ordering::Relaxed) {
|
|
return;
|
|
}
|
|
let _ = self.notify.changed().await;
|
|
self.shutdown.store(true, Ordering::Relaxed);
|
|
}
|
|
|
|
pub async fn recv_with_delay(&mut self) {
|
|
self.recv()
|
|
.then(|msg| async move {
|
|
sleep(Duration::from_secs(2)).await;
|
|
msg
|
|
})
|
|
.await
|
|
}
|
|
|
|
// legacy code
|
|
#[allow(clippy::panic)]
|
|
pub async fn recv_timeout(&mut self) {
|
|
if self.mode.is_dummy() {
|
|
return pending().await;
|
|
}
|
|
|
|
if let Err(timeout) = timeout(
|
|
Self::SHUTDOWN_TIMEOUT_WAITING_FOR_SIGNAL_ON_EXIT,
|
|
self.recv(),
|
|
)
|
|
.await
|
|
{
|
|
self.log(Level::Error, "Task stopped without shutdown called");
|
|
panic!("{:?}: {timeout}", self.name)
|
|
}
|
|
}
|
|
|
|
pub fn is_shutdown_poll(&self) -> bool {
|
|
if self.mode.is_dummy() {
|
|
return false;
|
|
}
|
|
if self.shutdown.load(Ordering::Relaxed) {
|
|
return true;
|
|
}
|
|
match self.notify.has_changed() {
|
|
Ok(has_changed) => {
|
|
if has_changed {
|
|
self.shutdown.store(true, Ordering::Relaxed);
|
|
}
|
|
has_changed
|
|
}
|
|
Err(err) => {
|
|
self.log(Level::Error, format!("Polling shutdown failed: {err}"));
|
|
self.log(Level::Error, "Assuming this means we should shutdown...");
|
|
|
|
true
|
|
}
|
|
}
|
|
}
|
|
|
|
// This listener should to *not* notify the ShutdownNotifier to shutdown when dropped. For
|
|
// example when we clone the listener for a task handling connections, we often want to drop
|
|
// without signal failure.
|
|
pub fn disarm(&mut self) {
|
|
self.mode.set_should_not_signal_on_drop();
|
|
}
|
|
|
|
pub fn rearm(&mut self) {
|
|
self.mode.set_should_signal_on_drop();
|
|
}
|
|
|
|
pub fn send_we_stopped(&mut self, err: SentError) {
|
|
if self.mode.is_dummy() {
|
|
return;
|
|
}
|
|
|
|
self.log(Level::Trace, format!("Notifying we stopped: {err}"));
|
|
|
|
if self.return_error.send(err).is_err() {
|
|
self.log(Level::Error, "failed to send back error message");
|
|
}
|
|
}
|
|
|
|
pub fn send_status_msg(&mut self, msg: SentStatus) {
|
|
if self.mode.is_dummy() {
|
|
return;
|
|
}
|
|
// Since it's not always the case that anyone is listening, just try send and ignore any
|
|
// failures.
|
|
self.status_msg.try_send(msg).ok();
|
|
}
|
|
}
|
|
|
|
#[allow(deprecated)]
|
|
impl Drop for TaskClient {
|
|
fn drop(&mut self) {
|
|
if !self.mode.should_signal_on_drop() {
|
|
self.log(
|
|
Level::Trace,
|
|
"the task client is getting dropped but instructed to not signal: this is expected during client shutdown",
|
|
);
|
|
return;
|
|
} else {
|
|
self.log(
|
|
Level::Debug,
|
|
"the task client is getting dropped: this is expected during client shutdown",
|
|
);
|
|
}
|
|
|
|
if !self.is_shutdown_poll() {
|
|
self.log(Level::Trace, "Notifying stop on unexpected drop");
|
|
|
|
// If we can't send, well then there is not much to do
|
|
self.drop_error
|
|
.send(Box::new(TaskError::UnexpectedHalt {
|
|
shutdown_name: self.name.clone(),
|
|
}))
|
|
.ok();
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
enum ClientOperatingMode {
|
|
// Normal operations
|
|
Listening,
|
|
// Normal operations, but we don't report back if the we stop by getting dropped.
|
|
ListeningButDontReportHalt,
|
|
// Dummy mode, for when we don't do anything at all.
|
|
Dummy,
|
|
}
|
|
|
|
impl ClientOperatingMode {
|
|
fn is_dummy(&self) -> bool {
|
|
self == &ClientOperatingMode::Dummy
|
|
}
|
|
|
|
fn should_signal_on_drop(&self) -> bool {
|
|
match self {
|
|
ClientOperatingMode::Listening => true,
|
|
ClientOperatingMode::ListeningButDontReportHalt | ClientOperatingMode::Dummy => false,
|
|
}
|
|
}
|
|
|
|
fn set_should_signal_on_drop(&mut self) {
|
|
use ClientOperatingMode::{Dummy, Listening, ListeningButDontReportHalt};
|
|
*self = match &self {
|
|
ListeningButDontReportHalt | Listening => Listening,
|
|
Dummy => Dummy,
|
|
};
|
|
}
|
|
|
|
fn set_should_not_signal_on_drop(&mut self) {
|
|
use ClientOperatingMode::{Dummy, Listening, ListeningButDontReportHalt};
|
|
*self = match &self {
|
|
ListeningButDontReportHalt | Listening => ListeningButDontReportHalt,
|
|
Dummy => Dummy,
|
|
};
|
|
}
|
|
}
|
|
|
|
#[deprecated]
|
|
#[allow(deprecated)]
|
|
#[derive(Debug)]
|
|
pub enum TaskHandle {
|
|
/// Full [`TaskManager`] that was created by the underlying task.
|
|
Internal(TaskManager),
|
|
|
|
/// `[TaskClient]` that was passed from an external task, that controls the shutdown process.
|
|
External(TaskClient),
|
|
}
|
|
|
|
#[allow(deprecated)]
|
|
impl From<TaskManager> for TaskHandle {
|
|
fn from(value: TaskManager) -> Self {
|
|
TaskHandle::Internal(value)
|
|
}
|
|
}
|
|
|
|
#[allow(deprecated)]
|
|
impl From<TaskClient> for TaskHandle {
|
|
fn from(value: TaskClient) -> Self {
|
|
TaskHandle::External(value)
|
|
}
|
|
}
|
|
|
|
#[allow(deprecated)]
|
|
impl Default for TaskHandle {
|
|
fn default() -> Self {
|
|
TaskHandle::Internal(TaskManager::default())
|
|
}
|
|
}
|
|
|
|
#[allow(deprecated)]
|
|
impl TaskHandle {
|
|
#[must_use]
|
|
pub fn name_if_unnamed<S: Into<String>>(self, name: S) -> Self {
|
|
match self {
|
|
TaskHandle::Internal(task_manager) => {
|
|
if task_manager.name.is_none() {
|
|
TaskHandle::Internal(task_manager.named(name))
|
|
} else {
|
|
TaskHandle::Internal(task_manager)
|
|
}
|
|
}
|
|
TaskHandle::External(task_client) => {
|
|
if task_client.name.is_none() {
|
|
TaskHandle::External(task_client.named(name))
|
|
} else {
|
|
TaskHandle::External(task_client)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[must_use]
|
|
pub fn named<S: Into<String>>(self, name: S) -> Self {
|
|
match self {
|
|
TaskHandle::Internal(task_manager) => TaskHandle::Internal(task_manager.named(name)),
|
|
TaskHandle::External(task_client) => TaskHandle::External(task_client.named(name)),
|
|
}
|
|
}
|
|
|
|
pub fn fork<S: Into<String>>(&self, child_suffix: S) -> TaskClient {
|
|
match self {
|
|
TaskHandle::External(shutdown) => shutdown.fork(child_suffix),
|
|
TaskHandle::Internal(shutdown) => shutdown.subscribe_named(child_suffix),
|
|
}
|
|
}
|
|
|
|
pub fn get_handle(&self) -> TaskClient {
|
|
match self {
|
|
TaskHandle::External(shutdown) => shutdown.clone(),
|
|
TaskHandle::Internal(shutdown) => shutdown.subscribe(),
|
|
}
|
|
}
|
|
|
|
pub fn try_into_task_manager(self) -> Option<TaskManager> {
|
|
match self {
|
|
TaskHandle::External(_) => None,
|
|
TaskHandle::Internal(shutdown) => Some(shutdown),
|
|
}
|
|
}
|
|
|
|
#[cfg(not(target_arch = "wasm32"))]
|
|
pub async fn wait_for_shutdown(self) -> Result<(), SentError> {
|
|
match self {
|
|
TaskHandle::Internal(mut task_manager) => task_manager.catch_interrupt().await,
|
|
TaskHandle::External(mut task_client) => {
|
|
task_client.recv().await;
|
|
Ok(())
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[tokio::test]
|
|
#[allow(deprecated)]
|
|
async fn signal_shutdown() {
|
|
let shutdown = TaskManager::default();
|
|
let mut listener = shutdown.subscribe();
|
|
|
|
let task = tokio::spawn(async move {
|
|
tokio::select! {
|
|
_ = listener.recv() => 42,
|
|
}
|
|
});
|
|
|
|
shutdown.signal_shutdown().unwrap();
|
|
assert_eq!(task.await.unwrap(), 42);
|
|
}
|
|
}
|