use std::sync::Arc;
use tokio::sync::{mpsc, oneshot};
use super::super::SupervisorCore;
use crate::{
core::registry::{RegistryCommand, RemoveReplyRx},
error::RuntimeError,
events::{Event, EventKind},
identity::TaskId,
};
impl SupervisorCore {
pub(crate) async fn remove(&self, id: TaskId) -> Result<bool, RuntimeError> {
let reply = self.enqueue_remove_wait(id, None).await?;
Self::await_remove_reply(reply).await
}
pub(crate) async fn try_remove(&self, id: TaskId) -> Result<bool, RuntimeError> {
let reply = self.enqueue_remove(id, None)?;
Self::await_remove_reply(reply).await
}
pub(in crate::core) async fn remove_by_label(
&self,
label: Arc<str>,
) -> Result<bool, RuntimeError> {
let reply = self.enqueue_remove_by_label_wait(label).await?;
Self::await_remove_reply(reply).await
}
pub(in crate::core) async fn try_remove_by_label(
&self,
label: Arc<str>,
) -> Result<bool, RuntimeError> {
let reply = self.enqueue_remove_by_label(label)?;
Self::await_remove_reply(reply).await
}
async fn await_remove_reply(reply: RemoveReplyRx) -> Result<bool, RuntimeError> {
match reply.await {
Ok(result) => result,
Err(_) => Err(RuntimeError::ShuttingDown),
}
}
pub(in crate::core::runtime) fn enqueue_remove(
&self,
id: TaskId,
reason: Option<&'static str>,
) -> Result<RemoveReplyRx, RuntimeError> {
if self.is_shutting_down() {
return Err(RuntimeError::ShuttingDown);
}
let permit = self.cmd_tx.try_reserve().map_err(|error| match error {
mpsc::error::TrySendError::Full(()) => RuntimeError::CommandQueueFull,
mpsc::error::TrySendError::Closed(()) => RuntimeError::ShuttingDown,
})?;
let Some(_admission) = self.command_admission() else {
drop(permit);
return Err(RuntimeError::ShuttingDown);
};
Ok(self.commit_remove(permit, id, reason))
}
async fn enqueue_remove_wait(
&self,
id: TaskId,
reason: Option<&'static str>,
) -> Result<RemoveReplyRx, RuntimeError> {
if self.is_shutting_down() {
return Err(RuntimeError::ShuttingDown);
}
let permit = self
.cmd_tx
.reserve()
.await
.map_err(|_| RuntimeError::ShuttingDown)?;
let Some(_admission) = self.command_admission() else {
drop(permit);
return Err(RuntimeError::ShuttingDown);
};
Ok(self.commit_remove(permit, id, reason))
}
fn enqueue_remove_by_label(&self, label: Arc<str>) -> Result<RemoveReplyRx, RuntimeError> {
if self.is_shutting_down() {
return Err(RuntimeError::ShuttingDown);
}
let permit = self.cmd_tx.try_reserve().map_err(|error| match error {
mpsc::error::TrySendError::Full(()) => RuntimeError::CommandQueueFull,
mpsc::error::TrySendError::Closed(()) => RuntimeError::ShuttingDown,
})?;
let Some(_admission) = self.command_admission() else {
drop(permit);
return Err(RuntimeError::ShuttingDown);
};
Ok(Self::commit_remove_by_label(permit, label))
}
async fn enqueue_remove_by_label_wait(
&self,
label: Arc<str>,
) -> Result<RemoveReplyRx, RuntimeError> {
if self.is_shutting_down() {
return Err(RuntimeError::ShuttingDown);
}
let permit = self
.cmd_tx
.reserve()
.await
.map_err(|_| RuntimeError::ShuttingDown)?;
let Some(_admission) = self.command_admission() else {
drop(permit);
return Err(RuntimeError::ShuttingDown);
};
Ok(Self::commit_remove_by_label(permit, label))
}
fn commit_remove_by_label(
permit: mpsc::Permit<'_, RegistryCommand>,
label: Arc<str>,
) -> RemoveReplyRx {
let (reply, reply_rx) = oneshot::channel();
permit.send(RegistryCommand::RemoveByLabel { label, reply });
reply_rx
}
fn commit_remove(
&self,
permit: mpsc::Permit<'_, RegistryCommand>,
id: TaskId,
reason: Option<&'static str>,
) -> RemoveReplyRx {
let (reply, reply_rx) = oneshot::channel();
self.bus.publish_lazy(|| {
let mut event = Event::new(EventKind::TaskRemoveRequested).with_id(id);
if let Some(reason) = reason {
event = event.with_reason(reason);
}
event
});
permit.send(RegistryCommand::Remove { id, reply });
reply_rx
}
}