use crate::error::ContextWrapper;
use crate::wait_token::WaitToken;
use futures::{Future, FutureExt};
use std::panic::AssertUnwindSafe;
use std::time::{Duration, Instant};
use tokio::select;
use tokio::task::JoinHandle;
pub trait Repeater<Payload, Out> {
fn repeat(self, payload: Payload) -> impl std::future::Future<Output = Out> + Send;
}
impl<Promise, Payload, T> Repeater<Payload, Promise::Output> for T
where
Promise: Future + Send + 'static,
Promise::Output: Send,
Payload: Send + Sync + Clone + 'static,
T: Fn(Payload) -> Promise + Send,
{
async fn repeat(self, payload: Payload) -> Promise::Output {
loop {
(self)(payload.clone())
.spawn()
.await
.anyhow_with("Task stopped, repeat ... ")
.ok();
}
}
}
pub trait TimeMeasurable {
type Output;
fn duration(self) -> impl std::future::Future<Output = Self::Output> + Send;
}
impl<T> TimeMeasurable for T
where
T: Future + Send + 'static,
{
type Output = (Duration, T::Output);
async fn duration(self) -> Self::Output {
let now = Instant::now();
let result = self.await;
(now.elapsed(), result)
}
}
pub trait Ticker {
fn next(&self) -> impl Future<Output = bool> + Send + Sync;
}
impl Ticker for Option<Duration> {
async fn next(&self) -> bool {
let Some(duration) = self else {
return false;
};
tokio::time::sleep(*duration).await;
true
}
}
impl Ticker for Duration {
async fn next(&self) -> bool {
tokio::time::sleep(*self).await;
true
}
}
pub trait Spawn {
type Result;
fn spawn(self) -> Self::Result;
}
impl<Fut: Future + Sized + Send + 'static> Spawn for Fut
where
Fut::Output: Send + 'static,
{
type Result = JoinHandle<Fut::Output>;
fn spawn(self) -> Self::Result {
tokio::spawn(self)
}
}
pub trait SpawnBlockingEmpty {
type R;
fn spawn_blocking(self) -> impl Future<Output = anyhow::Result<Self::R>>;
}
pub trait SpawnBlockingArgs<X, Y> {
type R;
type M;
fn spawn_blocking(
self,
thing: Self::M,
args: Y,
) -> impl Future<Output = anyhow::Result<Self::R>>;
}
#[allow(non_snake_case, unused_mut)]
impl<F, R> SpawnBlockingEmpty for F
where
F: FnOnce() -> R + Send + 'static,
R: Send + 'static,
{
type R = R;
async fn spawn_blocking(self) -> anyhow::Result<R> {
tokio::task::spawn_blocking(self).await.anyhow()
}
}
#[allow(non_snake_case, unused_mut)]
impl<F, R, M, T1> SpawnBlockingArgs<(M, T1), T1> for F
where
F: FnOnce(T1) -> R + Send + 'static,
R: Send + 'static,
T1: Send + 'static,
{
type R = R;
type M = M;
async fn spawn_blocking(self, _: M, args: T1) -> anyhow::Result<R> {
tokio::task::spawn_blocking(move || (self)(args))
.await
.anyhow()
}
}
#[allow(non_snake_case, unused_mut)]
impl<F, R, M, T1, T2> SpawnBlockingArgs<(M, T1, T2), (T1, T2)> for F
where
F: FnOnce(T1, T2) -> R + Send + 'static,
R: Send + 'static,
T1: Send + 'static,
T2: Send + 'static,
{
type R = R;
type M = M;
async fn spawn_blocking(self, _: M, args: (T1, T2)) -> anyhow::Result<R> {
tokio::task::spawn_blocking(move || (self)(args.0, args.1))
.await
.anyhow()
}
}
#[allow(non_snake_case, unused_mut)]
impl<F, R, M, T1, T2, T3> SpawnBlockingArgs<(M, T1, T2, T3), (T1, T2, T3)> for F
where
F: FnOnce(T1, T2, T3) -> R + Send + 'static,
R: Send + 'static,
T1: Send + 'static,
T2: Send + 'static,
T3: Send + 'static,
{
type R = R;
type M = M;
async fn spawn_blocking(self, _: M, args: (T1, T2, T3)) -> anyhow::Result<R> {
tokio::task::spawn_blocking(move || (self)(args.0, args.1, args.2))
.await
.anyhow()
}
}
pub trait Timeout {
type Output;
fn timeout(
self,
time: Duration,
) -> impl std::future::Future<Output = anyhow::Result<Self::Output>> + Send;
}
pub trait TimeoutOpt {
type Output;
fn timeout_opt(
self,
time: Option<Duration>,
) -> impl std::future::Future<Output = anyhow::Result<Self::Output>> + Send;
}
impl<T> Timeout for T
where
T: Future + Send,
{
type Output = T::Output;
async fn timeout(self, time: Duration) -> anyhow::Result<Self::Output> {
tokio::time::timeout(time, self).await.anyhow()
}
}
impl<T> TimeoutOpt for T
where
T: Future + Send,
{
type Output = T::Output;
async fn timeout_opt(self, time: Option<Duration>) -> anyhow::Result<Self::Output> {
let Some(duration) = time else {
return Ok(self.await);
};
tokio::time::timeout(duration, self).await.anyhow()
}
}
pub trait Handler: Send + Clone + 'static
where
Self: 'static,
{
fn run(&mut self) -> impl Future<Output = anyhow::Result<()>> + Send;
fn on_err(&mut self) -> impl Future<Output = anyhow::Result<()>> + Send + Sync {
async { Ok(()) }
}
}
pub trait Thread<Handler>: Sized + Send {
fn new(handler: Handler) -> Self;
fn kill(self) -> Self;
fn with_msg(self) -> Self;
fn with_name(self, name: impl AsRef<str>) -> Self;
fn with_restart(self, duration: Duration) -> Self;
fn with_delay(self, duration: Duration) -> Self;
fn with_timeout(self, duration: Duration) -> Self;
fn with_invoke_frq(self, duration: Duration) -> Self;
fn with_cancellation(self, token: WaitToken) -> Self;
fn run(self) -> JoinHandle<()>;
}
#[derive(Clone)]
pub struct ThreadController<H> {
kill: bool,
msg: bool,
timeout: Option<Duration>,
invoke_frq: Option<Duration>,
cancellation: Option<WaitToken>,
restart: Option<Duration>,
delay: Option<Duration>,
name: String,
handler: H,
}
impl<H> Thread<H> for ThreadController<H>
where
H: Handler + Sized,
H: Send + Sync + 'static,
{
fn new(handler: H) -> Self {
Self {
timeout: None,
invoke_frq: None,
cancellation: None,
restart: None,
delay: None,
name: "DefaultName".to_string(),
handler,
kill: false,
msg: false,
}
}
fn kill(mut self) -> Self {
self.kill = true;
self
}
fn with_msg(mut self) -> Self {
self.msg = true;
self
}
fn with_delay(mut self, duration: Duration) -> Self {
self.delay = Some(duration);
self
}
fn with_name(mut self, name: impl AsRef<str>) -> Self {
self.name = name.as_ref().to_string();
self
}
fn with_restart(mut self, duration: Duration) -> Self {
self.restart = Some(duration);
self
}
fn with_timeout(mut self, duration: Duration) -> Self {
self.timeout = Some(duration);
self
}
fn with_invoke_frq(mut self, duration: Duration) -> Self {
self.invoke_frq = Some(duration);
self
}
fn with_cancellation(mut self, token: WaitToken) -> Self {
self.cancellation = Some(token);
self
}
fn run(self) -> JoinHandle<()> {
async move {
let token = self.cancellation.clone().unwrap_or_default();
token
.clone()
.repeat_until_cancel(self.delay.unwrap_or_default())
.next()
.await;
let mut clock = token
.clone()
.repeat_until_cancel(self.restart.unwrap_or_default());
if self.msg {
tracing::warn!(target: "dutils", "Starting {}", self.name);
}
loop {
let mut handler = self.clone();
if self.kill {
select! {
_ = AssertUnwindSafe(handler.main_loop()).catch_unwind() => {
}
_ = token.cancelled() => {
break;
}
}
} else {
AssertUnwindSafe(handler.main_loop())
.catch_unwind()
.await
.ok();
}
if self.cancellation.is_none() && self.restart.is_none() || !clock.next().await {
break;
}
}
if self.msg {
tracing::warn!(target: "dutils", "Stopped {}", self.name);
}
}
.spawn()
}
}
impl<H> ThreadController<H>
where
H: Handler + Sized,
H: Send + Sync + 'static,
{
pub async fn main_loop(&mut self) {
let mut clock = self
.cancellation
.clone()
.unwrap_or_default()
.repeat_until_cancel(self.invoke_frq.unwrap_or_default());
loop {
let result = async { self.handler.run().timeout_opt(self.timeout).await? }
.await
.track_with(&self.name);
if result.is_err() {
self.handler
.on_err()
.await
.anyhow_with(&self.name)
.track()
.ok();
}
if self.cancellation.is_none() && self.invoke_frq.is_none() || !clock.next().await {
break;
}
}
}
}