Restructured cache. Added Redis|Valkey|Dragonfly support.

src/cache/mod.rs

@@ -0,0 +1,2 @@
+pub mod redis;
+pub mod types;

src/cache/redis.rs

@@ -0,0 +1,222 @@
+use async_trait::async_trait;
+use bytes::Bytes;
+use log::warn;
+use pingora::Result;
+use pingora_cache::storage::{HandleHit, HandleMiss, HitHandler, MissFinishType, MissHandler, Storage};
+use pingora_cache::trace::SpanHandle;
+use pingora_cache::{CacheKey, CacheMeta, PurgeOutcome, PurgeTarget, PurgeType};
+use redis::aio::ConnectionManager;
+use redis::{AsyncCommands, Client, RedisResult};
+use std::any::Any;
+use std::borrow::Cow;
+use std::sync::Arc;
+
+#[derive(Clone)]
+pub struct RedisCache {
+    pub client: Arc<ConnectionManager>,
+    pub ttl: u64,
+}
+
+pub struct RedisHitHandler {
+    body: Option<Bytes>,
+}
+
+pub struct RedisMissHandler {
+    client: Arc<ConnectionManager>,
+    body_key: String,
+    body_buffer: Vec<u8>,
+    ttl: u64,
+}
+
+impl RedisCache {
+    pub async fn new(url: &str, ttl: u64) -> RedisResult<Self> {
+        let client = Client::open(url)?;
+        let manager = ConnectionManager::new(client).await?;
+        Ok(Self { client: Arc::new(manager), ttl })
+    }
+    fn make_key(key: &CacheKey) -> Cow<'_, str> {
+        match key.primary_key_str() {
+            Some(s) => Cow::Borrowed(s),
+            None => Cow::Owned(key.primary_key().iter().map(|b| format!("{:02x}", b)).collect()),
+        }
+    }
+}
+
+#[async_trait]
+impl HandleHit for RedisHitHandler {
+    async fn read_body(&mut self) -> Result<Option<Bytes>> {
+        Ok(self.body.take())
+    }
+
+    async fn finish(self: Box<Self>, _storage: &'static (dyn Storage + Sync), _key: &CacheKey, _trace: &SpanHandle) -> Result<()> {
+        Ok(())
+    }
+
+    fn can_seek(&self) -> bool {
+        false
+    }
+
+    fn seek(&mut self, _start: usize, _end: Option<usize>) -> Result<()> {
+        Ok(())
+    }
+
+    fn should_count_access(&self) -> bool {
+        true
+    }
+
+    fn get_eviction_weight(&self) -> usize {
+        self.body.as_ref().map(|b| b.len()).unwrap_or(0)
+    }
+
+    fn entry_id(&self) -> Option<pingora_cache::eviction::CacheEntryId> {
+        None
+    }
+
+    fn as_any(&self) -> &(dyn Any + Send + Sync) {
+        self
+    }
+
+    fn as_any_mut(&mut self) -> &mut (dyn Any + Send + Sync) {
+        self
+    }
+}
+
+#[async_trait]
+impl HandleMiss for RedisMissHandler {
+    async fn write_body(&mut self, data: Bytes, _eof: bool) -> Result<()> {
+        self.body_buffer.extend_from_slice(&data);
+        Ok(())
+    }
+
+    async fn finish(self: Box<Self>) -> Result<MissFinishType> {
+        let len = self.body_buffer.len();
+        let mut conn = (*self.client).clone();
+        if let Err(e) = conn.set_ex::<_, _, ()>(&self.body_key, &self.body_buffer, self.ttl).await {
+            warn!("Failed to write cache entry to Redis: {e}");
+        }
+        Ok(MissFinishType::Appended(len, None))
+    }
+
+    fn entry_id(&self) -> Option<pingora_cache::eviction::CacheEntryId> {
+        None
+    }
+}
+
+#[async_trait]
+impl Storage for RedisCache {
+    async fn lookup(&'static self, key: &CacheKey, _trace: &SpanHandle) -> Result<Option<(CacheMeta, HitHandler)>> {
+        let base = Self::make_key(key);
+
+        // Pre-allocate exact capacities to prevent reallocations
+        let mut meta_key = String::with_capacity(base.len() + 5);
+        meta_key.push_str(&base);
+        meta_key.push_str(":meta");
+
+        let mut body_key = String::with_capacity(base.len() + 5);
+        body_key.push_str(&base);
+        body_key.push_str(":body");
+
+        let mut internal_key = String::with_capacity(meta_key.len() + 9);
+        internal_key.push_str(&meta_key);
+        internal_key.push_str(":internal");
+
+        let mut header_key = String::with_capacity(meta_key.len() + 7);
+        header_key.push_str(&meta_key);
+        header_key.push_str(":header");
+
+        let mut conn = (*self.client).clone();
+
+        let internal: Option<Vec<u8>> = conn.get(&internal_key).await.ok().flatten();
+        let header: Option<Vec<u8>> = conn.get(&header_key).await.ok().flatten();
+        let body: Option<Vec<u8>> = conn.get(&body_key).await.ok().flatten();
+
+        match (internal, header, body) {
+            (Some(internal_bytes), Some(header_bytes), Some(body_bytes)) => match CacheMeta::deserialize(&internal_bytes, &header_bytes) {
+                Ok(meta) => {
+                    let handler: HitHandler = Box::new(RedisHitHandler {
+                        body: Some(Bytes::from(body_bytes)),
+                    });
+                    Ok(Some((meta, handler)))
+                }
+                Err(_) => {
+                    let _: RedisResult<()> = conn.del(&[&internal_key, &header_key, &body_key]).await;
+                    Ok(None)
+                }
+            },
+            _ => Ok(None),
+        }
+    }
+
+    async fn get_miss_handler(&'static self, key: &CacheKey, meta: &CacheMeta, _trace: &SpanHandle) -> Result<MissHandler> {
+        let base = Self::make_key(key);
+
+        let mut meta_key = String::with_capacity(base.len() + 5);
+        meta_key.push_str(&base);
+        meta_key.push_str(":meta");
+
+        let mut body_key = String::with_capacity(base.len() + 5);
+        body_key.push_str(&base);
+        body_key.push_str(":body");
+
+        let mut internal_key = String::with_capacity(meta_key.len() + 9);
+        internal_key.push_str(&meta_key);
+        internal_key.push_str(":internal");
+
+        let mut header_key = String::with_capacity(meta_key.len() + 7);
+        header_key.push_str(&meta_key);
+        header_key.push_str(":header");
+
+        let ttl = self.ttl;
+        let (internal_bytes, header_bytes) = meta.serialize()?;
+        let mut conn = (*self.client).clone();
+
+        let _: RedisResult<()> = conn.set_ex(&internal_key, &internal_bytes, ttl).await;
+        let _: RedisResult<()> = conn.set_ex(&header_key, &header_bytes, ttl).await;
+
+        Ok(Box::new(RedisMissHandler {
+            client: self.client.clone(),
+            body_key,
+            body_buffer: Vec::with_capacity(4096),
+            ttl,
+        }))
+    }
+
+    async fn purge(&'static self, _target: PurgeTarget<'_>, _purge_type: PurgeType, _trace: &SpanHandle) -> Result<PurgeOutcome> {
+        Ok(PurgeOutcome::NotFound)
+    }
+
+    async fn update_meta(&'static self, key: &CacheKey, meta: &CacheMeta, _trace: &SpanHandle) -> Result<bool> {
+        let base = Self::make_key(key);
+
+        let mut meta_key = String::with_capacity(base.len() + 5);
+        meta_key.push_str(&base);
+        meta_key.push_str(":meta");
+
+        let mut internal_key = String::with_capacity(meta_key.len() + 9);
+        internal_key.push_str(&meta_key);
+        internal_key.push_str(":internal");
+
+        let mut header_key = String::with_capacity(meta_key.len() + 7);
+        header_key.push_str(&meta_key);
+        header_key.push_str(":header");
+
+        let ttl = self.ttl;
+        match meta.serialize() {
+            Ok((internal_bytes, header_bytes)) => {
+                let mut conn = (*self.client).clone();
+                let _: RedisResult<()> = conn.set_ex(&internal_key, &internal_bytes, ttl).await;
+                let _: RedisResult<()> = conn.set_ex(&header_key, &header_bytes, ttl).await;
+                Ok(true)
+            }
+            Err(_) => Ok(false),
+        }
+    }
+
+    fn support_streaming_partial_write(&self) -> bool {
+        false
+    }
+
+    fn as_any(&self) -> &(dyn Any + Send + Sync) {
+        self
+    }
+}

