use futures_util::stream::StreamExt;
use log::{error, info};
use redis::{Client, Commands, RedisResult};
use serde_json;
use std::future::Future;
use std::pin::Pin;
use crate::types::{
channels::{ChannelMessage, ListChannel, ListMessage, SubscriptionChannel},
provider::{
AsyncCallback, CacheProvider, Connectable, FifoProvider, ProviderTrait, PubSubProvider,
PubSubProviderError, SortedSetProvider,
},
};
pub struct Redis {
pub url: String,
pub client: Option<Client>,
}
impl ProviderTrait for Redis {
fn new(url: String) -> Self {
Redis { url, client: None }
}
}
impl Connectable for Redis {
fn set_client(&mut self) {
let client = redis::Client::open(self.url.clone()).expect("invalid Redis URL");
self.client = Some(client);
}
fn get_connection(&self) -> &Client {
self.client.as_ref().unwrap()
}
fn get_mut_connection(&mut self) -> &mut Client {
self.client.as_mut().unwrap()
}
}
impl PubSubProvider for Redis
where
Redis: Connectable,
{
type Channels = SubscriptionChannel;
fn subscribe(
&self,
callback: AsyncCallback<ChannelMessage>,
channel: Self::Channels,
) -> Pin<Box<dyn Future<Output = Result<(), PubSubProviderError>> + Send + '_>> {
Box::pin(async move {
let client = self.get_connection();
let mut subscriber = client
.get_async_pubsub()
.await
.map_err(|_| "Failed to get async pubsub")?;
subscriber
.subscribe(&[channel.get_channel()])
.await
.map_err(|_| "Failed to subscribe to channel")?;
let mut stream = subscriber.into_on_message();
while let Some(msg) = stream.next().await {
let payload = msg.get_payload::<String>();
if let Ok(json_str) = payload {
if let Ok(channel_msg) = serde_json::from_str::<ChannelMessage>(&json_str) {
info!(
"[{}] New message: {}",
channel_msg.channel.get_channel(),
&json_str
);
callback(channel_msg).await;
}
}
}
Ok(())
})
}
fn publish(&mut self, channel: Self::Channels, message: ChannelMessage) {
let client = self.get_mut_connection();
if let Ok(json_str) = serde_json::to_string(&message) {
let _ = client.publish::<&str, String, String>(channel.get_channel(), json_str);
}
}
}
impl FifoProvider for Redis
where
Redis: Connectable,
{
fn pop(&self, channel_name: ListChannel, job_key: &str) -> Result<ListMessage, &'static str> {
let list_name = self.get_list_name(channel_name.get_channel(), job_key);
let client = self.get_connection();
let mut client = client.clone();
let result = client
.brpop(&list_name, 0.0)
.map(|mut items: Vec<String>| items.pop().unwrap_or_default());
match result {
Ok(json_str) => serde_json::from_str::<serde_json::Value>(&json_str)
.map(|value| ListMessage {
meta: None,
body: Some(value),
})
.map_err(|_| "Failed to deserialize list message"),
Err(_) => Err("Failed to pop from list"),
}
}
fn push(&self, channel_name: ListChannel, job_key: &str, job_payload: serde_json::Value) {
let client: &Client = self.get_connection();
let mut client = client.clone();
let Ok(json_str) = serde_json::to_string(&job_payload) else {
return;
};
let list_name = self.get_list_name(channel_name.get_channel(), job_key);
if let Err(e) = client.rpush::<_, _, ()>(&list_name, json_str) {
error!("Redis RPUSH {list_name} failed: {e}");
}
}
}
impl CacheProvider for Redis
where
Redis: Connectable,
{
fn get(
&self,
cache_space: crate::types::channels::CacheSpace,
key: &str,
app_name: Option<String>,
) -> String {
let client: &Client = self.get_connection();
let mut client = client.clone();
let redis_key = self.get_key_name(cache_space, key, app_name);
let result: RedisResult<String> = client.get(&redis_key);
match result {
Ok(value) => value,
Err(_) => String::new(),
}
}
fn set(
&self,
cache_space: crate::types::channels::CacheSpace,
key: &str,
value: serde_json::Value,
app_name: Option<String>,
) {
let client: &Client = self.get_connection();
let mut client = client.clone();
let redis_key = self.get_key_name(cache_space, key, app_name);
let cache_value = self.get_value(None, value);
let Ok(json_str) = serde_json::to_string(&cache_value) else {
return;
};
if let Err(e) = client.set::<_, _, ()>(&redis_key, json_str) {
error!("Redis SET {redis_key} failed: {e}");
}
}
fn delete(
&self,
cache_space: crate::types::channels::CacheSpace,
key: &str,
app_name: Option<String>,
) {
let client: &Client = self.get_connection();
let mut client = client.clone();
let redis_key = self.get_key_name(cache_space, key, app_name);
if let Err(e) = client.del::<_, ()>(&redis_key) {
error!("Redis DEL {redis_key} failed: {e}");
}
}
fn set_with_ttl(
&self,
cache_space: crate::types::channels::CacheSpace,
key: &str,
value: serde_json::Value,
app_name: Option<String>,
ttl_secs: u64,
) {
let client: &Client = self.get_connection();
let mut client = client.clone();
let redis_key = self.get_key_name(cache_space, key, app_name);
let cache_value = self.get_value(None, value);
let Ok(json_str) = serde_json::to_string(&cache_value) else {
return;
};
if let Err(e) = client.set_ex::<_, _, ()>(&redis_key, json_str, ttl_secs) {
error!("Redis SETEX {redis_key} failed: {e}");
}
}
}
impl SortedSetProvider for Redis
where
Redis: Connectable,
{
fn zadd(&self, key: &str, score: f64, member: &str) -> Result<u32, &'static str> {
let client = self.get_connection();
let mut client = client.clone();
client
.zadd(key, member, score)
.map_err(|_| "Failed to ZADD")
}
fn zremrangebyscore(&self, key: &str, min: f64, max: f64) -> Result<u32, &'static str> {
let client = self.get_connection();
let mut client = client.clone();
client
.zrembyscore(key, min, max)
.map_err(|_| "Failed to ZREMRANGEBYSCORE")
}
fn zcard(&self, key: &str) -> Result<u32, &'static str> {
let client = self.get_connection();
let mut client = client.clone();
client.zcard(key).map_err(|_| "Failed to ZCARD")
}
fn zrangebyscore_withscores(
&self,
key: &str,
min: f64,
max: f64,
limit: Option<usize>,
) -> Result<Vec<(String, f64)>, &'static str> {
let client = self.get_connection();
let mut client = client.clone();
let result: RedisResult<Vec<(String, f64)>> = if let Some(count) = limit {
redis::cmd("ZRANGEBYSCORE")
.arg(key)
.arg(min)
.arg(max)
.arg("WITHSCORES")
.arg("LIMIT")
.arg(0)
.arg(count)
.query(&mut client)
} else {
redis::cmd("ZRANGEBYSCORE")
.arg(key)
.arg(min)
.arg(max)
.arg("WITHSCORES")
.query(&mut client)
};
result.map_err(|_| "Failed to ZRANGEBYSCORE")
}
fn expire(&self, key: &str, ttl_secs: u64) -> Result<(), &'static str> {
let client = self.get_connection();
let mut client = client.clone();
redis::cmd("EXPIRE")
.arg(key)
.arg(ttl_secs)
.query::<()>(&mut client)
.map_err(|_| "Failed to EXPIRE")
}
fn eval_script(
&self,
script: &str,
keys: &[&str],
args: &[&str],
) -> Result<Vec<i64>, &'static str> {
let client = self.get_connection();
let mut client = client.clone();
let cmd = redis::Script::new(script);
let mut invocation = cmd.prepare_invoke();
for k in keys {
invocation.key(*k);
}
for a in args {
invocation.arg(*a);
}
invocation
.invoke::<Vec<i64>>(&mut client)
.map_err(|_| "Failed to evaluate Lua script")
}
}