4208993
use export::{AppRedis, RedisConnection};
use redis::aio::MultiplexedConnection;
use crate::sync::error::{SyncError, 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,
consumer_conns: Vec<MultiplexedConnection>,
}
impl SyncQueue {
#[must_use]
pub fn new(redis: AppRedis) -> Self {
Self {
redis,
consumer_conns: Vec::new(),
}
}
pub async fn with_consumer_conns(redis: AppRedis, count: usize) -> SyncResult<Self> {
let count = count.max(1);
let mut consumer_conns = Vec::with_capacity(count);
for _ in 0..count {
let conn = open_multiplexed(&redis).await?;
consumer_conns.push(conn);
}
Ok(Self {
redis,
consumer_conns,
})
}
#[must_use]
pub fn redis(&self) -> &AppRedis {
&self.redis
}
#[must_use]
pub fn consumer_conns(&self) -> &[MultiplexedConnection] {
&self.consumer_conns
}
pub async fn consume(
&self,
key: &str,
timeout_secs: u64,
slot: usize,
) -> SyncResult<Option<Vec<u8>>> {
let res: Option<(String, Vec<u8>)> = if timeout_secs == 0 {
let mut conn = self.redis.connection_clone();
dispatch!(conn, |c| {
redis::cmd("RPOP").arg(key).query_async(c).await
})?
} else if let Some(conn) = self.consumer_conns.get(slot).cloned() {
let mut conn = conn;
redis::cmd("BRPOP")
.arg(key)
.arg(timeout_secs)
.query_async(&mut conn)
.await?
} else {
let mut conn = self.redis.connection_clone();
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(())
}
}
async fn open_multiplexed(redis: &AppRedis) -> SyncResult<MultiplexedConnection> {
let url = redis.config().url().to_string();
let client = redis::Client::open(url.as_str())
.map_err(|e| SyncError::Internal(format!("open redis client for consumer: {e}")))?;
let config = redis::AsyncConnectionConfig::new()
.set_response_timeout(None)
.set_connection_timeout(Some(redis.config().connection_timeout()));
client
.get_multiplexed_async_connection_with_config(&config)
.await
.map_err(|e| SyncError::Internal(format!("open multiplexed consumer connection: {e}")))
}