pingap_cache/
lib.rs

1// Copyright 2024-2025 Tree xie.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use bytesize::ByteSize;
16use memory_stats::memory_stats;
17use once_cell::sync::OnceCell;
18use pingap_core::convert_query_map;
19use snafu::Snafu;
20use std::sync::atomic::{AtomicBool, Ordering};
21use std::sync::Arc;
22use tracing::info;
23
24mod file;
25mod http_cache;
26mod tiny;
27
28pub static PAGE_SIZE: usize = 4096;
29
30/// Category name for cache related logging
31pub static LOG_CATEGORY: &str = "cache";
32
33#[derive(Debug, Snafu)]
34pub enum Error {
35    #[snafu(display("Io error: {source}"))]
36    Io { source: std::io::Error },
37    #[snafu(display("{message}"))]
38    Invalid { message: String },
39    #[snafu(display("Over quota error, max: {max}, {message}"))]
40    OverQuota { max: u32, message: String },
41    #[snafu(display("{message}"))]
42    Prometheus { message: String },
43}
44pub type Result<T, E = Error> = std::result::Result<T, E>;
45
46impl From<Error> for pingora::BError {
47    fn from(value: Error) -> Self {
48        pingap_core::new_internal_error(500, value.to_string())
49    }
50}
51
52fn new_tiny_ufo_cache(mode: &str, size: usize) -> HttpCache {
53    HttpCache {
54        directory: None,
55        cache: Arc::new(tiny::new_tiny_ufo_cache(mode, size / PAGE_SIZE, size)),
56    }
57}
58fn new_file_cache(dir: &str) -> Result<HttpCache> {
59    let cache = file::new_file_cache(dir)?;
60    Ok(HttpCache {
61        directory: Some(cache.directory.clone()),
62        cache: Arc::new(cache),
63    })
64}
65
66static CACHE_BACKEND: OnceCell<HttpCache> = OnceCell::new();
67const MAX_MEMORY_SIZE: usize = 100 * 1024 * 1024;
68static CACHED_INIT: AtomicBool = AtomicBool::new(false);
69
70/// Check if the cache backend has been initialized
71pub fn is_cache_backend_init() -> bool {
72    CACHED_INIT.load(Ordering::Relaxed)
73}
74
75#[derive(Debug, Default)]
76pub struct CacheBackendOption {
77    /// Directory to store cache files
78    pub cache_directory: Option<String>,
79    /// Maximum size of cache storage
80    pub cache_max_size: Option<ByteSize>,
81}
82
83/// Get the cache backend
84pub fn get_cache_backend(
85    option: Option<CacheBackendOption>,
86) -> Result<&'static HttpCache> {
87    if !is_cache_backend_init() && option.is_none() {
88        return Err(Error::Invalid {
89            message: "cache backend is not initialized".to_string(),
90        });
91    };
92    let option = option.unwrap_or_default();
93    // Get or initialize the global cache backend using OnceCell
94    CACHE_BACKEND.get_or_try_init(|| {
95        // let basic_conf = &get_current_config().basic;
96        // Determine cache size from config or use default MAX_MEMORY_SIZE
97        let mut size = if let Some(cache_max_size) = option.cache_max_size {
98            cache_max_size.as_u64() as usize
99        } else {
100            MAX_MEMORY_SIZE
101        };
102
103        let mut cache_type = "memory";
104        let mut cache_mode = "".to_string();
105        // Get optional cache directory from config
106        let cache_directory =
107            if let Some(cache_directory) = &option.cache_directory {
108                cache_directory.trim().to_string()
109            } else {
110                "".to_string()
111            };
112
113        // Choose between file-based or memory-based cache
114        let cache = if !cache_directory.is_empty()
115            && !cache_directory.starts_with("memory://")
116        {
117            // Use file-based cache if directory is specified
118            cache_type = "file";
119            new_file_cache(cache_directory.as_str()).map_err(|e| {
120                Error::Invalid {
121                    message: e.to_string(),
122                }
123            })?
124        } else {
125            // For memory cache, limit size to half of available physical memory
126            // or fallback to 256MB if memory stats unavailable
127            let max_memory = if let Some(value) = memory_stats() {
128                value.physical_mem * 1024 / 2
129            } else {
130                ByteSize::mb(256).as_u64() as usize
131            };
132
133            if let Some((_, query)) = cache_directory.split_once('?') {
134                let query_map = convert_query_map(query);
135                cache_mode = query_map.get("mode").cloned().unwrap_or_default();
136            }
137
138            size = size.min(max_memory);
139            // Create memory-based tiny UFO cache
140            new_tiny_ufo_cache(&cache_mode, size)
141        };
142        CACHED_INIT.store(true, Ordering::Relaxed);
143
144        // Log cache initialization details
145        info!(
146            category = LOG_CATEGORY,
147            size = ByteSize::b(size as u64).to_string(),
148            cache_type,
149            cache_mode,
150            support_clear = cache.cache.support_clear(),
151            "init cache backend success"
152        );
153        Ok(cache)
154    })
155}
156
157pub use http_cache::{new_storage_clear_service, HttpCache};
158
159#[cfg(feature = "full")]
160mod prom;
161#[cfg(feature = "full")]
162pub use prom::{CACHE_READING_TIME, CACHE_WRITING_TIME};
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167    use pretty_assertions::assert_eq;
168    use tempfile::TempDir;
169
170    #[test]
171    fn test_convert_error() {
172        let err = Error::Invalid {
173            message: "invalid error".to_string(),
174        };
175
176        let b_error: pingora::BError = err.into();
177
178        assert_eq!(
179            " HTTPStatus context: invalid error cause:  InternalError",
180            b_error.to_string()
181        );
182    }
183
184    #[test]
185    fn test_is_cache_backend_init() {
186        assert_eq!(false, is_cache_backend_init());
187    }
188    #[test]
189    fn test_cache() {
190        let _ = new_tiny_ufo_cache("compact", 1024);
191
192        let dir = TempDir::new().unwrap();
193        let result = new_file_cache(&dir.into_path().to_string_lossy());
194        assert_eq!(true, result.is_ok());
195    }
196}
pingap_cache/lib.rs

pingap_cache/
lib.rs
