use std::collections::HashMap;
use std::sync::{Arc, Mutex};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CachingStrategy {
Disabled,
Redis,
InMemory,
None,
}
pub trait StorageAccess<K, V>: Send + Sync
where
K: std::hash::Hash + Eq + Clone + Send + Sync + 'static,
V: Clone + Send + Sync + 'static,
{
fn get(&self, key: &K) -> Option<V>;
fn put(&self, key: K, value: V);
fn invalidate(&self, key: &K);
fn contains(&self, key: &K) -> bool {
self.get(key).is_some()
}
fn clear(&self);
}
pub struct InMemoryStorageAccess<K, V>
where
K: std::hash::Hash + Eq + Clone + Send + Sync + 'static,
V: Clone + Send + Sync + 'static,
{
cache: Arc<Mutex<HashMap<K, V>>>,
max_capacity: usize,
}
impl<K, V> InMemoryStorageAccess<K, V>
where
K: std::hash::Hash + Eq + Clone + Send + Sync + 'static,
V: Clone + Send + Sync + 'static,
{
pub fn new(max_capacity: u64) -> Self {
Self {
cache: Arc::new(Mutex::new(HashMap::new())),
max_capacity: max_capacity as usize,
}
}
pub fn for_region(_region: &str) -> Self {
Self::new(1024)
}
}
impl<K, V> StorageAccess<K, V> for InMemoryStorageAccess<K, V>
where
K: std::hash::Hash + Eq + Clone + Send + Sync + 'static,
V: Clone + Send + Sync + 'static,
{
fn get(&self, key: &K) -> Option<V> {
self.cache.lock().unwrap().get(key).cloned()
}
fn put(&self, key: K, value: V) {
let mut map = self.cache.lock().unwrap();
if map.len() >= self.max_capacity {
if let Some(k) = map.keys().next().cloned() {
map.remove(&k);
}
}
map.insert(key, value);
}
fn invalidate(&self, key: &K) {
self.cache.lock().unwrap().remove(key);
}
fn clear(&self) {
self.cache.lock().unwrap().clear();
}
}
use redis::{AsyncCommands, Client, RedisResult};
#[derive(Clone)]
pub struct RedisStorage {
client: Client,
}
impl RedisStorage {
pub fn new(url: &str) -> RedisResult<Self> {
Ok(Self {
client: Client::open(url)?,
})
}
pub fn from_env() -> RedisResult<Self> {
let url = std::env::var("REDIS_URL")
.or_else(|_| std::env::var("REDIS_URI"))
.unwrap_or_else(|_| "redis://127.0.0.1:6379/".to_string());
Self::new(&url)
}
pub fn client(&self) -> &Client {
&self.client
}
async fn conn(&self) -> RedisResult<redis::aio::MultiplexedConnection> {
self.client.get_multiplexed_async_connection().await
}
pub async fn get_value(&self, key: &str) -> RedisResult<Option<String>> {
let mut conn = self.conn().await?;
conn.get(key).await
}
pub async fn set_value(&self, key: &str, value: &str) -> RedisResult<()> {
let mut conn = self.conn().await?;
conn.set::<_, _, ()>(key, value).await
}
pub async fn set_value_with_expiration(&self, key: &str, value: &str, seconds: u64) -> RedisResult<()> {
let mut conn = self.conn().await?;
conn.set_ex::<_, _, ()>(key, value, seconds).await
}
pub async fn set_value_if_absent(&self, key: &str, value: &str) -> RedisResult<bool> {
let mut conn = self.conn().await?;
let res: Option<String> = redis::cmd("SET")
.arg(key)
.arg(value)
.arg("NX")
.query_async(&mut conn)
.await?;
Ok(res.is_some())
}
pub async fn delete_key(&self, key: &str) -> RedisResult<i64> {
let mut conn = self.conn().await?;
conn.del(key).await
}
pub async fn unlink_key(&self, key: &str) -> RedisResult<i64> {
let mut conn = self.conn().await?;
redis::cmd("UNLINK").arg(key).query_async(&mut conn).await
}
pub async fn delete_keys(&self, keys: &[String]) -> RedisResult<i64> {
if keys.is_empty() {
return Ok(0);
}
let mut conn = self.conn().await?;
conn.del(keys).await
}
pub async fn key_exists(&self, key: &str) -> RedisResult<bool> {
let mut conn = self.conn().await?;
let v: i64 = conn.exists(key).await?;
Ok(v == 1)
}
pub async fn increment_value(&self, key: &str) -> RedisResult<i64> {
let mut conn = self.conn().await?;
conn.incr(key, 1i64).await
}
pub async fn increment_value_by(&self, key: &str, delta: i64) -> RedisResult<i64> {
let mut conn = self.conn().await?;
conn.incr(key, delta).await
}
pub async fn decrement_value(&self, key: &str) -> RedisResult<i64> {
let mut conn = self.conn().await?;
conn.decr(key, 1i64).await
}
pub async fn decrement_value_by(&self, key: &str, delta: i64) -> RedisResult<i64> {
let mut conn = self.conn().await?;
conn.decr(key, delta).await
}
pub async fn idle_time(&self, key: &str) -> RedisResult<i64> {
let mut conn = self.conn().await?;
redis::cmd("OBJECT")
.arg("IDLETIME")
.arg(key)
.query_async(&mut conn)
.await
}
pub async fn set_expiration(&self, key: &str, seconds: u64) -> RedisResult<bool> {
let mut conn = self.conn().await?;
let v: i64 = conn.expire(key, seconds as i64).await?;
Ok(v == 1)
}
pub async fn get_time_to_live(&self, key: &str) -> RedisResult<i64> {
let mut conn = self.conn().await?;
conn.ttl(key).await
}
pub async fn get_multiple_values(&self, keys: &[String]) -> RedisResult<Vec<Option<String>>> {
if keys.is_empty() {
return Ok(vec![]);
}
let mut conn = self.conn().await?;
conn.mget(keys).await
}
pub async fn set_multiple_values(&self, kv: &HashMap<String, String>) -> RedisResult<()> {
if kv.is_empty() {
return Ok(());
}
let mut conn = self.conn().await?;
let mut args: Vec<String> = Vec::with_capacity(kv.len() * 2);
for (k, v) in kv {
args.push(k.clone());
args.push(v.clone());
}
redis::cmd("MSET").arg(args).query_async::<()>(&mut conn).await?;
Ok(())
}
pub async fn get_hash_value(&self, key: &str, field: &str) -> RedisResult<Option<String>> {
let mut conn = self.conn().await?;
conn.hget(key, field).await
}
pub async fn set_hash_value(&self, key: &str, field: &str, value: &str) -> RedisResult<()> {
let mut conn = self.conn().await?;
conn.hset::<_, _, _, ()>(key, field, value).await
}
pub async fn delete_hash_field(&self, key: &str, field: &str) -> RedisResult<i64> {
let mut conn = self.conn().await?;
conn.hdel(key, field).await
}
pub async fn set_hash_values(&self, key: &str, field_values: &HashMap<String, String>) -> RedisResult<()> {
if field_values.is_empty() {
return Ok(());
}
let mut conn = self.conn().await?;
let mut cmd = redis::cmd("HSET");
cmd.arg(key);
for (f, v) in field_values {
cmd.arg(f).arg(v);
}
cmd.query_async::<()>(&mut conn).await?;
Ok(())
}
pub async fn set_hash_value_if_absent(&self, key: &str, field: &str, value: &str) -> RedisResult<bool> {
let mut conn = self.conn().await?;
let v: i64 = conn.hset_nx(key, field, value).await?;
Ok(v == 1)
}
pub async fn increment_hash_field(&self, key: &str, field: &str, delta: i64) -> RedisResult<i64> {
let mut conn = self.conn().await?;
conn.hincr(key, field, delta).await
}
pub async fn decrement_hash_field(&self, key: &str, field: &str, delta: i64) -> RedisResult<i64> {
self.increment_hash_field(key, field, -delta).await
}
pub async fn get_all_hash_fields(&self, key: &str) -> RedisResult<HashMap<String, String>> {
let mut conn = self.conn().await?;
conn.hgetall(key).await
}
pub async fn get_hash_keys(&self, key: &str) -> RedisResult<Vec<String>> {
let mut conn = self.conn().await?;
conn.hkeys(key).await
}
pub async fn get_hash_values(&self, key: &str) -> RedisResult<Vec<String>> {
let mut conn = self.conn().await?;
conn.hvals(key).await
}
pub async fn hash_field_exists(&self, key: &str, field: &str) -> RedisResult<bool> {
let mut conn = self.conn().await?;
let v: bool = conn.hexists(key, field).await?;
Ok(v)
}
pub async fn push_to_list_start(&self, key: &str, value: &str) -> RedisResult<i64> {
let mut conn = self.conn().await?;
conn.lpush(key, value).await
}
pub async fn push_to_list_end(&self, key: &str, value: &str) -> RedisResult<i64> {
let mut conn = self.conn().await?;
conn.rpush(key, value).await
}
pub async fn pop_from_list_start(&self, key: &str) -> RedisResult<Option<String>> {
let mut conn = self.conn().await?;
conn.lpop(key, None).await
}
pub async fn pop_from_list_end(&self, key: &str) -> RedisResult<Option<String>> {
let mut conn = self.conn().await?;
conn.rpop(key, None).await
}
pub async fn get_list_length(&self, key: &str) -> RedisResult<i64> {
let mut conn = self.conn().await?;
conn.llen(key).await
}
pub async fn get_list_range(&self, key: &str, start: i64, stop: i64) -> RedisResult<Vec<String>> {
let mut conn = self.conn().await?;
conn.lrange(key, start as isize, stop as isize).await
}
pub async fn add_to_set(&self, key: &str, member: &str) -> RedisResult<i64> {
let mut conn = self.conn().await?;
conn.sadd(key, member).await
}
pub async fn remove_from_set(&self, key: &str, member: &str) -> RedisResult<i64> {
let mut conn = self.conn().await?;
conn.srem(key, member).await
}
pub async fn get_set_members(&self, key: &str) -> RedisResult<Vec<String>> {
let mut conn = self.conn().await?;
conn.smembers(key).await
}
pub async fn is_set_member(&self, key: &str, member: &str) -> RedisResult<bool> {
let mut conn = self.conn().await?;
conn.sismember(key, member).await
}
pub async fn get_set_size(&self, key: &str) -> RedisResult<i64> {
let mut conn = self.conn().await?;
conn.scard(key).await
}
pub async fn add_to_sorted_set(&self, key: &str, score: f64, member: &str) -> RedisResult<i64> {
let mut conn = self.conn().await?;
conn.zadd(key, member, score).await
}
pub async fn get_sorted_set_range(&self, key: &str, start: i64, stop: i64) -> RedisResult<Vec<String>> {
let mut conn = self.conn().await?;
conn.zrange(key, start as isize, stop as isize).await
}
pub async fn remove_from_sorted_set(&self, key: &str, member: &str) -> RedisResult<i64> {
let mut conn = self.conn().await?;
conn.zrem(key, member).await
}
pub async fn get_sorted_set_score(&self, key: &str, member: &str) -> RedisResult<Option<f64>> {
let mut conn = self.conn().await?;
conn.zscore(key, member).await
}
pub async fn get_sorted_set_size(&self, key: &str) -> RedisResult<i64> {
let mut conn = self.conn().await?;
conn.zcard(key).await
}
pub async fn remove_expiration(&self, key: &str) -> RedisResult<bool> {
let mut conn = self.conn().await?;
let v: i64 = redis::cmd("PERSIST").arg(key).query_async(&mut conn).await?;
Ok(v == 1)
}
pub async fn rename_key(&self, old_key: &str, new_key: &str) -> RedisResult<()> {
let mut conn = self.conn().await?;
redis::cmd("RENAME").arg(old_key).arg(new_key).query_async::<()>(&mut conn).await?;
Ok(())
}
pub async fn scan_keys(&self, pattern: &str, count: usize) -> RedisResult<Vec<String>> {
let mut conn = self.conn().await?;
let mut cursor: u64 = 0;
let mut all = Vec::new();
loop {
let (next_cursor, keys): (u64, Vec<String>) = redis::cmd("SCAN")
.arg(cursor)
.arg("MATCH")
.arg(pattern)
.arg("COUNT")
.arg(count)
.query_async(&mut conn)
.await?;
all.extend(keys);
if next_cursor == 0 {
break;
}
cursor = next_cursor;
}
Ok(all)
}
pub async fn scan_keys_default(&self, pattern: &str) -> RedisResult<Vec<String>> {
self.scan_keys(pattern, 250).await
}
#[deprecated(note = "Use scan_keys instead to avoid blocking Redis")]
pub async fn find_keys(&self, pattern: &str) -> RedisResult<Vec<String>> {
self.scan_keys_default(pattern).await
}
}
pub struct RedisStorageAccess {
storage: RedisStorage,
prefix: String,
ttl_seconds: u64,
}
impl RedisStorageAccess {
pub fn new(storage: RedisStorage, region_name: &str) -> Self {
Self {
storage,
prefix: format!("hibernate:cache:{}:", region_name),
ttl_seconds: 3600,
}
}
pub fn from_url(url: &str, region_name: &str) -> RedisResult<Self> {
Ok(Self::new(RedisStorage::new(url)?, region_name))
}
pub fn with_ttl(mut self, ttl_seconds: u64) -> Self {
self.ttl_seconds = ttl_seconds;
self
}
fn build_key<K: ToString>(&self, key: &K) -> String {
format!("{}{}", self.prefix, key.to_string())
}
fn serialize<T: serde::Serialize>(value: &T) -> Result<Vec<u8>, serde_json::Error> {
serde_json::to_vec(value)
}
fn deserialize<T: serde::de::DeserializeOwned>(bytes: &[u8]) -> Result<T, serde_json::Error> {
serde_json::from_slice(bytes)
}
pub async fn contains_async<K: ToString + Send + Sync>(&self, key: &K) -> RedisResult<bool> {
self.storage.key_exists(&self.build_key(key)).await
}
pub async fn get_from_cache_async<K, V>(&self, key: &K) -> RedisResult<Option<V>>
where
K: ToString + Send + Sync,
V: serde::de::DeserializeOwned,
{
let raw: Option<Vec<u8>> = {
let mut conn = self.storage.conn().await?;
let k = self.build_key(key);
conn.get(k).await?
};
match raw {
None => Ok(None),
Some(bytes) => match Self::deserialize::<V>(&bytes) {
Ok(v) => Ok(Some(v)),
Err(_) => Ok(None),
},
}
}
pub async fn put_into_cache_async<K, V>(&self, key: &K, value: &V) -> RedisResult<()>
where
K: ToString + Send + Sync,
V: serde::Serialize,
{
let bytes = Self::serialize(value).map_err(|e| {
redis::RedisError::from((
redis::ErrorKind::Io,
"serialization failed",
e.to_string(),
))
})?;
let mut conn = self.storage.conn().await?;
let k = self.build_key(key);
conn.set_ex::<_, _, ()>(k, bytes, self.ttl_seconds).await
}
pub async fn remove_from_cache_async<K: ToString + Send + Sync>(&self, key: &K) -> RedisResult<()> {
self.storage.unlink_key(&self.build_key(key)).await.map(|_| ())
}
pub async fn clear_cache_async(&self) -> RedisResult<()> {
let pattern = format!("{}*", self.prefix);
let keys = self.storage.scan_keys(&pattern, 750).await?;
if keys.is_empty() {
return Ok(());
}
let mut conn = self.storage.conn().await?;
for key in keys {
let _: () = redis::cmd("UNLINK").arg(key).query_async(&mut conn).await?;
}
Ok(())
}
fn block_on<F: Future>(fut: F) -> F::Output {
if let Ok(handle) = tokio::runtime::Handle::try_current() {
tokio::task::block_in_place(|| handle.block_on(fut))
} else {
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap()
.block_on(fut)
}
}
}
impl<K, V> StorageAccess<K, V> for RedisStorageAccess
where
K: std::hash::Hash + Eq + Clone + ToString + Send + Sync + 'static,
V: Clone + serde::Serialize + serde::de::DeserializeOwned + Send + Sync + 'static,
{
fn get(&self, key: &K) -> Option<V> {
Self::block_on(self.get_from_cache_async(key)).ok().flatten()
}
fn put(&self, key: K, value: V) {
let _ = Self::block_on(self.put_into_cache_async(&key, &value));
}
fn invalidate(&self, key: &K) {
let _ = Self::block_on(self.remove_from_cache_async(key));
}
fn clear(&self) {
let _ = Self::block_on(self.clear_cache_async());
}
}
use std::collections::hash_map::Entry;
pub struct InMemoryRegionFactory {
regions: Mutex<HashMap<String, Arc<InMemoryStorageAccess<String, Vec<u8>>>>>,
}
impl InMemoryRegionFactory {
pub fn new() -> Self {
Self {
regions: Mutex::new(HashMap::new()),
}
}
pub fn get_or_create(&self, region_name: &str) -> Arc<InMemoryStorageAccess<String, Vec<u8>>> {
let mut map = self.regions.lock().unwrap();
match map.entry(region_name.to_string()) {
Entry::Occupied(o) => o.get().clone(),
Entry::Vacant(v) => {
let access = Arc::new(InMemoryStorageAccess::for_region(region_name));
v.insert(access.clone());
access
}
}
}
pub fn clear_all(&self) {
let map = self.regions.lock().unwrap();
for access in map.values() {
access.clear();
}
}
}
impl Default for InMemoryRegionFactory {
fn default() -> Self {
Self::new()
}
}
pub struct RedisRegionFactory {
storage: RedisStorage,
regions: Mutex<HashMap<String, Arc<RedisStorageAccess>>>,
}
impl RedisRegionFactory {
pub fn new(storage: RedisStorage) -> Self {
Self {
storage,
regions: Mutex::new(HashMap::new()),
}
}
pub fn from_url(url: &str) -> RedisResult<Self> {
Ok(Self::new(RedisStorage::new(url)?))
}
pub fn get_or_create(&self, region_name: &str) -> Arc<RedisStorageAccess> {
let mut map = self.regions.lock().unwrap();
match map.entry(region_name.to_string()) {
Entry::Occupied(o) => o.get().clone(),
Entry::Vacant(v) => {
let access = Arc::new(RedisStorageAccess::new(self.storage.clone(), region_name));
v.insert(access.clone());
access
}
}
}
}