feat: init
This commit is contained in:
Vendored
+304
@@ -0,0 +1,304 @@
|
||||
use dashmap::DashMap;
|
||||
use std::collections::HashMap;
|
||||
use std::hash::Hash;
|
||||
use std::sync::Mutex;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
struct CacheEntry<V> {
|
||||
value: V,
|
||||
expires_at: Instant,
|
||||
}
|
||||
|
||||
struct LruNode<K> {
|
||||
key: Option<K>,
|
||||
prev: usize,
|
||||
next: usize,
|
||||
}
|
||||
|
||||
struct LruTracker<K> {
|
||||
nodes: Vec<LruNode<K>>,
|
||||
key_to_idx: HashMap<K, usize>,
|
||||
head: usize,
|
||||
tail: usize,
|
||||
}
|
||||
|
||||
impl<K: Eq + Hash + Clone> LruTracker<K> {
|
||||
fn new() -> Self {
|
||||
let sentinel = LruNode {
|
||||
key: None,
|
||||
prev: 0,
|
||||
next: 0,
|
||||
};
|
||||
Self {
|
||||
nodes: vec![sentinel],
|
||||
key_to_idx: HashMap::new(),
|
||||
head: 0,
|
||||
tail: 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn touch(&mut self, key: &K) {
|
||||
if let Some(&idx) = self.key_to_idx.get(key) {
|
||||
self.detach(idx);
|
||||
self.attach_front(idx);
|
||||
}
|
||||
}
|
||||
|
||||
fn push_front(&mut self, key: K) -> usize {
|
||||
let idx = self.nodes.len();
|
||||
self.nodes.push(LruNode {
|
||||
key: Some(key.clone()),
|
||||
prev: 0,
|
||||
next: 0,
|
||||
});
|
||||
self.key_to_idx.insert(key, idx);
|
||||
self.attach_front(idx);
|
||||
idx
|
||||
}
|
||||
|
||||
fn pop_back(&mut self) -> Option<K> {
|
||||
if self.tail == 0 {
|
||||
return None;
|
||||
}
|
||||
let lru = self.tail;
|
||||
let key = self.nodes[lru].key.take();
|
||||
self.detach(lru);
|
||||
if let Some(ref k) = key {
|
||||
self.key_to_idx.remove(k);
|
||||
}
|
||||
key
|
||||
}
|
||||
|
||||
fn remove(&mut self, key: &K) {
|
||||
if let Some(&idx) = self.key_to_idx.get(key) {
|
||||
self.detach(idx);
|
||||
self.key_to_idx.remove(key);
|
||||
}
|
||||
}
|
||||
|
||||
fn clear(&mut self) {
|
||||
self.key_to_idx.clear();
|
||||
self.nodes.truncate(1);
|
||||
self.head = 0;
|
||||
self.tail = 0;
|
||||
}
|
||||
|
||||
fn len(&self) -> usize {
|
||||
self.key_to_idx.len()
|
||||
}
|
||||
|
||||
fn detach(&mut self, idx: usize) {
|
||||
let prev = self.nodes[idx].prev;
|
||||
let next = self.nodes[idx].next;
|
||||
|
||||
if prev != 0 {
|
||||
self.nodes[prev].next = next;
|
||||
} else {
|
||||
self.head = next;
|
||||
}
|
||||
|
||||
if next != 0 {
|
||||
self.nodes[next].prev = prev;
|
||||
} else {
|
||||
self.tail = prev;
|
||||
}
|
||||
}
|
||||
|
||||
fn attach_front(&mut self, idx: usize) {
|
||||
self.nodes[idx].prev = 0;
|
||||
self.nodes[idx].next = self.head;
|
||||
|
||||
if self.head != 0 {
|
||||
self.nodes[self.head].prev = idx;
|
||||
} else {
|
||||
self.tail = idx;
|
||||
}
|
||||
|
||||
self.head = idx;
|
||||
}
|
||||
}
|
||||
|
||||
pub struct LruTtlCache<K, V> {
|
||||
map: DashMap<K, CacheEntry<V>>,
|
||||
lru: Mutex<LruTracker<K>>,
|
||||
capacity: usize,
|
||||
ttl: Duration,
|
||||
}
|
||||
|
||||
impl<K: Eq + Hash + Clone, V: Clone> LruTtlCache<K, V> {
|
||||
pub fn new(capacity: usize, ttl: Duration) -> Self {
|
||||
Self {
|
||||
map: DashMap::with_capacity(capacity),
|
||||
lru: Mutex::new(LruTracker::new()),
|
||||
capacity,
|
||||
ttl,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get(&self, key: &K) -> Option<V> {
|
||||
let entry = self.map.get(key)?;
|
||||
let expired = entry.expires_at <= Instant::now();
|
||||
let value = entry.value.clone();
|
||||
drop(entry);
|
||||
|
||||
if expired {
|
||||
self.remove(key);
|
||||
return None;
|
||||
}
|
||||
|
||||
if let Ok(mut lru) = self.lru.lock() {
|
||||
lru.touch(key);
|
||||
}
|
||||
|
||||
Some(value)
|
||||
}
|
||||
|
||||
pub fn insert(&self, key: K, value: V) {
|
||||
self.insert_with_ttl(key, value, self.ttl);
|
||||
}
|
||||
|
||||
pub fn insert_with_ttl(&self, key: K, value: V, ttl: Duration) {
|
||||
let now = Instant::now();
|
||||
|
||||
if self.map.contains_key(&key) {
|
||||
self.map.insert(
|
||||
key.clone(),
|
||||
CacheEntry {
|
||||
value,
|
||||
expires_at: now + ttl,
|
||||
},
|
||||
);
|
||||
if let Ok(mut lru) = self.lru.lock() {
|
||||
lru.touch(&key);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
let mut lru = self.lru.lock().unwrap();
|
||||
|
||||
if lru.len() >= self.capacity
|
||||
&& let Some(evicted_key) = lru.pop_back()
|
||||
{
|
||||
self.map.remove(&evicted_key);
|
||||
}
|
||||
|
||||
self.map.insert(
|
||||
key.clone(),
|
||||
CacheEntry {
|
||||
value,
|
||||
expires_at: now + ttl,
|
||||
},
|
||||
);
|
||||
lru.push_front(key);
|
||||
}
|
||||
|
||||
pub fn remove(&self, key: &K) -> Option<V> {
|
||||
if let Ok(mut lru) = self.lru.lock() {
|
||||
lru.remove(key);
|
||||
}
|
||||
self.map.remove(key).map(|(_, entry)| entry.value)
|
||||
}
|
||||
|
||||
pub fn contains(&self, key: &K) -> bool {
|
||||
self.map.contains_key(key)
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
self.map.len()
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.map.is_empty()
|
||||
}
|
||||
|
||||
pub fn clear(&self) {
|
||||
self.map.clear();
|
||||
if let Ok(mut lru) = self.lru.lock() {
|
||||
lru.clear();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_insert_and_get() {
|
||||
let cache = LruTtlCache::new(3, Duration::from_secs(60));
|
||||
cache.insert("a", 1);
|
||||
cache.insert("b", 2);
|
||||
cache.insert("c", 3);
|
||||
|
||||
assert_eq!(cache.get(&"a"), Some(1));
|
||||
assert_eq!(cache.get(&"b"), Some(2));
|
||||
assert_eq!(cache.get(&"c"), Some(3));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_lru_eviction() {
|
||||
let cache = LruTtlCache::new(2, Duration::from_secs(60));
|
||||
cache.insert("a", 1);
|
||||
cache.insert("b", 2);
|
||||
cache.get(&"a");
|
||||
cache.insert("c", 3);
|
||||
|
||||
assert_eq!(cache.get(&"a"), Some(1));
|
||||
assert_eq!(cache.get(&"b"), None);
|
||||
assert_eq!(cache.get(&"c"), Some(3));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ttl_expiry() {
|
||||
let cache = LruTtlCache::new(3, Duration::from_millis(10));
|
||||
cache.insert("a", 1);
|
||||
std::thread::sleep(Duration::from_millis(20));
|
||||
|
||||
assert_eq!(cache.get(&"a"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_update_existing() {
|
||||
let cache = LruTtlCache::new(3, Duration::from_secs(60));
|
||||
cache.insert("a", 1);
|
||||
cache.insert("a", 100);
|
||||
|
||||
assert_eq!(cache.get(&"a"), Some(100));
|
||||
assert_eq!(cache.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_remove() {
|
||||
let cache = LruTtlCache::new(3, Duration::from_secs(60));
|
||||
cache.insert("a", 1);
|
||||
cache.insert("b", 2);
|
||||
|
||||
assert_eq!(cache.remove(&"a"), Some(1));
|
||||
assert_eq!(cache.get(&"a"), None);
|
||||
assert_eq!(cache.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_concurrent_access() {
|
||||
let cache = std::sync::Arc::new(LruTtlCache::new(10, Duration::from_secs(60)));
|
||||
let c1 = cache.clone();
|
||||
let c2 = cache.clone();
|
||||
|
||||
let t1 = std::thread::spawn(move || {
|
||||
for i in 0..100 {
|
||||
c1.insert(i, i * 2);
|
||||
}
|
||||
});
|
||||
|
||||
let t2 = std::thread::spawn(move || {
|
||||
for i in 0..100 {
|
||||
let _ = c2.get(&i);
|
||||
}
|
||||
});
|
||||
|
||||
t1.join().unwrap();
|
||||
t2.join().unwrap();
|
||||
|
||||
assert!(cache.len() <= 10);
|
||||
}
|
||||
}
|
||||
Vendored
+97
@@ -0,0 +1,97 @@
|
||||
use crate::cache::lru::LruTtlCache;
|
||||
use crate::cache::redis::AppRedis;
|
||||
use crate::config::AppConfig;
|
||||
use crate::error::AppResult;
|
||||
use ::redis::Cmd;
|
||||
use serde::Serialize;
|
||||
use serde::de::DeserializeOwned;
|
||||
use std::time::Duration;
|
||||
|
||||
pub mod lru;
|
||||
pub mod redis;
|
||||
|
||||
pub struct AppCache {
|
||||
l1: LruTtlCache<String, String>,
|
||||
l2: AppRedis,
|
||||
key_prefix: String,
|
||||
default_ttl: Duration,
|
||||
}
|
||||
|
||||
impl AppCache {
|
||||
pub fn from_config(config: &AppConfig) -> AppResult<Self> {
|
||||
let cap = config.lru_default_capacity()?;
|
||||
let ttl = Duration::from_secs(config.lru_default_ttl_secs()?);
|
||||
let l2 = AppRedis::from_config(config)?;
|
||||
let key_prefix = config.redis_key_prefix()?;
|
||||
Ok(Self {
|
||||
l1: LruTtlCache::new(cap, ttl),
|
||||
l2,
|
||||
key_prefix,
|
||||
default_ttl: ttl,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn get<T: DeserializeOwned>(&self, key: &str) -> Option<T> {
|
||||
if let Some(json) = self.l1.get(&key.to_string()) {
|
||||
return serde_json::from_str(&json).ok();
|
||||
}
|
||||
|
||||
let full_key = self.full_key(key);
|
||||
let mut conn = self.l2.get_connection().ok()?;
|
||||
let json: String = Cmd::new()
|
||||
.arg("GET")
|
||||
.arg(&full_key)
|
||||
.query::<Option<String>>(&mut *conn.inner_mut())
|
||||
.ok()??;
|
||||
|
||||
let value: T = serde_json::from_str(&json).ok()?;
|
||||
self.l1.insert(key.to_string(), json);
|
||||
Some(value)
|
||||
}
|
||||
|
||||
pub fn set<T: Serialize>(&self, key: &str, value: &T, ttl: Option<Duration>) -> AppResult<()> {
|
||||
let json = serde_json::to_string(value)?;
|
||||
let full_key = self.full_key(key);
|
||||
let ttl_duration = ttl.unwrap_or(self.default_ttl);
|
||||
let ttl_secs = ttl_duration.as_secs() as usize;
|
||||
let mut conn = self.l2.get_connection()?;
|
||||
Cmd::new()
|
||||
.arg("SETEX")
|
||||
.arg(&full_key)
|
||||
.arg(ttl_secs)
|
||||
.arg(&json)
|
||||
.query::<()>(&mut *conn.inner_mut())?;
|
||||
self.l1.insert_with_ttl(key.to_string(), json, ttl_duration);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn delete(&self, key: &str) -> AppResult<()> {
|
||||
self.l1.remove(&key.to_string());
|
||||
let full_key = self.full_key(key);
|
||||
let mut conn = self.l2.get_connection()?;
|
||||
Cmd::new()
|
||||
.arg("DEL")
|
||||
.arg(&full_key)
|
||||
.query::<()>(&mut *conn.inner_mut())?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn exists(&self, key: &str) -> bool {
|
||||
if self.l1.get(&key.to_string()).is_some() {
|
||||
return true;
|
||||
}
|
||||
let full_key = self.full_key(key);
|
||||
if let Ok(mut conn) = self.l2.get_connection() {
|
||||
return Cmd::new()
|
||||
.arg("EXISTS")
|
||||
.arg(&full_key)
|
||||
.query(&mut *conn.inner_mut())
|
||||
.unwrap_or(false);
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
fn full_key(&self, key: &str) -> String {
|
||||
format!("{}{}", self.key_prefix, key)
|
||||
}
|
||||
}
|
||||
Vendored
+117
@@ -0,0 +1,117 @@
|
||||
use crate::config::AppConfig;
|
||||
use crate::error::AppError;
|
||||
use crate::error::AppResult;
|
||||
use r2d2::Pool;
|
||||
use redis::cluster::ClusterClient;
|
||||
use redis::{Client, ConnectionLike, RedisError};
|
||||
use std::time::Duration;
|
||||
|
||||
#[derive(Clone)]
|
||||
enum RedisBackend {
|
||||
Single(Pool<Client>),
|
||||
Cluster(Pool<ClusterClient>),
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AppRedis {
|
||||
backend: RedisBackend,
|
||||
}
|
||||
|
||||
impl AppRedis {
|
||||
pub fn from_config(config: &AppConfig) -> AppResult<Self> {
|
||||
let backend = if config.redis_cluster_enabled()? {
|
||||
let nodes = config.redis_cluster_nodes()?;
|
||||
let cluster_client =
|
||||
ClusterClient::new(nodes.iter().map(|s| s.as_str()).collect::<Vec<_>>())?;
|
||||
let pool = Self::build_pool(config, cluster_client)?;
|
||||
RedisBackend::Cluster(pool)
|
||||
} else {
|
||||
let url = config
|
||||
.redis_url()?
|
||||
.ok_or_else(|| AppError::Config("APP_REDIS_URL is not set".into()))?;
|
||||
let client = Client::open(url.as_str())?;
|
||||
let pool = Self::build_pool(config, client)?;
|
||||
RedisBackend::Single(pool)
|
||||
};
|
||||
|
||||
Ok(Self { backend })
|
||||
}
|
||||
|
||||
fn build_pool<M: r2d2::ManageConnection>(config: &AppConfig, manager: M) -> AppResult<Pool<M>> {
|
||||
let max_conn = config.redis_max_connections()?;
|
||||
let min_conn = config.redis_min_connections()?;
|
||||
let idle_timeout = config.redis_idle_timeout()?;
|
||||
let conn_timeout = config.redis_connection_timeout()?;
|
||||
|
||||
Ok(r2d2::Builder::new()
|
||||
.max_size(max_conn)
|
||||
.min_idle(Some(min_conn))
|
||||
.idle_timeout(Some(Duration::from_secs(idle_timeout)))
|
||||
.connection_timeout(Duration::from_secs(conn_timeout))
|
||||
.build(manager)?)
|
||||
}
|
||||
|
||||
pub fn get_connection(&self) -> Result<PooledRedisConnection, r2d2::Error> {
|
||||
match &self.backend {
|
||||
RedisBackend::Single(pool) => pool.get().map(PooledRedisConnection::Single),
|
||||
RedisBackend::Cluster(pool) => pool.get().map(PooledRedisConnection::Cluster),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::large_enum_variant)]
|
||||
pub enum PooledRedisConnection {
|
||||
Single(r2d2::PooledConnection<Client>),
|
||||
Cluster(r2d2::PooledConnection<ClusterClient>),
|
||||
}
|
||||
|
||||
impl PooledRedisConnection {
|
||||
pub fn inner_mut(&mut self) -> &mut dyn ConnectionLike {
|
||||
match self {
|
||||
PooledRedisConnection::Single(conn) => conn,
|
||||
PooledRedisConnection::Cluster(conn) => conn,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ConnectionLike for PooledRedisConnection {
|
||||
fn req_packed_command(&mut self, cmd: &[u8]) -> Result<redis::Value, RedisError> {
|
||||
match self {
|
||||
PooledRedisConnection::Single(conn) => conn.req_packed_command(cmd),
|
||||
PooledRedisConnection::Cluster(conn) => conn.req_packed_command(cmd),
|
||||
}
|
||||
}
|
||||
|
||||
fn req_packed_commands(
|
||||
&mut self,
|
||||
cmd: &[u8],
|
||||
offset: usize,
|
||||
count: usize,
|
||||
) -> Result<Vec<redis::Value>, RedisError> {
|
||||
match self {
|
||||
PooledRedisConnection::Single(conn) => conn.req_packed_commands(cmd, offset, count),
|
||||
PooledRedisConnection::Cluster(conn) => conn.req_packed_commands(cmd, offset, count),
|
||||
}
|
||||
}
|
||||
|
||||
fn get_db(&self) -> i64 {
|
||||
match self {
|
||||
PooledRedisConnection::Single(conn) => conn.get_db(),
|
||||
PooledRedisConnection::Cluster(conn) => conn.get_db(),
|
||||
}
|
||||
}
|
||||
|
||||
fn check_connection(&mut self) -> bool {
|
||||
match self {
|
||||
PooledRedisConnection::Single(conn) => conn.check_connection(),
|
||||
PooledRedisConnection::Cluster(conn) => conn.check_connection(),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_open(&self) -> bool {
|
||||
match self {
|
||||
PooledRedisConnection::Single(conn) => conn.is_open(),
|
||||
PooledRedisConnection::Cluster(conn) => conn.is_open(),
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user