mirror of
https://github.com/LemmyNet/lemmy.git
synced 2024-09-16 19:29:09 -06:00
create struct to hold cancellable task for readability
This commit is contained in:
parent
82c22439b8
commit
14479cefd2
@ -1,5 +1,5 @@
|
|||||||
use crate::{
|
use crate::{
|
||||||
util::{retry_sleep_duration, spawn_cancellable},
|
util::{retry_sleep_duration, CancellableTask},
|
||||||
worker::instance_worker,
|
worker::instance_worker,
|
||||||
};
|
};
|
||||||
use activitypub_federation::config::FederationConfig;
|
use activitypub_federation::config::FederationConfig;
|
||||||
@ -71,7 +71,7 @@ async fn start_stop_federation_workers<T: Clone + Send + Sync + 'static>(
|
|||||||
let stats_sender = stats_sender.clone();
|
let stats_sender = stats_sender.clone();
|
||||||
workers.insert(
|
workers.insert(
|
||||||
instance.id,
|
instance.id,
|
||||||
spawn_cancellable(WORKER_EXIT_TIMEOUT, |stop| {
|
CancellableTask::spawn(WORKER_EXIT_TIMEOUT, |stop| {
|
||||||
instance_worker(
|
instance_worker(
|
||||||
pool.clone(),
|
pool.clone(),
|
||||||
instance,
|
instance,
|
||||||
@ -83,7 +83,7 @@ async fn start_stop_federation_workers<T: Clone + Send + Sync + 'static>(
|
|||||||
);
|
);
|
||||||
} else if !should_federate {
|
} else if !should_federate {
|
||||||
if let Some(worker) = workers.remove(&instance.id) {
|
if let Some(worker) = workers.remove(&instance.id) {
|
||||||
if let Err(e) = worker.await {
|
if let Err(e) = worker.cancel().await {
|
||||||
tracing::error!("error stopping worker: {e}");
|
tracing::error!("error stopping worker: {e}");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@ -102,7 +102,8 @@ async fn start_stop_federation_workers<T: Clone + Send + Sync + 'static>(
|
|||||||
workers.len(),
|
workers.len(),
|
||||||
WORKER_EXIT_TIMEOUT
|
WORKER_EXIT_TIMEOUT
|
||||||
);
|
);
|
||||||
futures::future::join_all(workers.into_values()).await;
|
// the cancel futures need to be awaited concurrently for the shutdown processes to be triggered concurrently
|
||||||
|
futures::future::join_all(workers.into_values().map(|e| e.cancel())).await;
|
||||||
exit_print.await?;
|
exit_print.await?;
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
@ -113,8 +114,8 @@ pub fn start_stop_federation_workers_cancellable(
|
|||||||
opts: Opts,
|
opts: Opts,
|
||||||
pool: ActualDbPool,
|
pool: ActualDbPool,
|
||||||
config: FederationConfig<impl Clone + Send + Sync + 'static>,
|
config: FederationConfig<impl Clone + Send + Sync + 'static>,
|
||||||
) -> impl Future<Output = anyhow::Result<()>> {
|
) -> CancellableTask<(), impl Future<Output = anyhow::Result<()>>> {
|
||||||
spawn_cancellable(WORKER_EXIT_TIMEOUT, move |c| {
|
CancellableTask::spawn(WORKER_EXIT_TIMEOUT, move |c| {
|
||||||
start_stop_federation_workers(opts, pool, config, c)
|
start_stop_federation_workers(opts, pool, config, c)
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
@ -40,7 +40,7 @@ async fn main() -> anyhow::Result<()> {
|
|||||||
let mut interrupt = tokio::signal::unix::signal(SignalKind::interrupt())?;
|
let mut interrupt = tokio::signal::unix::signal(SignalKind::interrupt())?;
|
||||||
let mut terminate = tokio::signal::unix::signal(SignalKind::terminate())?;
|
let mut terminate = tokio::signal::unix::signal(SignalKind::terminate())?;
|
||||||
|
|
||||||
let cancel =
|
let task =
|
||||||
lemmy_federate::start_stop_federation_workers_cancellable(opts, pool, federation_config);
|
lemmy_federate::start_stop_federation_workers_cancellable(opts, pool, federation_config);
|
||||||
tokio::select! {
|
tokio::select! {
|
||||||
_ = tokio::signal::ctrl_c() => {
|
_ = tokio::signal::ctrl_c() => {
|
||||||
@ -53,6 +53,6 @@ async fn main() -> anyhow::Result<()> {
|
|||||||
tracing::warn!("Received terminate, shutting down gracefully...");
|
tracing::warn!("Received terminate, shutting down gracefully...");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
cancel.await?;
|
task.cancel().await?;
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
@ -29,13 +29,22 @@ use std::{
|
|||||||
use tokio::{task::JoinHandle, time::sleep};
|
use tokio::{task::JoinHandle, time::sleep};
|
||||||
use tokio_util::sync::CancellationToken;
|
use tokio_util::sync::CancellationToken;
|
||||||
|
|
||||||
|
pub struct CancellableTask<R: Send + 'static, F>
|
||||||
|
where
|
||||||
|
F: Future<Output = Result<R>>,
|
||||||
|
{
|
||||||
|
f: F,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<R: Send + 'static, F> CancellableTask<R, F>
|
||||||
|
where
|
||||||
|
F: Future<Output = Result<R>>,
|
||||||
|
{
|
||||||
/// spawn a task but with graceful shutdown
|
/// spawn a task but with graceful shutdown
|
||||||
///
|
pub fn spawn(
|
||||||
/// only await the returned future when you want to cancel the task
|
|
||||||
pub fn spawn_cancellable<R: Send + 'static, F>(
|
|
||||||
timeout: Duration,
|
timeout: Duration,
|
||||||
task: impl FnOnce(CancellationToken) -> F,
|
task: impl FnOnce(CancellationToken) -> F,
|
||||||
) -> impl Future<Output = Result<R>>
|
) -> CancellableTask<R, impl Future<Output = Result<R>>>
|
||||||
where
|
where
|
||||||
F: Future<Output = Result<R>> + Send + 'static,
|
F: Future<Output = Result<R>> + Send + 'static,
|
||||||
{
|
{
|
||||||
@ -52,7 +61,8 @@ where
|
|||||||
}
|
}
|
||||||
});
|
});
|
||||||
let abort = task.abort_handle();
|
let abort = task.abort_handle();
|
||||||
async move {
|
CancellableTask {
|
||||||
|
f: async move {
|
||||||
stop.cancel();
|
stop.cancel();
|
||||||
tokio::select! {
|
tokio::select! {
|
||||||
r = task => {
|
r = task => {
|
||||||
@ -64,6 +74,13 @@ where
|
|||||||
Err(anyhow!("task aborted due to timeout"))
|
Err(anyhow!("task aborted due to timeout"))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 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> {
|
||||||
|
self.f.await
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user