209a1c0
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>,
/// cluster 模式：每个消费者一组到所有 master 的连接（按 slot 选择）
cluster_consumer_groups: Vec<Vec<MultiplexedConnection>>,
}
impl SyncQueue {
#[must_use]
pub fn new(redis: AppRedis) -> Self {
Self {
redis,
consumer_conns: Vec::new(),
cluster_consumer_groups: Vec::new(),
}
}
pub async fn with_consumer_conns(redis: AppRedis, count: usize) -> SyncResult<Self> {
let count = count.max(1);
// cluster 模式：BRPOP 是对单个 key 的操作，必须打到该 key slot 的 master，
// 否则会收到 MOVED。这里按哨兵轮询的方式给每个消费者预建到所有 master 的连接，
// 消费时再按 slot 选择对应连接。standalone 模式维持单连接。
let mut consumer_conns = Vec::with_capacity(count);
if redis.is_cluster() {
let masters = redis.cluster_master_connections().await.map_err(|e| {
SyncError::Internal(format!("open cluster master connections: {e}"))
})?;
if masters.is_empty() {
return Err(SyncError::Internal(
"redis cluster reported no master nodes".to_string(),
));
}
// 每个消费者一组 master 连接；取模轮询避免所有消费者压在同一节点
let mut groups = Vec::with_capacity(count);
for i in 0..count {
let mut group = Vec::with_capacity(masters.len());
for offset in 0..masters.len() {
let idx = (i + offset) % masters.len();
group.push(masters[idx].clone());
}
groups.push(group);
}
tracing::debug!(
count,
masters = masters.len(),
"sync queue cluster consumer connections opened"
);
return Ok(Self {
redis,
consumer_conns: Vec::new(),
cluster_consumer_groups: groups,
});
}
for _ in 0..count {
let conn = open_multiplexed(&redis).await?;
consumer_conns.push(conn);
}
tracing::debug!(count, "sync queue consumer connections opened");
Ok(Self {
redis,
consumer_conns,
cluster_consumer_groups: Vec::new(),
})
}
#[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 !self.cluster_consumer_groups.is_empty() {
// cluster：BRPOP 必须落在持有该 key slot 的 master 上，否则会 MOVED。
// 先算 KEYSLOT，再到负责该 slot 的 master 上消费。
let conn = self.cluster_master_for_key(key).await?;
let mut conn = conn;
redis::cmd("BRPOP")
.arg(key)
.arg(timeout_secs)
.query_async(&mut conn)
.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
})?;
tracing::debug!(
topic = key,
queue_len = n,
payload_bytes = payload.len(),
"sync queue message pushed"
);
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)
}
/// 计算 key 所属 slot，并返回负责该 slot 的 master 连接（cluster 模式）。
async fn cluster_master_for_key(&self, key: &str) -> SyncResult<MultiplexedConnection> {
let group = self
.cluster_consumer_groups
.first()
.ok_or_else(|| SyncError::Internal("no cluster consumer connections".to_string()))?;
let mut conn = self.redis.connection_clone();
let key_slot: u16 = {
let RedisConnection::Cluster(cluster) = &mut conn else {
return Err(SyncError::Internal(
"cluster master lookup on non-cluster connection".to_string(),
));
};
redis::cmd("CLUSTER")
.arg("KEYSLOT")
.arg(key)
.query_async(cluster)
.await
.map_err(|e| SyncError::Internal(format!("CLUSTER KEYSLOT failed: {e}")))?
};
let ranges = self
.redis
.cluster_slot_ranges()
.await
.map_err(|e| SyncError::Internal(format!("cluster slot ranges: {e}")))?;
let index = ranges
.iter()
.find(|(start, end, _)| (*start..=*end).contains(&(key_slot as i64)))
.map(|(_, _, idx)| *idx)
.ok_or_else(|| SyncError::Internal(format!("no master owns slot {key_slot}")))?;
group
.get(index)
.or_else(|| group.first())
.cloned()
.ok_or_else(|| SyncError::Internal("cluster consumer connection missing".to_string()))
}
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> {
// cluster 模式配置是逗号分隔的节点列表；BRPOP 消费者连单个节点即可
// （`redis::Client::open` 无法解析多节点串）。
let url = redis
.config()
.url()
.split(',')
.map(str::trim)
.find(|node| !node.is_empty())
.ok_or_else(|| SyncError::Internal("redis url contains no nodes".to_string()))?
.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}")))
}