5ba9d7d
use export::{AppRedis, RedisConnection};
use crate::sync::error::SyncResult;
macro_rules! dispatch {
($conn:expr, |$c:ident| $body:expr) => {
match &mut $conn {
RedisConnection::Standalone($c) => $body,
RedisConnection::Cluster($c) => $body,
}
};
}
const INCR_EXPIRE_SCRIPT: &str = "local c = redis.call('INCR', KEYS[1])\nif c == 1 then\n redis.call('EXPIRE', KEYS[1], ARGV[1])\nend\nreturn c";
#[derive(Clone)]
pub struct SyncQueue {
redis: AppRedis,
}
impl SyncQueue {
#[must_use]
pub fn new(redis: AppRedis) -> Self {
Self { redis }
}
#[must_use]
pub fn redis(&self) -> &AppRedis {
&self.redis
}
pub async fn consume(&self, key: &str, timeout_secs: u64) -> SyncResult<Option<Vec<u8>>> {
let mut conn = self.redis.connection_clone();
let res: Option<(String, Vec<u8>)> = if timeout_secs == 0 {
dispatch!(conn, |c| {
redis::cmd("RPOP").arg(key).query_async(c).await
})
} else {
dispatch!(conn, |c| {
redis::cmd("BRPOP")
.arg(key)
.arg(timeout_secs)
.query_async(c)
.await
})
}?;
Ok(res.map(|(_, v)| v))
}
pub async fn publish_raw(
&self,
key: &str,
_partition: Option<&[u8]>,
payload: &[u8],
) -> SyncResult<i64> {
let mut conn = self.redis.connection_clone();
let n: i64 = dispatch!(conn, |c| {
redis::cmd("LPUSH")
.arg(key)
.arg(payload)
.query_async(c)
.await
})?;
Ok(n)
}
pub async fn requeue(&self, key: &str, payload: &[u8]) -> SyncResult<i64> {
self.publish_raw(key, None, payload).await
}
pub async fn incr_with_ttl(&self, key: &str, ttl_secs: u64) -> SyncResult<i64> {
let mut conn = self.redis.connection_clone();
let n: i64 = dispatch!(conn, |c| {
redis::cmd("EVAL")
.arg(INCR_EXPIRE_SCRIPT)
.arg(1)
.arg(key)
.arg(ttl_secs)
.query_async(c)
.await
})?;
Ok(n)
}
pub async fn delete(&self, key: &str) -> SyncResult<()> {
let mut conn = self.redis.connection_clone();
let _: i64 = dispatch!(conn, |c| {
redis::cmd("DEL").arg(key).query_async(c).await
})?;
Ok(())
}
}