src/cache/types.rs

@@ -0,0 +1,41 @@
+use crate::cache::redis::RedisCache;
+use crate::utils::lazylock::{CACHE_TTL, EVICTION};
+use crate::utils::types::AppConfig;
+use log::{info, warn};
+use pingora_cache::eviction::simple_lru::Manager;
+use pingora_cache::{MemCache, Storage};
+use std::sync::{Arc, OnceLock};
+
+pub static GLOBAL_STORAGE: OnceLock<Box<dyn Storage + Send + Sync>> = OnceLock::new();
+
+pub async fn init_cache(config: Arc<AppConfig>) {
+    let cache_storage = config.cache_storage.clone().unwrap_or_else(|| {
+        info!("No cache_storage configured, defaulting to memory");
+        "memory".to_string()
+    });
+    let cache_type = cache_storage.splitn(2, ':').next().unwrap_or("memory");
+    let ttl = config.cache_ttl.unwrap_or(60);
+
+    match cache_type {
+        "memory" => {
+            CACHE_TTL.set(ttl).ok();
+            if let Some(size) = config.cache_size_mb {
+                let cache_size = size * 1024 * 1024;
+                drop(EVICTION.set(Manager::new(cache_size)));
+            }
+            GLOBAL_STORAGE.set(Box::new(MemCache::new())).ok();
+            info!("Cache initialized successfully. Storage: {}", cache_storage);
+        }
+        "redis" => match RedisCache::new(&cache_storage, ttl).await {
+            Ok(cache) => {
+                CACHE_TTL.set(ttl).ok();
+                GLOBAL_STORAGE.set(Box::new(cache)).ok();
+                info!("Cache initialized successfully. Storage: {}", cache_storage);
+            }
+            Err(e) => log::error!("Cache storage initialization failed: {}", e),
+        },
+        _ => {
+            warn!("Unsupported storage, disabling caching: {}", cache_storage);
+        }
+    }
+}

