use super::{
AbortHandle, AssertUnwindSafe, Cell, Context, DriverControl, DriverTask, Duration, Future,
FutureExt, LocalBoxFuture, LocalTask, ModuleDependencies, ModuleLifecyclePhase, Pin, Poll, Rc,
RefCell, RuntimeDriver, RuntimeFailure, SpawnError, TaskOutcome, oneshot, wait_until,
};
#[derive(Clone, Debug)]
pub struct AppReadyGate {
pub(super) state: Rc<AppReadyState>,
}
#[derive(Debug)]
pub(super) struct AppReadyState {
pub(super) open: Cell<bool>,
pub(super) waiters: RefCell<Vec<oneshot::Sender<()>>>,
}
impl AppReadyGate {
pub fn new() -> Self {
Self {
state: Rc::new(AppReadyState {
open: Cell::new(false),
waiters: RefCell::new(Vec::new()),
}),
}
}
pub fn is_open(&self) -> bool {
self.state.open.get()
}
pub fn wait(&self) -> LocalBoxFuture<'static, ()> {
if self.is_open() {
return Box::pin(futures::future::ready(()));
}
let (wakeup, waiter) = oneshot::channel();
self.state.waiters.borrow_mut().push(wakeup);
Box::pin(async move {
let _ = waiter.await;
})
}
pub(super) fn open(&self) {
if self.state.open.replace(true) {
return;
}
for waiter in self.state.waiters.borrow_mut().drain(..) {
let _ = waiter.send(());
}
}
}
impl Default for AppReadyGate {
fn default() -> Self {
Self::new()
}
}
#[derive(Clone, Debug)]
pub struct AppAdmission {
pub(super) state: Rc<AppAdmissionState>,
}
#[derive(Debug)]
pub(super) struct AppAdmissionState {
pub(super) open: Cell<bool>,
}
impl AppAdmission {
pub(super) fn new() -> Self {
Self {
state: Rc::new(AppAdmissionState {
open: Cell::new(false),
}),
}
}
pub fn is_open(&self) -> bool {
self.state.open.get()
}
pub fn is_closed(&self) -> bool {
!self.is_open()
}
pub(super) fn open(&self) {
self.state.open.set(true);
}
pub(super) fn close(&self) {
self.state.open.set(false);
}
}
#[derive(Clone, Debug)]
pub struct CancellationToken {
pub(super) state: Rc<CancellationState>,
}
#[derive(Debug)]
pub(super) struct CancellationState {
pub(super) cancelled: Cell<bool>,
pub(super) next_waiter_id: Cell<usize>,
pub(super) waiters: RefCell<Vec<(usize, oneshot::Sender<()>)>>,
}
impl CancellationToken {
pub fn new() -> Self {
Self {
state: Rc::new(CancellationState {
cancelled: Cell::new(false),
next_waiter_id: Cell::new(0),
waiters: RefCell::new(Vec::new()),
}),
}
}
pub fn is_cancelled(&self) -> bool {
self.state.cancelled.get()
}
pub fn cancelled(&self) -> LocalBoxFuture<'static, ()> {
if self.is_cancelled() {
return Box::pin(futures::future::ready(()));
}
let (wakeup, waiter) = oneshot::channel();
let waiter_id = self.state.next_waiter_id.get();
self.state.next_waiter_id.set(waiter_id.saturating_add(1));
self.state.waiters.borrow_mut().push((waiter_id, wakeup));
Box::pin(CancellationWaiter {
state: self.state.clone(),
waiter_id,
receiver: waiter,
registered: true,
})
}
pub fn cancel(&self) {
if self.state.cancelled.replace(true) {
return;
}
for (_, waiter) in self.state.waiters.borrow_mut().drain(..) {
let _ = waiter.send(());
}
}
}
#[derive(Debug)]
pub(super) struct CancellationWaiter {
pub(super) state: Rc<CancellationState>,
pub(super) waiter_id: usize,
pub(super) receiver: oneshot::Receiver<()>,
pub(super) registered: bool,
}
impl Future for CancellationWaiter {
type Output = ();
fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
match Pin::new(&mut self.receiver).poll(context) {
Poll::Ready(_) => {
self.registered = false;
Poll::Ready(())
}
Poll::Pending => Poll::Pending,
}
}
}
impl Drop for CancellationWaiter {
fn drop(&mut self) {
if !self.registered {
return;
}
self.state
.waiters
.borrow_mut()
.retain(|(waiter_id, _)| *waiter_id != self.waiter_id);
}
}
impl Default for CancellationToken {
fn default() -> Self {
Self::new()
}
}
pub type ResourceFuture = LocalBoxFuture<'static, Result<(), RuntimeFailure>>;
pub trait ManagedResource: std::fmt::Debug + 'static {
fn release(&self) -> ResourceFuture;
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ResourceRegistrationError {
ScopeClosed,
}
pub(super) struct ManagedResourceEntry {
pub(super) resource: Rc<dyn ManagedResource>,
pub(super) release: RefCell<ManagedResourceRelease>,
}
pub(super) enum ManagedResourceRelease {
Pending,
Running(ResourceFuture),
Complete(Result<(), RuntimeFailure>),
}
impl std::fmt::Debug for ManagedResourceEntry {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let state = match &*self.release.borrow() {
ManagedResourceRelease::Pending => "pending",
ManagedResourceRelease::Running(_) => "running",
ManagedResourceRelease::Complete(Ok(())) => "released",
ManagedResourceRelease::Complete(Err(_)) => "failed",
};
formatter
.debug_struct("ManagedResourceEntry")
.field("release", &state)
.finish_non_exhaustive()
}
}
#[derive(Clone, Debug)]
pub struct ManagedResourceHandle {
pub(super) entry: Rc<ManagedResourceEntry>,
}
impl ManagedResourceHandle {
pub fn is_released(&self) -> bool {
matches!(
&*self.entry.release.borrow(),
ManagedResourceRelease::Complete(_)
)
}
pub async fn release(&self) -> Result<(), RuntimeFailure> {
ManagedResourceReleaseOperation {
entry: self.entry.clone(),
}
.await
}
}
pub(super) struct ManagedResourceReleaseOperation {
pub(super) entry: Rc<ManagedResourceEntry>,
}
impl Future for ManagedResourceReleaseOperation {
type Output = Result<(), RuntimeFailure>;
fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
let mut release = self.entry.release.borrow_mut();
if matches!(*release, ManagedResourceRelease::Pending) {
*release = ManagedResourceRelease::Running(self.entry.resource.release());
}
match &mut *release {
ManagedResourceRelease::Running(future) => match future.as_mut().poll(context) {
Poll::Ready(result) => {
*release = ManagedResourceRelease::Complete(result.clone());
Poll::Ready(result)
}
Poll::Pending => Poll::Pending,
},
ManagedResourceRelease::Complete(result) => Poll::Ready(result.clone()),
ManagedResourceRelease::Pending => unreachable!("pending release was started"),
}
}
}
#[derive(Clone)]
pub struct ManagedResourceScope {
pub(super) state: Rc<ManagedResourceScopeState>,
}
#[derive(Debug, Default)]
pub(super) struct ManagedResourceScopeState {
pub(super) resources: RefCell<Vec<ManagedResourceHandle>>,
pub(super) closed: Cell<bool>,
}
impl std::fmt::Debug for ManagedResourceScope {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("ManagedResourceScope")
.field("resource_count", &self.resource_count())
.finish()
}
}
impl ManagedResourceScope {
pub(super) fn new() -> Self {
Self {
state: Rc::new(ManagedResourceScopeState::default()),
}
}
pub fn register(
&self,
resource: impl ManagedResource,
) -> Result<ManagedResourceHandle, ResourceRegistrationError> {
if self.state.closed.get() {
return Err(ResourceRegistrationError::ScopeClosed);
}
let handle = ManagedResourceHandle {
entry: Rc::new(ManagedResourceEntry {
resource: Rc::new(resource),
release: RefCell::new(ManagedResourceRelease::Pending),
}),
};
self.state.resources.borrow_mut().push(handle.clone());
Ok(handle)
}
pub fn resource_count(&self) -> usize {
self.state
.resources
.borrow()
.iter()
.filter(|resource| !resource.is_released())
.count()
}
pub(super) fn close(&self) {
self.state.closed.set(true);
}
pub(super) async fn release_all(&self) -> Option<RuntimeFailure> {
let resources = std::mem::take(&mut *self.state.resources.borrow_mut());
let mut first_error = None;
for resource in resources {
if let Err(error) = resource.release().await
&& first_error.is_none()
{
first_error = Some(error);
}
}
first_error
}
pub(super) async fn release_all_until(
&self,
driver: &DriverControl,
deadline: Duration,
) -> Result<Option<RuntimeFailure>, ()> {
let resources = std::mem::take(&mut *self.state.resources.borrow_mut());
let mut first_error = None;
for (index, resource) in resources.iter().enumerate() {
match wait_until(driver, deadline, resource.release()).await {
Some(Ok(())) => {}
Some(Err(error)) => {
if first_error.is_none() {
first_error = Some(error);
}
}
None => {
self.state
.resources
.borrow_mut()
.extend(resources.into_iter().skip(index));
return Err(());
}
}
}
Ok(first_error)
}
}
#[derive(Clone, Debug)]
pub struct ManagedTask {
pub(super) task: Rc<RefCell<Option<DriverTask>>>,
pub(super) abort: AbortHandle,
pub(super) failed: Rc<Cell<bool>>,
}
impl ManagedTask {
pub(super) fn from_driver_task(task: DriverTask) -> Self {
Self {
abort: task.abort_handle(),
task: Rc::new(RefCell::new(Some(task))),
failed: Rc::new(Cell::new(false)),
}
}
pub fn cancel(&self) {
self.abort.abort();
}
pub(super) async fn join(&self) -> TaskOutcome {
let task = self.task.borrow_mut().take();
if let Some(task) = task {
let outcome = task.await;
if self.failed.get() {
TaskOutcome::Failed
} else {
outcome
}
} else if self.failed.get() {
TaskOutcome::Failed
} else {
TaskOutcome::Completed
}
}
}
#[derive(Debug)]
pub enum ManagedTaskError {
ScopeClosed,
Driver(SpawnError),
}
impl From<SpawnError> for ManagedTaskError {
fn from(error: SpawnError) -> Self {
Self::Driver(error)
}
}
#[derive(Clone)]
pub struct ManagedTaskScope {
pub(super) spawn: Rc<dyn Fn(LocalTask) -> Result<DriverTask, SpawnError>>,
pub(super) state: Rc<ManagedTaskScopeState>,
}
pub(super) struct ManagedTaskScopeState {
pub(super) tasks: RefCell<Vec<ManagedTask>>,
pub(super) closed: Cell<bool>,
pub(super) cancellation: CancellationToken,
pub(super) failure_handler: RefCell<Option<Rc<dyn Fn()>>>,
pub(super) unreported_failure: Cell<bool>,
}
impl Default for ManagedTaskScopeState {
fn default() -> Self {
Self {
tasks: RefCell::new(Vec::new()),
closed: Cell::new(false),
cancellation: CancellationToken::new(),
failure_handler: RefCell::new(None),
unreported_failure: Cell::new(false),
}
}
}
impl std::fmt::Debug for ManagedTaskScopeState {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("ManagedTaskScopeState")
.field("task_count", &self.tasks.borrow().len())
.field("closed", &self.closed.get())
.field("unreported_failure", &self.unreported_failure.get())
.finish_non_exhaustive()
}
}
impl std::fmt::Debug for ManagedTaskScope {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("ManagedTaskScope")
.field("task_count", &self.task_count())
.finish()
}
}
impl ManagedTaskScope {
pub(super) fn new<D: RuntimeDriver>(driver: &D) -> Self {
let spawner = driver.clone();
Self {
spawn: Rc::new(move |task| spawner.spawn_local(task)),
state: Rc::new(ManagedTaskScopeState::default()),
}
}
pub(super) fn new_from_driver_control(driver: &DriverControl) -> Self {
let spawn = driver.spawn_local.clone();
Self {
spawn,
state: Rc::new(ManagedTaskScopeState::default()),
}
}
pub fn spawn_local(&self, task: LocalTask) -> Result<ManagedTask, ManagedTaskError> {
if self.state.closed.get() {
return Err(ManagedTaskError::ScopeClosed);
}
let failed = Rc::new(Cell::new(false));
let task_failed = failed.clone();
let state = self.state.clone();
let monitored = Box::pin(async move {
if AssertUnwindSafe(task).catch_unwind().await.is_err() {
task_failed.set(true);
state.report_failure();
}
});
let driver_task = (self.spawn)(monitored)?;
let handle = ManagedTask {
failed,
..ManagedTask::from_driver_task(driver_task)
};
self.state.tasks.borrow_mut().push(handle.clone());
Ok(handle)
}
pub fn task_count(&self) -> usize {
self.state.tasks.borrow().len()
}
pub fn cancellation(&self) -> CancellationToken {
self.state.cancellation.clone()
}
pub(super) fn close(&self) {
self.state.closed.set(true);
self.state.cancellation.cancel();
}
pub(super) fn set_failure_handler(&self, handler: &Rc<dyn Fn()>) {
self.state.failure_handler.replace(Some(handler.clone()));
if self.state.unreported_failure.replace(false) {
handler();
}
}
pub(super) fn cancel(&self) {
self.state.cancellation.cancel();
}
pub(super) fn abort_all(&self) {
for task in self.state.tasks.borrow().iter() {
task.cancel();
}
}
pub(super) async fn cancel_all(&self) {
self.close();
let tasks = std::mem::take(&mut *self.state.tasks.borrow_mut());
for task in tasks {
task.cancel();
let _ = task.join().await;
}
}
pub(super) async fn drain_until(&self, driver: &DriverControl, deadline: Duration) -> bool {
self.cancel();
let tasks = std::mem::take(&mut *self.state.tasks.borrow_mut());
for (index, task) in tasks.iter().enumerate() {
if wait_until(driver, deadline, task.join()).await.is_none() {
for pending in tasks.iter().skip(index) {
pending.cancel();
}
return false;
}
}
true
}
}
impl ManagedTaskScopeState {
pub(super) fn report_failure(&self) {
let handler = self.failure_handler.borrow().clone();
if let Some(handler) = handler {
handler();
} else {
self.unreported_failure.set(true);
}
}
}
#[derive(Clone, Debug)]
pub struct PrepareContext {
pub(super) instance_key: String,
pub(super) entrypoint: String,
pub(super) configuration: String,
pub(super) dependencies: ModuleDependencies,
pub(super) resources: ManagedResourceScope,
pub(super) cancellation: CancellationToken,
pub(super) admission: AppAdmission,
}
impl PrepareContext {
pub fn instance_key(&self) -> &str {
&self.instance_key
}
pub fn entrypoint(&self) -> &str {
&self.entrypoint
}
pub fn configuration(&self) -> &str {
&self.configuration
}
pub const fn phase(&self) -> ModuleLifecyclePhase {
ModuleLifecyclePhase::Prepare
}
pub fn dependencies(&self) -> &ModuleDependencies {
&self.dependencies
}
pub fn resources(&self) -> &ManagedResourceScope {
&self.resources
}
pub fn cancellation(&self) -> CancellationToken {
self.cancellation.clone()
}
pub fn admission(&self) -> AppAdmission {
self.admission.clone()
}
}
#[derive(Clone, Debug)]
pub struct ActivateContext {
pub(super) instance_key: String,
pub(super) dependencies: ModuleDependencies,
pub(super) ready_gate: AppReadyGate,
pub(super) tasks: ManagedTaskScope,
pub(super) resources: ManagedResourceScope,
pub(super) cancellation: CancellationToken,
pub(super) admission: AppAdmission,
}
impl ActivateContext {
pub fn instance_key(&self) -> &str {
&self.instance_key
}
pub const fn phase(&self) -> ModuleLifecyclePhase {
ModuleLifecyclePhase::Activate
}
pub fn dependencies(&self) -> &ModuleDependencies {
&self.dependencies
}
pub fn ready_gate(&self) -> AppReadyGate {
self.ready_gate.clone()
}
pub fn readiness(&self) -> ReadinessContext {
ReadinessContext {
instance_key: self.instance_key.clone(),
dependencies: self.dependencies.clone(),
ready_gate: self.ready_gate.clone(),
tasks: self.tasks.clone(),
resources: self.resources.clone(),
cancellation: self.cancellation.clone(),
admission: self.admission.clone(),
}
}
pub fn tasks(&self) -> &ManagedTaskScope {
&self.tasks
}
pub fn resources(&self) -> &ManagedResourceScope {
&self.resources
}
pub fn cancellation(&self) -> CancellationToken {
self.cancellation.clone()
}
pub fn admission(&self) -> AppAdmission {
self.admission.clone()
}
}
#[derive(Clone, Debug)]
pub struct ReadinessContext {
pub(super) instance_key: String,
pub(super) dependencies: ModuleDependencies,
pub(super) ready_gate: AppReadyGate,
pub(super) tasks: ManagedTaskScope,
pub(super) resources: ManagedResourceScope,
pub(super) cancellation: CancellationToken,
pub(super) admission: AppAdmission,
}
impl ReadinessContext {
pub fn instance_key(&self) -> &str {
&self.instance_key
}
pub const fn phase(&self) -> ModuleLifecyclePhase {
ModuleLifecyclePhase::Ready
}
pub fn dependencies(&self) -> &ModuleDependencies {
&self.dependencies
}
pub fn ready_gate(&self) -> AppReadyGate {
self.ready_gate.clone()
}
pub fn wait(&self) -> LocalBoxFuture<'static, ()> {
self.ready_gate.wait()
}
pub fn is_open(&self) -> bool {
self.ready_gate.is_open()
}
pub fn tasks(&self) -> &ManagedTaskScope {
&self.tasks
}
pub fn resources(&self) -> &ManagedResourceScope {
&self.resources
}
pub fn cancellation(&self) -> CancellationToken {
self.cancellation.clone()
}
pub fn is_accepting(&self) -> bool {
self.admission.is_open()
}
pub fn admission(&self) -> AppAdmission {
self.admission.clone()
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum DeactivationReason {
StartupRollback,
Shutdown,
SupervisionRestart,
}
#[derive(Clone, Debug)]
pub struct DeactivateContext {
pub(super) instance_key: String,
pub(super) dependencies: ModuleDependencies,
pub(super) reason: DeactivationReason,
pub(super) tasks: ManagedTaskScope,
pub(super) resources: ManagedResourceScope,
pub(super) cancellation: CancellationToken,
pub(super) admission: AppAdmission,
}
impl DeactivateContext {
pub fn instance_key(&self) -> &str {
&self.instance_key
}
pub const fn phase(&self) -> ModuleLifecyclePhase {
ModuleLifecyclePhase::Deactivate
}
pub fn dependencies(&self) -> &ModuleDependencies {
&self.dependencies
}
pub const fn reason(&self) -> DeactivationReason {
self.reason
}
pub fn tasks(&self) -> &ManagedTaskScope {
&self.tasks
}
pub fn resources(&self) -> &ManagedResourceScope {
&self.resources
}
pub fn cancellation(&self) -> CancellationToken {
self.cancellation.clone()
}
pub fn admission(&self) -> AppAdmission {
self.admission.clone()
}
}
pub type ModuleFuture = LocalBoxFuture<'static, Result<(), RuntimeFailure>>;
pub trait ModuleLifecycle: std::fmt::Debug + 'static {
fn prepare(&self, _context: PrepareContext) -> ModuleFuture {
Box::pin(futures::future::ready(Ok(())))
}
fn activate(&self, _context: ActivateContext) -> ModuleFuture {
Box::pin(futures::future::ready(Ok(())))
}
fn deactivate(&self, _context: DeactivateContext) -> ModuleFuture {
Box::pin(futures::future::ready(Ok(())))
}
}
#[derive(Debug, Default)]
pub struct NoopModuleLifecycle;
impl ModuleLifecycle for NoopModuleLifecycle {}