use nanoid::nanoid;
use std::{
cmp,
marker::PhantomData,
sync::Arc,
time::{Duration, Instant},
};
use tracing::{Instrument, Level, debug, info, span, trace, warn};
use chrono::{DateTime, Utc};
use deadpool_redis::{Pool, PoolError};
use redis::{AsyncCommands, RedisError};
use serde::de::DeserializeOwned;
use tokio::{
spawn,
sync::{
Semaphore,
mpsc::{self, Sender},
},
time::sleep,
};
use crate::{
job::JobWorkHandle,
luacommands::{InvokeLuaScript as _, MoveToActive, MoveToActiveOk, RateLimiter},
queue::QueueName,
worker::workererror::WorkerError,
};
pub async fn pull_job_thread<D, R>(
pool: Pool,
queue_name: QueueName,
job_sender: Sender<Result<JobWorkHandle<D, R>, WorkerError>>,
semaphore: Arc<Semaphore>,
failure_cooldown: Duration,
pull_worker_id: String,
) where
D: DeserializeOwned + std::fmt::Debug,
{
let (marker_send, mut marker_recv) = mpsc::channel(1);
let poll_span = span!(Level::TRACE, "poll-marker");
spawn(poll_marker(pool.clone(), queue_name.clone(), marker_send).instrument(poll_span));
let mut counter: usize = 0;
loop {
let permit = semaphore
.clone()
.acquire_owned()
.await
.expect("semaphore is never closed");
let start = Instant::now();
let con = pool.get().await;
if let Err(e) = con {
warn!("Failed to get Redis connection: {}", e);
let fatal = match &e {
&PoolError::Closed => true,
_ => false,
};
if job_sender.send(Err(e.into())).await.is_err() {
warn!("Receiver dropped, terminate");
return;
}
if fatal {
warn!("Error is fatal, terminate");
return;
}
sleep(failure_cooldown).await;
continue;
}
let mut con = con.unwrap();
trace!("Acquired connection after {:?}", start.elapsed());
let lock_token = format!("{pull_worker_id}-{counter}");
counter += 1;
let mts = MoveToActive::<D> {
queue: &queue_name,
worker_id: &pull_worker_id,
limiter: RateLimiter {
max: 0,
duration: Duration::from_millis(0),
},
lock_duration: Duration::from_secs(30),
token: &lock_token,
phantom: PhantomData, };
let get_job = mts.call(&mut con).await.unwrap();
let sleep_timer: Option<Duration> = match get_job {
MoveToActiveOk::JobData { id, data } => {
let lock_refresh_handle = tokio::spawn(lock_refresh());
info!(
"Fetched job {id}, preload channel has capacity of {}",
job_sender.capacity()
);
let _closed = job_sender
.send(Ok(JobWorkHandle::new(
queue_name.clone(),
pool.clone(),
id,
data.name,
permit,
data.data,
lock_refresh_handle,
lock_token.clone(),
pull_worker_id.clone(),
)))
.await
.is_err();
None
}
MoveToActiveOk::Delay { delay } => Some(delay),
MoveToActiveOk::WaitUntil { timestamp } => {
Some(Duration::from_millis(
cmp::max(0, (timestamp - Utc::now()).num_milliseconds()) as u64,
))
}
MoveToActiveOk::NothingToDo => Some(Duration::from_secs(10)),
};
if let Some(sleep_timer) = sleep_timer {
trace!("Nothing in queue, sleep for {:?}", sleep_timer);
let timeout = sleep(sleep_timer);
let marker = marker_recv.recv();
tokio::select! {
_ = timeout => (),
event = marker => {
let (job_id, ts) = event.expect("poll never terminates first");
debug!("Received marker {} {}", job_id, ts);
},
};
}
}
}
async fn poll_marker(
pool: Pool,
queue_name: QueueName,
sender: mpsc::Sender<(String, DateTime<Utc>)>,
) {
let marker_name = queue_name.marker();
loop {
if sender.is_closed() {
info!("Terminate gracefully");
return;
}
let con = pool.get().await;
if let Err(e) = con {
warn!("Marker poll could not get connection from pool: {:?}", e);
sleep(Duration::from_secs(1)).await;
continue;
}
let mut con = con.unwrap();
let res: Result<Option<(String, String, i64)>, RedisError> =
con.bzpopmin(&marker_name, 30.).await;
if let Err(e) = res {
warn!("Marker poll failed to get next timestamp: {:?}", e);
sleep(Duration::from_secs(1)).await;
continue;
}
let res = res.unwrap();
if res.is_none() {
continue;
}
let (_key, job_id, timestamp) = res.unwrap();
let ts: Option<DateTime<Utc>> = DateTime::from_timestamp_millis(timestamp);
if ts.is_none() {
warn!("Marker poll failed to parse next timestamp: {}", timestamp);
sleep(Duration::from_secs(1)).await;
continue;
}
let ts = ts.unwrap();
if let Err(_e) = sender.send((job_id, ts)).await {
return;
}
}
}
async fn lock_refresh() {}