use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, Weak};
use std::task::{Context, Poll};
use std::time::Duration;
use tokio::sync::watch;
use super::Runtime;
use crate::storage::StorageManager;
use crate::{Application, Result};
type StartupFuture<A> = Pin<Box<dyn Future<Output = std::result::Result<Runtime<A>, super::Incompatible>> + Send>>;
type CleanupFuture = Pin<Box<dyn Future<Output = Result<()>> + Send>>;
type ReconcileFuture = Pin<Box<dyn Future<Output = Result<()>> + Send>>;
type InitializeFuture = Pin<Box<dyn Future<Output = std::result::Result<(), Box<dyn std::error::Error + Send + Sync>>> + Send>>;
type Initialize<A> = Box<dyn FnOnce(super::Handle<A>) -> InitializeFuture + Send>;
pub(super) type Reconciliation<A> =
Box<dyn FnOnce(super::StateManager<A>, StorageManager<crate::storage::ErasedCatalog>) -> ReconcileFuture + Send>;
pub(super) struct Initializer<A: Application> {
pub(super) timeout: Duration,
pub(super) initialize: Initialize<A>,
}
#[derive(Clone)]
pub(super) struct Cleanup {
close: Arc<dyn Fn() -> CleanupFuture + Send + Sync>,
receipt: Weak<CleanupState>,
}
pub struct Startup<A: Application> {
future: Option<StartupFuture<A>>,
cleanup: BootstrapCleanup,
}
#[derive(Clone)]
pub struct BootstrapCleanup {
shared: Arc<CleanupState>,
}
struct CleanupState {
retained: Mutex<Retained>,
status: watch::Sender<BootstrapStatus>,
}
struct Retained {
owner: Option<Cleanup>,
failure: Option<Arc<crate::Error>>,
}
pub(super) struct Boot<A: Application> {
pub(super) application: Option<super::launch::ApplicationCleanup<A>>,
pub(super) identity: Arc<()>,
pub(super) storage: Option<StorageManager<crate::storage::ErasedCatalog>>,
pub(super) runtime: Option<Runtime<A>>,
pub(super) writer: Option<tokio::task::JoinHandle<()>>,
pub(super) state: Option<super::StateManager<A>>,
armed: bool,
cleanup: BootstrapCleanup,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BootstrapStatus {
Starting,
Cleaning,
Failed,
Settled,
HandedOff,
}
impl<A: Application> Startup<A> {
pub(super) fn new(builder: super::Builder<A>, storage: Option<StorageManager<crate::storage::ErasedCatalog>>) -> Self {
let cleanup = builder.cleanup.clone();
let receipt = cleanup.clone();
let future = Box::pin(async move {
let outcome = builder.boot(storage.as_ref(), receipt.clone()).await;
receipt.settle_unowned();
outcome
});
Self { future: Some(future), cleanup }
}
}
impl<A: Application> Future for Startup<A> {
type Output = std::result::Result<Runtime<A>, super::Incompatible>;
fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
let future = self.future.as_mut().expect("startup is polled before completion");
let outcome = future.as_mut().poll(context);
if outcome.is_ready() {
self.future.take();
}
outcome
}
}
impl<A: Application> Drop for Startup<A> {
fn drop(&mut self) {
self.future.take();
self.cleanup.settle_unowned();
}
}
impl BootstrapCleanup {
pub(super) fn new() -> Self {
let (status, _) = watch::channel(BootstrapStatus::Starting);
Self {
shared: Arc::new(CleanupState {
retained: Mutex::new(Retained { owner: None, failure: None }),
status,
}),
}
}
pub fn status(&self) -> BootstrapStatus {
*self.shared.status.borrow()
}
pub fn failure(&self) -> Option<Arc<crate::Error>> {
self.shared.retained.lock().expect("bootstrap cleanup lock").failure.clone()
}
pub async fn wait(&self) -> Result<()> {
let mut status = self.shared.status.subscribe();
loop {
let current = *status.borrow_and_update();
match current {
BootstrapStatus::Settled | BootstrapStatus::HandedOff => return Ok(()),
BootstrapStatus::Failed => {
let source = {
let retained = self.shared.retained.lock().expect("bootstrap cleanup lock");
if self.status() != BootstrapStatus::Failed {
continue;
}
retained.failure.clone().expect("failed cleanup retains its cause")
};
return Err(crate::Error::TaskFailed { domain: "bootstrap cleanup", source });
}
BootstrapStatus::Starting | BootstrapStatus::Cleaning => {}
}
status.changed().await.expect("receipt retains its status owner");
}
}
pub async fn retry(&self) -> Result<()> {
let owner = {
let mut retained = self.shared.retained.lock().expect("bootstrap cleanup lock");
if self.status() == BootstrapStatus::Failed {
let owner = retained.owner.clone().expect("failed cleanup owns resources");
retained.failure = None;
self.shared.status.send_replace(BootstrapStatus::Cleaning);
Some(owner)
} else {
None
}
};
if let Some(owner) = owner {
self.run(owner);
}
self.wait().await
}
fn run(&self, owner: Cleanup) {
tokio::spawn(async move {
let _ = owner.close().await;
});
}
pub(super) fn failed(&self, owner: Option<Cleanup>, failure: Arc<crate::Error>) {
let mut retained = self.shared.retained.lock().expect("bootstrap cleanup lock");
retained.owner = owner;
retained.failure = Some(failure);
self.shared.status.send_replace(BootstrapStatus::Failed);
}
fn cancelled(&self, owner: Cleanup) {
{
let mut retained = self.shared.retained.lock().expect("bootstrap cleanup lock");
retained.owner = Some(owner.clone());
self.shared.status.send_replace(BootstrapStatus::Cleaning);
}
self.run(owner);
}
pub(super) fn settle_unowned(&self) {
let retained = self.shared.retained.lock().expect("bootstrap cleanup lock");
if retained.owner.is_none() && self.status() == BootstrapStatus::Starting {
self.shared.status.send_replace(BootstrapStatus::Settled);
}
}
}
impl<A: Application> Initializer<A> {
pub(super) fn new<F, Fut, E>(timeout: Duration, initialize: F) -> Self
where
F: FnOnce(super::Handle<A>) -> Fut + Send + 'static,
Fut: Future<Output = std::result::Result<(), E>> + Send + 'static,
E: std::error::Error + Send + Sync + 'static,
{
Self {
timeout,
initialize: Box::new(move |handle| Box::pin(async move { initialize(handle).await.map_err(|source| Box::new(source) as _) })),
}
}
}
impl Cleanup {
pub(super) async fn close(&self) -> Result<()> {
let outcome = (self.close)().await;
let Some(receipt) = self.receipt.upgrade() else {
return outcome;
};
let mut retained = receipt.retained.lock().expect("bootstrap cleanup lock");
match outcome {
Ok(()) => {
retained.owner = None;
retained.failure = None;
receipt.status.send_replace(BootstrapStatus::Settled);
Ok(())
}
Err(error) => {
let source = Arc::new(error);
retained.failure = Some(Arc::clone(&source));
receipt.status.send_replace(BootstrapStatus::Failed);
Err(crate::Error::TaskFailed { domain: "bootstrap cleanup", source })
}
}
}
}
impl std::fmt::Debug for Cleanup {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.debug_struct("Cleanup").finish_non_exhaustive()
}
}
impl<A: Application> Boot<A> {
pub(super) fn new(storage: Option<&StorageManager<crate::storage::ErasedCatalog>>, cleanup: BootstrapCleanup) -> Result<Self> {
let identity = Arc::new(());
if let Some(storage) = storage {
storage.claim_owner(identity.clone())?;
}
Ok(Self {
application: None,
identity,
storage: storage.cloned(),
runtime: None,
writer: None,
state: None,
armed: true,
cleanup,
})
}
pub(super) async fn acquire(&self) -> Result<()> {
if let Some(storage) = &self.storage {
storage.acquire().await?;
}
Ok(())
}
pub(super) async fn close(&mut self) -> Result<()> {
if self.runtime.is_none()
&& let Some(application) = &self.application
{
application.settle().await.map_err(|source| crate::Error::TaskFailed { domain: "application cleanup", source })?;
}
if self.runtime.is_none() {
if let Some(state) = self.state.take() {
state.seal();
}
if let Some(writer) = self.writer.as_mut() {
let _ = writer.await;
}
self.writer.take();
}
let result = if let Some(runtime) = &self.runtime {
runtime.shutdown().await.map(|_| ())
} else if let Some(storage) = &self.storage {
storage.close().await
} else {
Ok(())
};
self.armed = false;
if result.is_ok() {
self.runtime.take();
}
result
}
pub(super) async fn reject(mut self, source: super::Incompatible) -> super::Incompatible {
match self.close().await {
Ok(()) => {
self.cleanup.shared.status.send_replace(BootstrapStatus::Settled);
source
}
Err(failure) => {
let failure = Arc::new(failure);
let retained = self.retain_cleanup(true);
if retained.is_some() {
self.cleanup.failed(retained.clone(), Arc::clone(&failure));
} else {
self.cleanup.shared.status.send_replace(BootstrapStatus::Settled);
}
super::Incompatible {
cause: super::Cause::Cleanup {
source: Box::new(source),
failure: crate::Error::TaskFailed {
domain: "bootstrap cleanup",
source: failure,
},
},
cleanup: retained,
}
}
}
}
pub(super) fn retain_cleanup(&mut self, failure_reported: bool) -> Option<Cleanup> {
let state = self.state.take();
let writer = self.writer.take();
let runtime = self.runtime.take().filter(|runtime| !failure_reported || !physically_settled(runtime));
let application = self.application.take();
let storage = self.storage.take().filter(|storage| !storage.status().closed);
self.armed = false;
if runtime.is_none() && storage.is_none() && writer.is_none() && state.is_none() && application.is_none() {
return None;
}
let writer = Arc::new(tokio::sync::Mutex::new(writer));
let runtime = runtime.map(Arc::new);
let failure_reported = Arc::new(AtomicBool::new(false));
Some(Cleanup {
receipt: Arc::downgrade(&self.cleanup.shared),
close: Arc::new(move || {
let state = state.clone();
let writer = Arc::clone(&writer);
let runtime = runtime.clone();
let storage = storage.clone();
let application = application.clone();
let failure_reported = failure_reported.clone();
Box::pin(async move {
if runtime.is_none() {
if let Some(application) = application {
application
.settle()
.await
.map_err(|source| crate::Error::TaskFailed { domain: "application cleanup", source })?;
}
if let Some(state) = state {
state.seal();
}
let mut writer = writer.lock().await;
if let Some(task) = writer.as_mut() {
let _ = task.await;
}
writer.take();
}
if let Some(runtime) = runtime {
match runtime.shutdown().await {
Err(source) if physically_settled(&runtime) => {
if failure_reported.swap(true, Ordering::AcqRel) {
Ok(())
} else {
Err(source)
}
}
outcome => outcome.map(|_| ()),
}
} else if let Some(storage) = storage {
storage.close().await
} else {
Ok(())
}
})
}),
})
}
pub(super) fn finish(mut self) -> Runtime<A> {
self.armed = false;
self.cleanup.shared.status.send_replace(BootstrapStatus::HandedOff);
self.runtime.take().expect("successful boot owns a runtime")
}
}
impl<A: Application> Drop for Boot<A> {
fn drop(&mut self) {
if !self.armed {
return;
}
if let Some(owner) = self.retain_cleanup(false) {
self.cleanup.cancelled(owner);
}
}
}
fn physically_settled<A: Application>(runtime: &Runtime<A>) -> bool {
let stopping = runtime.stopping.lock().expect("runtime transition is healthy");
matches!(&*stopping, super::Stopping::Closing { outcome, .. } if outcome.get().is_some())
}