use redis::aio::ConnectionManager;
use serde::{Deserialize, Serialize};
use crate::config::BotConfig;
use crate::error::Result;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PageEntry {
pub index: usize,
pub item: serde_json::Value,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PageList {
pub session_id: String,
pub title: String,
pub entries: Vec<PageEntry>,
pub created_at: String,
}
impl PageList {
pub fn empty() -> Self {
Self {
session_id: String::new(),
title: String::new(),
entries: Vec::new(),
created_at: String::new(),
}
}
}
#[derive(Clone)]
pub struct PageCache {
backend: PageCacheBackend,
ttl_secs: u64,
}
#[derive(Clone)]
enum PageCacheBackend {
Redis(ConnectionManager),
Offline,
}
impl PageCache {
pub fn offline() -> Self {
Self { backend: PageCacheBackend::Offline, ttl_secs: 300 }
}
pub async fn connect(config: &BotConfig) -> Result<Self> {
let client = redis::Client::open(config.redis_url.as_str())?;
let conn = ConnectionManager::new(client).await?;
Ok(Self {
backend: PageCacheBackend::Redis(conn),
ttl_secs: config.link_code_ttl_secs.max(300),
})
}
pub async fn store(&self, title: &str, entries: Vec<PageEntry>) -> Result<String> {
let session_id = uuid::Uuid::new_v4().to_string();
let PageCacheBackend::Redis(conn) = &self.backend else {
return Ok(session_id);
};
let list = PageList {
session_id: session_id.clone(),
title: title.to_string(),
entries,
created_at: chrono::Utc::now().to_rfc3339(),
};
let key = format!("archivist:session:{session_id}");
let data = serde_json::to_vec(&list)?;
redis::cmd("SETEX")
.arg(&key)
.arg(self.ttl_secs)
.arg(data)
.exec_async(&mut conn.clone())
.await?;
Ok(session_id)
}
pub async fn store_for_user(
&self,
user_id: u64,
title: &str,
entries: Vec<PageEntry>,
) -> Result<String> {
let session_id = self.store(title, entries).await?;
let PageCacheBackend::Redis(conn) = &self.backend else {
return Ok(session_id);
};
let key = format!("archivist:user:{user_id}:latest");
redis::cmd("SETEX")
.arg(&key)
.arg(self.ttl_secs)
.arg(&session_id)
.exec_async(&mut conn.clone())
.await?;
Ok(session_id)
}
pub async fn latest_for_user(&self, user_id: u64) -> Result<Option<String>> {
let PageCacheBackend::Redis(conn) = &self.backend else {
return Ok(None);
};
let key = format!("archivist:user:{user_id}:latest");
let val: Option<String> = redis::cmd("GET")
.arg(&key)
.query_async(&mut conn.clone())
.await?;
Ok(val)
}
pub async fn current_page(&self, session_id: &str) -> Result<Option<usize>> {
let PageCacheBackend::Redis(conn) = &self.backend else {
return Ok(None);
};
let key = format!("archivist:page:{session_id}");
let val: Option<u64> = redis::cmd("GET")
.arg(&key)
.query_async(&mut conn.clone())
.await?;
Ok(val.map(|v| v as usize))
}
pub async fn set_current_page(&self, session_id: &str, page: usize) -> Result<()> {
let PageCacheBackend::Redis(conn) = &self.backend else {
return Ok(());
};
let key = format!("archivist:page:{session_id}");
redis::cmd("SETEX")
.arg(&key)
.arg(self.ttl_secs)
.arg(page as u64)
.exec_async(&mut conn.clone())
.await?;
Ok(())
}
pub async fn fetch(&self, session_id: &str) -> Result<Option<PageList>> {
let PageCacheBackend::Redis(conn) = &self.backend else {
return Ok(None);
};
let key = format!("archivist:session:{session_id}");
let raw: Option<Vec<u8>> = redis::cmd("GET")
.arg(&key)
.query_async(&mut conn.clone())
.await?;
match raw {
Some(bytes) => {
let list: PageList = serde_json::from_slice(&bytes)?;
Ok(Some(list))
}
None => Ok(None),
}
}
pub async fn delete(&self, session_id: &str) -> Result<()> {
let PageCacheBackend::Redis(conn) = &self.backend else {
return Ok(());
};
let key = format!("archivist:session:{session_id}");
redis::cmd("DEL")
.arg(&key)
.exec_async(&mut conn.clone())
.await?;
Ok(())
}
pub async fn log_search(&self, entry: &SearchLogEntry) {
let PageCacheBackend::Redis(conn) = &self.backend else {
return;
};
let key = format!(
"archivist:searchlog:{}",
chrono::Utc::now().format("%Y-%m-%d")
);
let data = match serde_json::to_vec(entry) {
Ok(d) => d,
Err(e) => {
tracing::warn!("search log serialization failed: {e}");
return;
}
};
let mut conn = conn.clone();
let res: redis::RedisResult<i64> = redis::cmd("RPUSH")
.arg(&key)
.arg(&data)
.query_async(&mut conn)
.await;
match res {
Ok(_) => {
let _ = redis::cmd("LTRIM")
.arg(&key)
.arg(-(SEARCH_LOG_CAP as isize))
.arg(-1)
.exec_async(&mut conn.clone())
.await;
let _ = redis::cmd("EXPIRE")
.arg(&key)
.arg(7 * 24 * 60 * 60)
.exec_async(&mut conn.clone())
.await;
}
Err(e) => tracing::warn!("search log append failed: {e}"),
}
}
pub async fn cached_response(&self, kind: &str, query: &str) -> Option<serde_json::Value> {
let PageCacheBackend::Redis(conn) = &self.backend else {
return None;
};
let key = response_cache_key(kind, query);
let res: redis::RedisResult<Option<Vec<u8>>> = redis::cmd("GET")
.arg(&key)
.query_async(&mut conn.clone())
.await;
match res {
Ok(Some(bytes)) => match serde_json::from_slice(&bytes) {
Ok(v) => Some(v),
Err(e) => {
tracing::warn!("response cache deserialize failed for {key}: {e}");
None
}
},
Ok(None) => None,
Err(e) => {
tracing::warn!("response cache get failed for {key}: {e}");
None
}
}
}
pub async fn cache_response(&self, kind: &str, query: &str, value: &serde_json::Value) {
let PageCacheBackend::Redis(conn) = &self.backend else {
return;
};
let key = response_cache_key(kind, query);
let data = match serde_json::to_vec(value) {
Ok(d) => d,
Err(e) => {
tracing::warn!("response cache serialization failed for {key}: {e}");
return;
}
};
let ttl = match kind {
"ask" => ASK_CACHE_TTL_SECS,
_ => SEARCH_CACHE_TTL_SECS,
};
let res: redis::RedisResult<()> = redis::cmd("SETEX")
.arg(&key)
.arg(ttl)
.arg(data)
.exec_async(&mut conn.clone())
.await;
if let Err(e) = res {
tracing::warn!("response cache set failed for {key}: {e}");
}
}
pub async fn cached_raw(&self, key: &str) -> Option<serde_json::Value> {
let PageCacheBackend::Redis(conn) = &self.backend else {
return None;
};
let res: redis::RedisResult<Option<Vec<u8>>> = redis::cmd("GET")
.arg(key)
.query_async(&mut conn.clone())
.await;
match res {
Ok(Some(bytes)) => match serde_json::from_slice(&bytes) {
Ok(v) => Some(v),
Err(e) => {
tracing::warn!("raw cache deserialize failed for {key}: {e}");
None
}
},
Ok(None) => None,
Err(e) => {
tracing::warn!("raw cache get failed for {key}: {e}");
None
}
}
}
pub async fn cache_raw(&self, key: &str, ttl_secs: u64, value: &serde_json::Value) {
let PageCacheBackend::Redis(conn) = &self.backend else {
return;
};
let data = match serde_json::to_vec(value) {
Ok(d) => d,
Err(e) => {
tracing::warn!("raw cache serialization failed for {key}: {e}");
return;
}
};
let res: redis::RedisResult<()> = redis::cmd("SETEX")
.arg(key)
.arg(ttl_secs)
.arg(data)
.exec_async(&mut conn.clone())
.await;
if let Err(e) = res {
tracing::warn!("raw cache set failed for {key}: {e}");
}
}
}
pub const ASK_CACHE_TTL_SECS: u64 = 24 * 60 * 60;
pub const SEARCH_CACHE_TTL_SECS: u64 = 5 * 60;
fn response_cache_key(kind: &str, query: &str) -> String {
use sha2::{Digest, Sha256};
let norm = query.trim().to_lowercase();
let mut hasher = Sha256::new();
hasher.update(norm.as_bytes());
let hex = hasher.finalize();
let h = hex.iter().map(|b| format!("{b:02x}")).collect::<String>();
format!("archivist:{kind}cache:{h}")
}
pub const SEARCH_LOG_CAP: usize = 500;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchLogEntry {
pub ts: String,
pub user_id: u64,
pub command: String,
pub query: String,
pub api_query: String,
pub status: u16,
pub result_count: usize,
pub latency_ms: u64,
pub error: String,
}
pub fn slice_page<T>(items: &[T], page: usize, page_size: usize) -> Vec<&T> {
let page = page.max(1);
let start = (page - 1) * page_size;
if start >= items.len() {
return Vec::new();
}
let end = (start + page_size).min(items.len());
items[start..end].iter().collect()
}
pub fn page_count(len: usize, page_size: usize) -> usize {
if len == 0 {
return 0;
}
(len + page_size - 1) / page_size
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn slice_page_first() {
let items = vec![1, 2, 3, 4, 5, 6, 7];
let page = slice_page(&items, 1, 3);
assert_eq!(page, vec![&1, &2, &3]);
}
#[test]
fn slice_page_middle() {
let items = vec![1, 2, 3, 4, 5, 6, 7];
let page = slice_page(&items, 2, 3);
assert_eq!(page, vec![&4, &5, &6]);
}
#[test]
fn slice_page_last_partial() {
let items = vec![1, 2, 3, 4, 5, 6, 7];
let page = slice_page(&items, 3, 3);
assert_eq!(page, vec![&7]);
}
#[test]
fn slice_page_out_of_range() {
let items = vec![1, 2, 3, 4, 5, 6, 7];
let page = slice_page(&items, 99, 3);
assert!(page.is_empty());
}
#[test]
fn page_count_calc() {
assert_eq!(page_count(0, 3), 0);
assert_eq!(page_count(7, 3), 3);
assert_eq!(page_count(9, 3), 3);
assert_eq!(page_count(10, 3), 4);
assert_eq!(page_count(5, 5), 1);
}
#[test]
fn response_cache_key_prefixes_by_kind() {
let ask = response_cache_key("ask", "dark harry");
let search = response_cache_key("search", "dark harry");
assert!(ask.starts_with("archivist:askcache:"));
assert!(search.starts_with("archivist:searchcache:"));
assert_ne!(ask, search);
assert_eq!(ask.len(), "archivist:askcache:".len() + 64);
}
#[test]
fn response_cache_key_normalizes_query() {
let a = response_cache_key("ask", " Dark Harry ");
let b = response_cache_key("ask", "dark harry");
let c = response_cache_key("ask", "DARK HARRY");
assert_eq!(a, b);
assert_eq!(b, c);
}
#[test]
fn response_cache_key_differs_for_diff_queries() {
let a = response_cache_key("ask", "drarry");
let b = response_cache_key("ask", "enemies to lovers");
assert_ne!(a, b);
}
}