use std::time::Duration;
use chrono::Utc;
use rainier_drivers::{RedisClient, RedisConnector};
use rainier_support::{BoxFuture, Error, Result};
use serde_json::json;
use crate::job::QueuedJob;
use crate::queue::{FailedJob, Queue};
const ENTRY_ID: &str = "__redis_entry_id";
const GROUP: &str = "workers";
pub struct RedisQueue {
client: RedisClient,
keys: Keys,
consumer: String,
reservation: Duration,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Keys {
prefix: String,
}
impl Keys {
pub fn new(prefix: impl Into<String>) -> Self {
Self { prefix: prefix.into() }
}
pub fn stream(&self, queue: &str) -> String {
format!("{}{queue}", self.prefix)
}
pub fn delayed(&self, queue: &str) -> String {
format!("{}{queue}:delayed", self.prefix)
}
pub fn failed(&self, queue: &str) -> String {
format!("{}{queue}:failed", self.prefix)
}
pub fn all(&self, queue: &str) -> Vec<String> {
vec![self.stream(queue), self.delayed(queue), self.failed(queue)]
}
}
impl RedisQueue {
pub async fn connect(connector: &RedisConnector) -> Result<Self> {
Ok(Self::new(RedisClient::connect(connector).await?))
}
pub fn new(client: RedisClient) -> Self {
Self {
client,
keys: Keys::new("rainier:queue:"),
consumer: consumer_name(),
reservation: Duration::from_secs(90),
}
}
pub fn with_prefix(mut self, prefix: impl Into<String>) -> Self {
self.keys = Keys::new(prefix);
self
}
pub fn with_reservation(mut self, reservation: Duration) -> Self {
self.reservation = reservation;
self
}
pub fn as_consumer(mut self, name: impl Into<String>) -> Self {
self.consumer = name.into();
self
}
pub async fn check_eviction_policy(&self) -> Option<String> {
let policy = self.client.config_get("maxmemory-policy").await.ok().flatten()?;
if policy != "noeviction" {
tracing::warn!(
policy,
"this Redis evicts keys under memory pressure, and a queued job is a key it can \
evict — set `maxmemory-policy noeviction` before relying on this queue"
);
}
Some(policy)
}
pub fn keys(&self) -> &Keys {
&self.keys
}
async fn promote_delayed(&self, queue: &str) -> Result<u64> {
self.client
.promote_due(
&self.keys.delayed(queue),
&self.keys.stream(queue),
Utc::now().timestamp_millis(),
)
.await
}
fn entry_id(job: &QueuedJob) -> Result<&str> {
job.payload.get(ENTRY_ID).and_then(|id| id.as_str()).ok_or_else(|| {
Error::internal(format!(
"job `{}` has no Redis entry id — it was not reserved from this queue",
job.id
))
})
}
fn without_entry_id(job: &QueuedJob) -> QueuedJob {
let mut job = job.clone();
if let Some(payload) = job.payload.as_object_mut() {
payload.remove(ENTRY_ID);
}
job
}
fn encode(job: &QueuedJob) -> Result<Vec<u8>> {
serde_json::to_vec(job)
.map_err(|e| Error::internal(format!("a queued job must serialise: {e}")))
}
}
#[allow(clippy::needless_lifetimes, reason = "the trait's signature names them")]
impl Queue for RedisQueue {
fn name(&self) -> &str {
"redis"
}
fn push<'a>(&'a self, job: QueuedJob) -> BoxFuture<'a, Result<String>> {
Box::pin(async move {
let body = Self::encode(&job)?;
let now = Utc::now();
if job.available_at > now {
self.client
.zadd(
&self.keys.delayed(&job.queue),
job.available_at.timestamp_millis(),
&body,
)
.await?;
} else {
self.client.xadd(&self.keys.stream(&job.queue), &body).await?;
}
Ok(job.id)
})
}
fn reserve<'a>(&'a self, queue: &'a str) -> BoxFuture<'a, Result<Option<QueuedJob>>> {
Box::pin(async move {
let stream = self.keys.stream(queue);
self.client.xgroup_create(&stream, GROUP).await?;
self.promote_delayed(queue).await?;
let entry = match self
.client
.xautoclaim_one(&stream, GROUP, &self.consumer, self.reservation)
.await?
{
Some(entry) => Some(entry),
None => self.client.xreadgroup_one(&stream, GROUP, &self.consumer).await?,
};
let Some(entry) = entry else { return Ok(None) };
let mut job: QueuedJob = match serde_json::from_slice(&entry.body) {
Ok(job) => job,
Err(e) => {
tracing::error!(entry = entry.id, error = %e, "discarding an unreadable entry");
self.client.xack_delete(&stream, GROUP, &entry.id).await?;
return Ok(None);
}
};
job.payload[ENTRY_ID] = json!(entry.id);
job.attempts += 1;
Ok(Some(job))
})
}
fn acknowledge<'a>(&'a self, job: &'a QueuedJob) -> BoxFuture<'a, Result<()>> {
Box::pin(async move {
self.client
.xack_delete(&self.keys.stream(&job.queue), GROUP, Self::entry_id(job)?)
.await
})
}
fn release<'a>(&'a self, job: &'a QueuedJob, delay: Duration) -> BoxFuture<'a, Result<()>> {
Box::pin(async move {
let stream = self.keys.stream(&job.queue);
let entry = Self::entry_id(job)?;
let mut retry = Self::without_entry_id(job);
retry.available_at = Utc::now()
+ chrono::Duration::from_std(delay).unwrap_or_else(|_| chrono::Duration::zero());
let body = Self::encode(&retry)?;
if delay.is_zero() {
self.client.xadd(&stream, &body).await?;
} else {
self.client
.zadd(
&self.keys.delayed(&job.queue),
retry.available_at.timestamp_millis(),
&body,
)
.await?;
}
self.client.xack_delete(&stream, GROUP, entry).await
})
}
fn fail<'a>(&'a self, job: &'a QueuedJob, error: &'a str) -> BoxFuture<'a, Result<()>> {
Box::pin(async move {
let failed = FailedJob {
job: Self::without_entry_id(job),
error: error.to_string(),
failed_at: Utc::now(),
};
let body = serde_json::to_vec(&json!({
"job": failed.job,
"error": failed.error,
"failed_at": failed.failed_at,
}))
.map_err(|e| Error::internal(format!("a failed job must serialise: {e}")))?;
self.client.lpush(&self.keys.failed(&job.queue), &body).await?;
self.client
.xack_delete(&self.keys.stream(&job.queue), GROUP, Self::entry_id(job)?)
.await
})
}
fn size<'a>(&'a self, queue: &'a str) -> BoxFuture<'a, Result<u64>> {
Box::pin(async move {
let ready = self.client.xlen(&self.keys.stream(queue)).await?;
let delayed = self.client.zcard(&self.keys.delayed(queue)).await?;
Ok(ready + delayed)
})
}
fn clear<'a>(&'a self, queue: &'a str) -> BoxFuture<'a, Result<u64>> {
Box::pin(async move {
let size = self.size(queue).await.unwrap_or(0);
self.client.delete_all(&self.keys.all(queue)).await?;
Ok(size)
})
}
}
impl RedisQueue {
pub async fn failed_jobs(&self, queue: &str) -> Result<Vec<FailedJob>> {
let raw = self.client.lrange_all(&self.keys.failed(queue)).await?;
Ok(raw
.iter()
.filter_map(|body| {
let value: serde_json::Value = serde_json::from_slice(body).ok()?;
Some(FailedJob {
job: serde_json::from_value(value.get("job")?.clone()).ok()?,
error: value.get("error")?.as_str()?.to_string(),
failed_at: serde_json::from_value(value.get("failed_at")?.clone()).ok()?,
})
})
.collect())
}
}
fn consumer_name() -> String {
let host = std::env::var("HOSTNAME")
.or_else(|_| std::env::var("COMPUTERNAME"))
.unwrap_or_else(|_| "worker".to_string());
format!("{host}-{}", std::process::id())
}
#[cfg(test)]
mod tests {
use super::*;
fn job() -> QueuedJob {
QueuedJob {
id: "j1".into(),
name: "test".into(),
payload: json!({ "a": 1 }),
queue: "default".into(),
attempts: 0,
max_attempts: 3,
available_at: Utc::now(),
created_at: Utc::now(),
unique_key: None,
}
}
#[test]
fn one_queue_owns_three_keys_and_the_prefix_reaches_all_of_them() {
let keys = Keys::new("app:q:");
assert_eq!(keys.stream("default"), "app:q:default");
assert_eq!(keys.delayed("default"), "app:q:default:delayed");
assert_eq!(keys.failed("default"), "app:q:default:failed");
assert_eq!(keys.all("default").len(), 3);
assert!(keys.all("default").iter().all(|key| key.starts_with("app:q:")));
}
#[test]
fn two_queues_do_not_share_a_key() {
let keys = Keys::new("app:q:");
assert_ne!(keys.stream("default"), keys.stream("mail"));
assert_ne!(keys.stream("default"), keys.delayed("default"));
}
#[test]
fn the_entry_id_rides_on_the_payload_and_comes_off_again() {
let mut reserved = job();
reserved.payload[ENTRY_ID] = json!("1700000000000-0");
assert_eq!(RedisQueue::entry_id(&reserved).unwrap(), "1700000000000-0");
let retry = RedisQueue::without_entry_id(&reserved);
assert!(retry.payload.get(ENTRY_ID).is_none());
assert_eq!(retry.payload["a"], 1, "the job's own payload survives");
}
#[test]
fn acknowledging_a_job_that_was_never_reserved_says_so() {
let err = RedisQueue::entry_id(&job()).unwrap_err();
assert!(err.message().contains("not reserved"), "{}", err.message());
assert!(err.message().contains("j1"), "{}", err.message());
}
#[test]
fn a_consumer_names_the_host_and_the_process() {
let name = consumer_name();
assert!(name.contains(&std::process::id().to_string()), "{name}");
}
}