src/core/bgservice.rs

@@ -1,3 +1,4 @@
+use crate::cache::types::init_cache;
 use crate::core::proxyhttp::LB;
 use crate::ingress::{APIUpstreamProvider, ConsulProvider, Discovery, FromFileProvider, KubernetesProvider};
 use crate::logging::core::{init_access_log, init_access_logging, log_builder};
@@ -84,6 +85,7 @@ impl BackgroundService for LB {
         let acme_provider = self.config.acme_dns_provider.clone();
         tokio::spawn(async move { refresh_order(certdir, confdir, acme_provider).await });
 
+        init_cache(self.config.clone()).await;
         init_access_logging(self.config.access_log.clone()).await;
         loop {
             tokio::select! {

src/core/proxyhttp.rs

@@ -1,7 +1,8 @@
 use crate::auth::authenticate;
+use crate::cache::types::GLOBAL_STORAGE;
 use crate::core::gethosts::{GetHost, GetHostsReturHeaders};
 use crate::logging::core::access_log;
-use crate::utils::lazylock::{CACHE_LOCK, CACHE_TTL, EVICTION, LOCALHOST, MEM_CACHE, RATE_LIMITER, REQUESTS_4XX, REVERSE_STORE};
+use crate::utils::lazylock::{CACHE_LOCK, CACHE_TTL, EVICTION, LOCALHOST, RATE_LIMITER, REQUESTS_4XX, REVERSE_STORE};
 use crate::utils::metrics::*;
 use crate::utils::types::{AppConfig, Extraparams, Headers, InnerMap, UpstreamsDashMap, UpstreamsIdMap};
 use arc_swap::ArcSwap;
@@ -67,8 +68,7 @@ impl ProxyHttp for LB {
 
     fn cache_key_callback(&self, session: &Session, ctx: &mut Self::CTX) -> Result<CacheKey> {
         let host = ctx.hostname.as_deref().unwrap_or("default");
-        let path = session.req_header().uri.path();
-
+        let path = session.req_header().uri.path_and_query().map_or("/", |pq| pq.as_str());
         let primary_key = KEY_BUF.with(|buf| {
             let mut b = buf.borrow_mut();
             b.clear();
@@ -76,15 +76,14 @@ impl ProxyHttp for LB {
             b.push_str(path);
             b.clone()
         });
-
-        let secondary_key = session.req_header().uri.path_and_query().map(|pq| pq.as_str()).unwrap_or("/");
-        Ok(CacheKey::new(primary_key, secondary_key))
+        Ok(CacheKey::new(primary_key, host /* host is just for debugging */))
     }
-
     fn request_cache_filter(&self, session: &mut Session, _ctx: &mut Self::CTX) -> Result<()> {
-        if self.cache_enabled && session.req_header().method == Method::GET {
-            if let Some(eviction) = EVICTION.get() {
-                session.cache.enable(&*MEM_CACHE, Some(eviction), None, Some(&*CACHE_LOCK), None);
+        let has_auth = _ctx.upstream_peer.as_ref().map(|m| m.authorization.is_some()).unwrap_or(false) || _ctx.extraparams.authentication.is_some();
+        if self.cache_enabled && session.req_header().method == Method::GET && !has_auth {
+            if let Some(storage) = GLOBAL_STORAGE.get() {
+                let eviction_mgr = EVICTION.get().map(|e| e as &(dyn pingora_cache::eviction::EvictionManager + Sync));
+                session.cache.enable(storage.as_ref(), eviction_mgr, None, Some(&*CACHE_LOCK), None);
             }
         }
         Ok(())

src/core/start.rs

@@ -29,20 +29,24 @@ use std::time::Duration;
 use std::{fs, thread};
 
 pub fn run() {
-    default_provider().install_default().expect("Failed to install RusTLS crypto provider");
+    default_provider().install_default().die("Failed to install RusTLS crypto provider:");
     let parameters = Opt::parse_args();
-    let file = parameters.conf.clone().unwrap();
+    let file = parameters.conf.clone().die("Failed to clone the config:");
     let maincfg = crate::utils::parceyaml::parce_main_config(file.as_str());
 
     let master_server = maincfg.cluster_member.as_deref() == Some("master");
-    let cache_enabled = maincfg.cache_size_mb.is_some();
-    CACHE_TTL.set(maincfg.cache_ttl.unwrap_or(60)).unwrap();
-    if let Some(size) = maincfg.cache_size_mb {
-        let cache_size = size * 1024 * 1024;
-        drop(EVICTION.set(Manager::new(cache_size)));
-    }
 
-    let mut server = Server::new(parameters).unwrap();
+    let cache_enabled = maincfg.cache_enabled.unwrap_or(false);
+
+    // ========================================================================== //
+    CACHE_TTL.set(maincfg.cache_ttl.unwrap_or(60)).die("Failed to set cache TTL");
+    if cache_enabled {
+        let cache_size = maincfg.cache_size_mb.unwrap_or(50);
+        let _ = EVICTION.set(Manager::new(cache_size * 1024 * 1024));
+    }
+    // ========================================================================== //
+
+    let mut server = Server::new(parameters).die("Failed to start the server:");
     server.bootstrap();
 
     let uf_config = Arc::new(DashMap::new());
@@ -106,22 +110,22 @@ pub fn run() {
     if let Some(bind_address_tls) = bind_address_tls {
         check_priv(bind_address_tls.as_str());
         let (tx, rx): (Sender<Vec<CertificateConfig>>, Receiver<Vec<CertificateConfig>>) = channel();
-        let certs_path = cfg.proxy_configs.clone().unwrap() + "/certificates";
+        let certs_path = cfg.proxy_configs.clone().die("Failed certificates path") + "/certificates";
 
         if fs::metadata(certs_path.clone()).is_err() {
-            fs::create_dir_all(certs_path.clone()).unwrap();
+            fs::create_dir_all(certs_path.clone()).die("Failed to create directories");
         }
         thread::spawn(move || {
-            folder_watch(certs_path, tx).unwrap();
+            folder_watch(certs_path, tx).die("Failed to spawn certificates channel");
         });
-        let certificate_configs = rx.recv().unwrap();
+        let certificate_configs = rx.recv().die("Failed to spawn certificates channel");
         let first_set = load::Certificates::new(&certificate_configs, grade.as_str()).unwrap_or_else(|| panic!("Unable to load initial certificate info"));
         let certificates = Arc::new(ArcSwap::from_pointee(first_set));
         let certs_for_callback = certificates.clone();
 
         let certs_for_watcher = certificates.clone();
         let new_certs = load::Certificates::new(&certificate_configs, grade.as_str());
-        certs_for_watcher.store(Arc::new(new_certs.unwrap()));
+        certs_for_watcher.store(Arc::new(new_certs.die("Failed to build the watcher")));
 
         let mut tls_settings =
             TlsSettings::intermediate(&certs_for_callback.load().default_cert_path, &certs_for_callback.load().default_key_path).expect("unable to load or parse cert/key");
@@ -162,7 +166,7 @@ pub fn run() {
     if let Err(e) = write_pid_file(pf.as_str()) {
         panic!("Failed to write PID file: {} : {}", pf, e);
     }
-    let mut signals = Signals::new([SIGINT, SIGTERM, SIGQUIT]).unwrap();
+    let mut signals = Signals::new([SIGINT, SIGTERM, SIGQUIT]).die("Failed to create signals");
     if let Some(sig) = signals.forever().next() {
         match sig {
             SIGINT => info!("SIGINT received! Exiting..."),

src/lib.rs

@@ -1,4 +1,5 @@
 pub mod auth;
+pub mod cache;
 pub mod core;
 pub mod ingress;
 pub mod logging;

src/utils/parceyaml.rs

@@ -2,7 +2,7 @@ use crate::utils::deser::{send_to_worker, serialize_config};
 use crate::utils::healthcheck;
 use crate::utils::lazylock::{DOMAINS, GLOBAL_PARAMS, IM_A_MASTER, MASTER_KEY, NODE_STATUS, REVERSE_STORE, WORKERS};
 use crate::utils::state::{is_first_run, mark_not_first_run};
-use crate::utils::tools::clone_dashmap_into;
+use crate::utils::tools::{Die, clone_dashmap_into};
 use crate::utils::types::*;
 use dashmap::DashMap;
 use log::{error, info, warn};
@@ -245,7 +245,7 @@ async fn populate_file_upstreams(config: &mut Configuration, parsed: &Config) {
         }
 
         {
-            let mut config_guard = GLOBAL_PARAMS.write().expect("Global config lock poisoned");
+            let mut config_guard = GLOBAL_PARAMS.write().die("Global config lock poisoned");
             config_guard.client_headers = config.client_headers.clone();
             config_guard.server_headers = config.server_headers.clone();
             config_guard.extraparams = config.extraparams.clone();
@@ -274,8 +274,8 @@ async fn populate_file_upstreams(config: &mut Configuration, parsed: &Config) {
     }
 }
 pub fn parce_main_config(path: &str) -> AppConfig {
-    let data = fs::read_to_string(path).expect("Failed to read main config file");
-    let mut cfo: AppConfig = noyalib::from_str(&data).expect("Failed to parse main config file");
+    let data = fs::read_to_string(path).die("Failed to read main config file");
+    let mut cfo: AppConfig = noyalib::from_str(&data).die("Failed to parse main config file");
     if let Ok(jwt_key) = env::var("MASTER_KEY") {
         cfo.master_key = Some(jwt_key);
     } else if let Ok(jwt_key) = env::var("JWT_KEY") {
@@ -301,19 +301,19 @@ pub fn parce_main_config(path: &str) -> AppConfig {
 
     cfo.proxy_tls_grade = parce_tls_grades(cfo.proxy_tls_grade.clone());
     if let Some(member) = cfo.cluster_member.clone() {
-        NODE_STATUS.set(Some(Arc::from(member.as_str()))).expect("Config can only be initialized once");
-        let mk = cfo.master_key.clone().expect("In CLuster mode master key must be set");
-        MASTER_KEY.set(Some(Arc::from(mk.as_str()))).expect("Config can only be initialized once");
+        NODE_STATUS.set(Some(Arc::from(member.as_str()))).die("Config can only be initialized once");
+        let mk = cfo.master_key.clone().die("In CLuster mode master key must be set");
+        MASTER_KEY.set(Some(Arc::from(mk.as_str()))).die("Config can only be initialized once");
         if member == "master" {
-            IM_A_MASTER.set(true).expect("Master status can only be initialized once");
-            let worker_str = cfo.worker_nodes.as_deref().expect("At least one worker must be set in master mode");
+            IM_A_MASTER.set(true).die("Master status can only be initialized once");
+            let worker_str = cfo.worker_nodes.as_deref().die("At least one worker must be set in master mode");
             let mut wks = Vec::new();
             for worker in worker_str.split(',') {
                 let clean = worker.replace("//", "");
                 let wk_vec = clean.trim().split(':').collect::<Vec<&str>>();
                 let wk = WorkerNodes {
                     address: Arc::from(wk_vec[1]),
-                    port: wk_vec[2].parse::<u16>().expect("Invalid worker port number"),
+                    port: wk_vec[2].parse::<u16>().die("Invalid worker port number"),
                     tls: match wk_vec[0] {
                         "https" => true,
                         "http" => false,
@@ -324,13 +324,13 @@ pub fn parce_main_config(path: &str) -> AppConfig {
                 wks.push(wk);
                 // println!("{:?}", wk);
                 // if let Some((host, port_str)) = worker.trim().split_once(':') {
-                //     let port = port_str.parse::<u16>().expect("Invalid worker port number");
+                //     let port = port_str.parse::<u16>().die("Invalid worker port number");
                 //     wks.push((host.to_string(), port));
                 // } else {
                 //     panic!("Invalid worker format, expected host:port");
                 // }
             }
-            WORKERS.set(wks).expect("Workers list can only be initialized once");
+            WORKERS.set(wks).die("Workers list can only be initialized once");
         }
     }
     cfo

src/utils/tools.rs

@@ -7,7 +7,7 @@ use privdrop::PrivDrop;
 use serde_json::{Value, json};
 use sha2::{Digest, Sha256};
 use std::collections::{HashMap, HashSet};
-use std::fmt::Write;
+use std::fmt::{Debug, Write};
 use std::fs::OpenOptions;
 use std::io::Write as IoWrite;
 use std::net::SocketAddr;
@@ -364,3 +364,38 @@ pub fn get_hostname() -> String {
         .filter(|s| !s.is_empty())
         .unwrap_or_else(|| "unknown".to_string())
 }
+
+// pub fn die(message: &str) -> ! {
+//     log::error!("{message}");
+//     process::exit(1);
+// }
+// ========================================= //
+
+pub trait Die<T> {
+    fn die(self, message: &str) -> T;
+}
+
+impl<T, E: Debug> Die<T> for Result<T, E> {
+    fn die(self, message: &str) -> T {
+        match self {
+            Ok(value) => value,
+            Err(e) => {
+                eprintln!("\x1b[1;31mERROR:\x1b[0m {message}: {e:?}");
+                process::exit(1);
+            }
+        }
+    }
+}
+
+impl<T> Die<T> for Option<T> {
+    fn die(self, message: &str) -> T {
+        match self {
+            Some(value) => value,
+            None => {
+                eprintln!("\x1b[1;31mERROR:\x1b[0m {message}");
+                process::exit(1);
+            }
+        }
+    }
+}
+// ========================================= //

src/utils/types.rs

@@ -112,6 +112,8 @@ pub struct AppConfig {
     pub hc_method: String,
     pub upstreams_conf: String,
     pub log_level: String,
+    pub cache_enabled: Option<bool>,
+    pub cache_storage: Option<String>,
     pub cache_size_mb: Option<usize>,
     pub cache_ttl: Option<u64>,
     pub access_log: Option<String>,