mod control;
mod intercom;
mod stats;
mod status;
pub use self::{
control::{Control, ControlReader, Controller},
intercom::{
Intercom, IntercomMsg, IntercomReceiver, IntercomSender, IntercomStats, IntercomStatus,
NoIntercom,
},
stats::Stats,
status::{Status, StatusReader, StatusUpdater},
};
use crate::{runtime::Runtime, watchdog::WatchdogQuery};
use async_trait::async_trait;
use futures_util::future::abortable;
use std::future::Future;
use thiserror::Error;
use tokio::{runtime::Handle, task::JoinHandle};
use tracing_futures::Instrument as _;
pub type ServiceIdentifier = &'static str;
#[async_trait]
pub trait Service: Send + Sized + 'static {
const SERVICE_IDENTIFIER: ServiceIdentifier;
type IntercomMsg: IntercomMsg;
fn prepare(service_state: ServiceState<Self>) -> Self;
async fn start(self);
}
pub trait ManageService {
const SERVICE_IDENTIFIER: ServiceIdentifier;
type IntercomMsg: IntercomMsg;
}
impl<T: Service> ManageService for ServiceManager<T> {
const SERVICE_IDENTIFIER: ServiceIdentifier = T::SERVICE_IDENTIFIER;
type IntercomMsg = T::IntercomMsg;
}
#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum ServiceError {
#[error("Service cannot be started because status is: {status}")]
CannotStart { status: Status },
}
#[derive(Debug, Clone)]
pub struct StatusReport {
pub identifier: ServiceIdentifier,
pub status: Status,
pub intercom: IntercomStatus,
pub started: u64,
}
pub struct ServiceManager<T: Service> {
identifier: ServiceIdentifier,
intercom_sender: IntercomSender<T::IntercomMsg>,
intercom_stats: IntercomStats,
started: u64,
status: StatusReader,
controller: Controller,
runtime: Handle,
}
pub struct ServiceRuntime<T: Service> {
service_state: ServiceState<T>,
status: StatusUpdater,
control: ControlReader,
}
pub struct ServiceState<T: Service> {
identifier: ServiceIdentifier,
handle: Handle,
intercom_receiver: IntercomReceiver<T::IntercomMsg>,
watchdog_query: WatchdogQuery,
status: StatusReader,
}
impl<T: Service> ServiceState<T> {
pub fn identifier(&self) -> ServiceIdentifier {
self.identifier
}
pub fn intercom_with<O: Service>(&self) -> Intercom<O> {
self.watchdog_query.intercom::<O>()
}
pub fn watchdog_controller(&self) -> &WatchdogQuery {
&self.watchdog_query
}
pub fn intercom_mut(&mut self) -> &mut IntercomReceiver<T::IntercomMsg> {
&mut self.intercom_receiver
}
pub fn status_reader(&self) -> &StatusReader {
&self.status
}
pub fn runtime_handle(&self) -> &Handle {
&self.handle
}
pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
self.runtime_handle().spawn(future)
}
}
impl<T: Service> ServiceManager<T> {
pub fn with_runtime(runtime: &mut Runtime) -> Self {
let identifier = T::SERVICE_IDENTIFIER;
let status = StatusReader::new(Status::shutdown());
let controller = runtime.block_on(async { Controller::new().await });
let (intercom_sender, _, intercom_stats) = intercom::channel();
Self {
identifier,
intercom_sender,
intercom_stats,
status,
controller,
runtime: runtime.handle().clone(),
started: 0,
}
}
pub fn intercom(&self) -> IntercomSender<T::IntercomMsg> {
self.intercom_sender.clone()
}
pub async fn status(&self) -> StatusReport {
StatusReport {
identifier: self.identifier,
status: self.status.status(),
intercom: self.intercom_stats.status().await,
started: self.started,
}
}
pub fn shutdown(&mut self) {
match self.status.status() {
Status::Shutdown { .. } | Status::ShuttingDown { .. } => {
}
Status::Starting { .. } | Status::Started { .. } => {
self.controller.send(Control::Shutdown)
}
}
}
pub fn runtime(
&mut self,
watchdog_query: WatchdogQuery,
) -> Result<ServiceRuntime<T>, ServiceError> {
let status = self.status.status();
if !status.is_shutdown() {
Err(ServiceError::CannotStart { status })
} else {
let (intercom_sender, intercom_receiver, intercom_stats) =
intercom::channel::<T::IntercomMsg>();
self.intercom_sender = intercom_sender;
self.intercom_stats = intercom_stats;
self.started += 1;
Ok(ServiceRuntime {
service_state: ServiceState {
identifier: self.identifier,
handle: self.runtime.clone(),
status: self.status.clone(),
intercom_receiver,
watchdog_query,
},
status: self.status.updater(),
control: self.controller.reader(),
})
}
}
}
impl<T: Service> ServiceRuntime<T> {
pub fn start(self) {
let ServiceRuntime {
service_state,
status,
mut control,
} = self;
let service_identifier: &'static str = service_state.identifier;
status.update(Status::starting());
let watchdog_query = service_state.watchdog_query.clone();
let handle = service_state.handle.clone();
let runner = T::prepare(service_state);
let (runner, abort_handle) = abortable(async move {
let span = tracing::info_span!("service", service_identifier);
let _enter = span.enter();
runner.start().in_current_span().await
});
let mut service_join_handle = handle.spawn(runner);
watchdog_query.spawn(async move {
status.update(Status::started());
let span = tracing::debug_span!("service control", service_identifier);
let _enter = span.enter();
loop {
tokio::select! {
join_result = &mut service_join_handle => {
if let Err(join_error) = join_result {
tracing::error!(
"main process failed with following error: {:#?}",
join_error
);
} else {
}
status.update(Status::shutdown());
break;
}
control = control.updated() => {
match control {
Some(Control::Shutdown) => {
tracing::info!("shutting down...");
status.update(Status::shutting_down());
}
None | Some(Control::Kill) => {
tracing::info!("Terminating...");
status.update(Status::shutdown());
abort_handle.abort();
break;
}
}
}
};
}
});
}
}
impl<T: Service> Drop for ServiceManager<T> {
fn drop(&mut self) {
if !self.status.status().is_shutdown() {
self.controller.send(Control::Kill)
}
}
}