2022-07-31 13:32:37 +00:00
|
|
|
use crate::asynk::async_queue::AsyncQueueable;
|
2022-08-16 10:25:19 +00:00
|
|
|
use crate::asynk::async_queue::DEFAULT_TASK_TYPE;
|
2022-07-31 13:32:37 +00:00
|
|
|
use crate::asynk::async_worker::AsyncWorker;
|
2022-08-31 17:45:13 +00:00
|
|
|
use crate::FangError;
|
2022-08-01 16:55:36 +00:00
|
|
|
use crate::{RetentionMode, SleepParams};
|
2022-07-31 13:32:37 +00:00
|
|
|
use async_recursion::async_recursion;
|
|
|
|
use log::error;
|
2022-08-04 15:22:53 +00:00
|
|
|
use tokio::task::JoinHandle;
|
2022-07-31 13:32:37 +00:00
|
|
|
use typed_builder::TypedBuilder;
|
|
|
|
|
|
|
|
#[derive(TypedBuilder, Clone)]
|
2022-08-02 14:32:58 +00:00
|
|
|
pub struct AsyncWorkerPool<AQueue>
|
2022-07-31 13:32:37 +00:00
|
|
|
where
|
2022-08-02 14:32:58 +00:00
|
|
|
AQueue: AsyncQueueable + Clone + Sync + 'static,
|
2022-07-31 13:32:37 +00:00
|
|
|
{
|
|
|
|
#[builder(setter(into))]
|
2022-12-22 15:38:56 +00:00
|
|
|
/// the AsyncWorkerPool uses a queue to control the tasks that will be executed.
|
2022-08-02 14:32:58 +00:00
|
|
|
pub queue: AQueue,
|
2022-12-22 15:38:56 +00:00
|
|
|
/// sleep_params controls how much time a worker will sleep while waiting for tasks
|
2022-08-01 16:55:36 +00:00
|
|
|
#[builder(default, setter(into))]
|
|
|
|
pub sleep_params: SleepParams,
|
2022-12-22 15:38:56 +00:00
|
|
|
/// retention_mode controls if tasks should be persisted after execution
|
2022-08-01 16:55:36 +00:00
|
|
|
#[builder(default, setter(into))]
|
|
|
|
pub retention_mode: RetentionMode,
|
2022-12-22 15:38:56 +00:00
|
|
|
/// the number of workers of the AsyncWorkerPool.
|
2022-07-31 13:32:37 +00:00
|
|
|
#[builder(setter(into))]
|
2022-08-01 19:37:32 +00:00
|
|
|
pub number_of_workers: u32,
|
2022-12-22 15:38:56 +00:00
|
|
|
/// The type of tasks that will be executed by `AsyncWorkerPool`.
|
2022-08-16 10:25:19 +00:00
|
|
|
#[builder(default=DEFAULT_TASK_TYPE.to_string(), setter(into))]
|
|
|
|
pub task_type: String,
|
2022-07-31 13:32:37 +00:00
|
|
|
}
|
|
|
|
|
2022-08-02 14:32:58 +00:00
|
|
|
impl<AQueue> AsyncWorkerPool<AQueue>
|
2022-07-31 13:32:37 +00:00
|
|
|
where
|
2022-08-02 14:32:58 +00:00
|
|
|
AQueue: AsyncQueueable + Clone + Sync + 'static,
|
2022-07-31 13:32:37 +00:00
|
|
|
{
|
2022-12-22 15:38:56 +00:00
|
|
|
/// Starts the configured number of workers
|
|
|
|
/// This is necessary in order to execute tasks.
|
2022-07-31 13:32:37 +00:00
|
|
|
pub async fn start(&mut self) {
|
2022-08-04 15:22:53 +00:00
|
|
|
for idx in 0..self.number_of_workers {
|
|
|
|
let pool = self.clone();
|
|
|
|
tokio::spawn(Self::supervise_task(pool, 0, idx));
|
2022-07-31 13:32:37 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
#[async_recursion]
|
2022-12-22 15:38:56 +00:00
|
|
|
async fn supervise_task(pool: AsyncWorkerPool<AQueue>, restarts: u64, worker_number: u32) {
|
2022-08-04 15:22:53 +00:00
|
|
|
let restarts = restarts + 1;
|
|
|
|
let join_handle = Self::spawn_worker(
|
|
|
|
pool.queue.clone(),
|
|
|
|
pool.sleep_params.clone(),
|
|
|
|
pool.retention_mode.clone(),
|
2022-08-16 10:25:19 +00:00
|
|
|
pool.task_type.clone(),
|
2022-08-04 15:22:53 +00:00
|
|
|
)
|
|
|
|
.await;
|
2022-07-31 13:32:37 +00:00
|
|
|
|
2022-08-04 15:22:53 +00:00
|
|
|
if (join_handle.await).is_err() {
|
|
|
|
error!(
|
|
|
|
"Worker {} stopped. Restarting. the number of restarts {}",
|
|
|
|
worker_number, restarts,
|
|
|
|
);
|
|
|
|
Self::supervise_task(pool, restarts, worker_number).await;
|
2022-07-31 13:32:37 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2022-12-22 15:38:56 +00:00
|
|
|
async fn spawn_worker(
|
2022-08-04 15:22:53 +00:00
|
|
|
queue: AQueue,
|
|
|
|
sleep_params: SleepParams,
|
|
|
|
retention_mode: RetentionMode,
|
2022-08-16 10:25:19 +00:00
|
|
|
task_type: String,
|
2022-08-31 17:45:13 +00:00
|
|
|
) -> JoinHandle<Result<(), FangError>> {
|
2022-08-16 10:25:19 +00:00
|
|
|
tokio::spawn(async move {
|
|
|
|
Self::run_worker(queue, sleep_params, retention_mode, task_type).await
|
|
|
|
})
|
2022-08-04 15:22:53 +00:00
|
|
|
}
|
2022-12-22 15:38:56 +00:00
|
|
|
async fn run_worker(
|
2022-08-02 14:32:58 +00:00
|
|
|
queue: AQueue,
|
2022-08-01 16:55:36 +00:00
|
|
|
sleep_params: SleepParams,
|
|
|
|
retention_mode: RetentionMode,
|
2022-08-16 10:25:19 +00:00
|
|
|
task_type: String,
|
2022-08-31 17:45:13 +00:00
|
|
|
) -> Result<(), FangError> {
|
2022-08-02 14:32:58 +00:00
|
|
|
let mut worker: AsyncWorker<AQueue> = AsyncWorker::builder()
|
|
|
|
.queue(queue)
|
2022-08-01 16:55:36 +00:00
|
|
|
.sleep_params(sleep_params)
|
|
|
|
.retention_mode(retention_mode)
|
2022-08-16 10:25:19 +00:00
|
|
|
.task_type(task_type)
|
2022-07-31 13:32:37 +00:00
|
|
|
.build();
|
|
|
|
|
|
|
|
worker.run_tasks().await
|
|
|
|
}
|
|
|
|
}
|