|
|
|
@ -20,12 +20,7 @@ use moka::future::Cache;
|
|
|
|
|
use once_cell::sync::Lazy;
|
|
|
|
|
use reqwest::Url;
|
|
|
|
|
use serde_json::Value;
|
|
|
|
|
use std::{
|
|
|
|
|
future::Future,
|
|
|
|
|
pin::Pin,
|
|
|
|
|
sync::{Arc, RwLock},
|
|
|
|
|
time::Duration,
|
|
|
|
|
};
|
|
|
|
|
use std::{fmt::Debug, future::Future, pin::Pin, sync::Arc, time::Duration};
|
|
|
|
|
use tokio::{task::JoinHandle, time::sleep};
|
|
|
|
|
use tokio_util::sync::CancellationToken;
|
|
|
|
|
|
|
|
|
@ -49,41 +44,41 @@ pub(crate) static WORK_FINISHED_RECHECK_DELAY: Lazy<Duration> = Lazy::new(|| {
|
|
|
|
|
}
|
|
|
|
|
});
|
|
|
|
|
|
|
|
|
|
pub struct CancellableTask<R: Send + 'static> {
|
|
|
|
|
f: Pin<Box<dyn Future<Output = Result<R, anyhow::Error>> + Send + 'static>>,
|
|
|
|
|
ended: Arc<RwLock<bool>>,
|
|
|
|
|
/// A task that will be run in an infinite loop, unless it is cancelled.
|
|
|
|
|
/// If the task exits without being cancelled, an error will be logged and the task will be restarted.
|
|
|
|
|
pub struct CancellableTask {
|
|
|
|
|
f: Pin<Box<dyn Future<Output = Result<(), anyhow::Error>> + Send + 'static>>,
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
impl<R: Send + 'static> CancellableTask<R> {
|
|
|
|
|
impl CancellableTask {
|
|
|
|
|
/// spawn a task but with graceful shutdown
|
|
|
|
|
pub fn spawn<F>(
|
|
|
|
|
pub fn spawn<F, R: Debug>(
|
|
|
|
|
timeout: Duration,
|
|
|
|
|
task: impl FnOnce(CancellationToken) -> F,
|
|
|
|
|
) -> CancellableTask<R>
|
|
|
|
|
task: impl Fn(CancellationToken) -> F + Send + 'static,
|
|
|
|
|
) -> CancellableTask
|
|
|
|
|
where
|
|
|
|
|
F: Future<Output = Result<R>> + Send + 'static,
|
|
|
|
|
F: Future<Output = R> + Send + 'static,
|
|
|
|
|
{
|
|
|
|
|
let stop = CancellationToken::new();
|
|
|
|
|
let task = task(stop.clone());
|
|
|
|
|
let ended = Arc::new(RwLock::new(false));
|
|
|
|
|
let ended_write = ended.clone();
|
|
|
|
|
let task: JoinHandle<Result<R>> = tokio::spawn(async move {
|
|
|
|
|
match task.await {
|
|
|
|
|
Ok(o) => Ok(o),
|
|
|
|
|
Err(e) => {
|
|
|
|
|
*ended_write.write().expect("poisoned") = true;
|
|
|
|
|
Err(e)
|
|
|
|
|
let stop2 = stop.clone();
|
|
|
|
|
let task: JoinHandle<()> = tokio::spawn(async move {
|
|
|
|
|
loop {
|
|
|
|
|
let res = task(stop2.clone()).await;
|
|
|
|
|
if stop2.is_cancelled() {
|
|
|
|
|
return;
|
|
|
|
|
} else {
|
|
|
|
|
tracing::warn!("task exited, restarting: {res:?}");
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
});
|
|
|
|
|
let abort = task.abort_handle();
|
|
|
|
|
CancellableTask {
|
|
|
|
|
ended,
|
|
|
|
|
f: Box::pin(async move {
|
|
|
|
|
stop.cancel();
|
|
|
|
|
tokio::select! {
|
|
|
|
|
r = task => {
|
|
|
|
|
Ok(r.context("could not join")??)
|
|
|
|
|
r.context("could not join")?;
|
|
|
|
|
Ok(())
|
|
|
|
|
},
|
|
|
|
|
_ = sleep(timeout) => {
|
|
|
|
|
abort.abort();
|
|
|
|
@ -96,12 +91,9 @@ impl<R: Send + 'static> CancellableTask<R> {
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// cancel the cancel signal, wait for timeout for the task to stop gracefully, otherwise abort it
|
|
|
|
|
pub async fn cancel(self) -> Result<R, anyhow::Error> {
|
|
|
|
|
pub async fn cancel(self) -> Result<(), anyhow::Error> {
|
|
|
|
|
self.f.await
|
|
|
|
|
}
|
|
|
|
|
pub fn has_ended(&self) -> bool {
|
|
|
|
|
*self.ended.read().expect("poisoned")
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// assuming apub priv key and ids are immutable, then we don't need to have TTL
|
|
|
|
|