use std::sync::Arc;
use futures::stream::StreamExt;
use tokio::task::JoinHandle;
use crate::error::{OxCacheError, OxCacheResult};
const DEFAULT_RECONNECT_ATTEMPTS: usize = 5;
const RECONNECT_BACKOFF_BASE_MS: u64 = 500;
const RECONNECT_BACKOFF_MAX_MS: u64 = 5000;
pub struct RedisPubSub {
client: redis::Client,
publish_conn: redis::aio::ConnectionManager,
tasks: std::sync::Mutex<Vec<JoinHandle<()>>>,
max_reconnect_attempts: usize,
}
impl RedisPubSub {
pub async fn new(url: &str) -> OxCacheResult<Self> {
let client = redis::Client::open(url)
.map_err(|e| OxCacheError::BackendError(format!("pubsub client open: {e}")))?;
let publish_conn = redis::aio::ConnectionManager::new(client.clone())
.await
.map_err(|e| OxCacheError::BackendError(format!("pubsub connection: {e}")))?;
Ok(Self {
client,
publish_conn,
tasks: std::sync::Mutex::new(Vec::new()),
max_reconnect_attempts: DEFAULT_RECONNECT_ATTEMPTS,
})
}
#[must_use]
pub fn with_reconnect_attempts(mut self, attempts: usize) -> Self {
self.max_reconnect_attempts = attempts;
self
}
pub async fn publish(&self, channel: &str, message: &str) -> OxCacheResult<i64> {
let mut conn = self.publish_conn.clone();
redis::cmd("PUBLISH")
.arg(channel)
.arg(message)
.query_async::<i64>(&mut conn)
.await
.map_err(|e| OxCacheError::BackendError(format!("pubsub publish: {e}")))
}
pub async fn subscribe(
&self,
channel: &str,
handler: Arc<dyn Fn(String) + Send + Sync>,
) -> OxCacheResult<()> {
let channel = channel.to_string();
let client = self.client.clone();
let max_attempts = self.max_reconnect_attempts;
let (ready_tx, ready_rx) = tokio::sync::oneshot::channel::<OxCacheResult<()>>();
let handle = tokio::spawn(async move {
let mut pubsub = match client.get_async_pubsub().await {
Ok(p) => p,
Err(e) => {
let _ = ready_tx.send(Err(OxCacheError::BackendError(format!(
"pubsub connect: {e}"
))));
return;
}
};
if let Err(e) = pubsub.subscribe(&channel).await {
let _ = ready_tx.send(Err(OxCacheError::BackendError(format!(
"pubsub subscribe: {e}"
))));
return;
}
let _ = ready_tx.send(Ok(()));
let mut attempt = 0usize;
loop {
{
let mut msg_stream = pubsub.on_message();
while let Some(msg) = msg_stream.next().await {
let payload: Result<String, _> = msg.get_payload();
match payload {
Ok(payload_str) => {
let handler_clone = handler.clone();
let result = std::panic::catch_unwind(
std::panic::AssertUnwindSafe(move || {
handler_clone(payload_str);
}),
);
if result.is_err() {
tracing::warn!(
"pubsub handler panicked: channel={channel}, continue"
);
}
}
Err(e) => {
tracing::warn!(
"pubsub payload parse failed: channel={channel}, err={e}"
);
}
}
}
}
if attempt >= max_attempts {
tracing::error!(
"pubsub stream ended, max reconnect attempts ({max_attempts}) reached: channel={channel}"
);
return;
}
attempt += 1;
let backoff = std::time::Duration::from_millis(
(RECONNECT_BACKOFF_BASE_MS * attempt as u64).min(RECONNECT_BACKOFF_MAX_MS),
);
tracing::warn!(
"pubsub stream ended (connection lost): channel={channel}, reconnect {attempt}/{max_attempts} in {backoff:?}"
);
tokio::time::sleep(backoff).await;
match client.get_async_pubsub().await {
Ok(mut p) => {
if let Err(e) = p.subscribe(&channel).await {
tracing::error!(
"pubsub resubscribe failed: channel={channel}, err={e}"
);
return;
}
pubsub = p;
attempt = 0;
}
Err(e) => {
tracing::error!("pubsub reconnect failed: channel={channel}, err={e}");
return;
}
}
}
});
let result = match ready_rx.await {
Ok(r) => r,
Err(_) => Err(OxCacheError::BackendError(
"pubsub subscribe task ended before ready".to_string(),
)),
};
if result.is_ok() {
let mut tasks = self.tasks.lock().unwrap();
tasks.retain(|h| !h.is_finished());
tasks.push(handle);
}
result
}
pub fn shutdown(&self) -> usize {
let mut tasks = self.tasks.lock().unwrap();
let stopped = tasks.len();
for handle in tasks.drain(..) {
handle.abort();
}
stopped
}
}
impl Drop for RedisPubSub {
fn drop(&mut self) {
for handle in self.tasks.lock().unwrap().drain(..) {
handle.abort();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
async fn redis_reachable() -> bool {
tokio::net::TcpStream::connect("127.0.0.1:6379")
.await
.is_ok()
}
#[tokio::test]
async fn new_with_unreachable_port_fails_fast() {
let result = RedisPubSub::new("redis://127.0.0.1:1").await;
assert!(result.is_err(), "不可达端口构造应 fail-fast 返回 Err");
}
#[tokio::test]
async fn new_rejects_malformed_url() {
assert!(RedisPubSub::new("not-a-redis-url").await.is_err());
}
#[tokio::test]
async fn subscribe_receives_published_message() {
if !redis_reachable().await {
eprintln!("[SKIP] Redis 不可达(127.0.0.1:6379 未监听)");
return;
}
let ps = RedisPubSub::new("redis://127.0.0.1:6379").await.unwrap();
let (tx, rx) = std::sync::mpsc::channel::<String>();
ps.subscribe(
"oxcache-pubsub-test-e2e",
Arc::new(move |msg| {
let _ = tx.send(msg);
}),
)
.await
.expect("subscribe 应成功");
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
ps.publish("oxcache-pubsub-test-e2e", "hello-pubsub")
.await
.unwrap();
let received = rx.recv_timeout(std::time::Duration::from_secs(2));
assert_eq!(received.ok().as_deref(), Some("hello-pubsub"));
ps.shutdown();
}
#[tokio::test]
async fn subscribe_handler_panic_does_not_interrupt() {
if !redis_reachable().await {
eprintln!("[SKIP] Redis 不可达(127.0.0.1:6379 未监听)");
return;
}
let ps = RedisPubSub::new("redis://127.0.0.1:6379").await.unwrap();
let (tx, rx) = std::sync::mpsc::channel::<String>();
let counter = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let counter_clone = counter.clone();
ps.subscribe(
"oxcache-pubsub-test-panic",
Arc::new(move |msg| {
counter_clone.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
if msg == "panic-trigger" {
panic!("intentional handler panic");
}
let _ = tx.send(msg);
}),
)
.await
.unwrap();
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
ps.publish("oxcache-pubsub-test-panic", "panic-trigger")
.await
.unwrap();
ps.publish("oxcache-pubsub-test-panic", "after-panic")
.await
.unwrap();
let received = rx.recv_timeout(std::time::Duration::from_secs(2));
assert_eq!(
received.ok().as_deref(),
Some("after-panic"),
"panic 后续消息应继续投递"
);
assert!(
counter.load(std::sync::atomic::Ordering::SeqCst) >= 2,
"panic 前后的消息都应到达 handler"
);
ps.shutdown();
}
#[tokio::test]
async fn shutdown_stops_message_delivery() {
if !redis_reachable().await {
eprintln!("[SKIP] Redis 不可达(127.0.0.1:6379 未监听)");
return;
}
let ps = RedisPubSub::new("redis://127.0.0.1:6379").await.unwrap();
let (tx, rx) = std::sync::mpsc::channel::<String>();
ps.subscribe(
"oxcache-pubsub-test-shutdown",
Arc::new(move |msg| {
let _ = tx.send(msg);
}),
)
.await
.unwrap();
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
let stopped = ps.shutdown();
assert_eq!(stopped, 1, "应停止 1 个订阅任务");
ps.publish("oxcache-pubsub-test-shutdown", "after-shutdown")
.await
.unwrap();
let leaked = rx.recv_timeout(std::time::Duration::from_millis(300));
assert!(
leaked.is_err(),
"shutdown 后不应再收到消息,实际: {leaked:?}"
);
}
}