tuxmain/rust-rpxy
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #80 from junkurihara/refactor

Browse source 
Refactor: Make cache feature code safer and clean
This commit is contained in:
Jun Kurihara 2023-08-22 21:20:01 +09:00 • committed by GitHub
commit 2beb627786
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 156 additions and 76 deletions
Download patch file Download diff file Expand all files Collapse all files

231
rpxy-lib/src/handler/cache.rs
View file

@ -11,7 +11,10 @@ use sha2::{Digest, Sha256};
use std::{ use std::{
fmt::Debug, fmt::Debug,
path::{Path, PathBuf}, path::{Path, PathBuf},
sync::{Arc, Mutex}, sync::{
atomic::{AtomicUsize, Ordering},
Arc, Mutex,
},
time::SystemTime, time::SystemTime,
}; };
use tokio::{ use tokio::{
@ -21,41 +24,53 @@ use tokio::{
}; };
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
/// Cache target in hybrid manner of on-memory and file system
pub enum CacheFileOrOnMemory { pub enum CacheFileOrOnMemory {
/// Pointer to the temporary cache file
File(PathBuf), File(PathBuf),
/// Cached body itself
OnMemory(Vec<u8>), OnMemory(Vec<u8>),
} }
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
/// Cache object definition
struct CacheObject { struct CacheObject {
/// Cache policy to determine if the stored cache can be used as a response to a new incoming request
pub policy: CachePolicy, pub policy: CachePolicy,
/// Cache target: on-memory object or temporary file
pub target: CacheFileOrOnMemory, pub target: CacheFileOrOnMemory,
} }
#[derive(Debug)] #[derive(Debug)]
struct CacheFileManager { /// Manager inner for cache on file system
struct CacheFileManagerInner {
/// Directory of temporary files
cache_dir: PathBuf, cache_dir: PathBuf,
/// Counter of current cached files
cnt: usize, cnt: usize,
/// Async runtime
runtime_handle: tokio::runtime::Handle, runtime_handle: tokio::runtime::Handle,
} }
impl CacheFileManager { impl CacheFileManagerInner {
async fn new(path: &PathBuf, runtime_handle: &tokio::runtime::Handle) -> Self { /// Build new cache file manager.
// Create cache file dir /// This first creates cache file dir if not exists, and cleans up the file inside the directory.
// Clean up the file dir before init /// TODO: Persistent cache is really difficult. `sqlite` or something like that is needed.
// TODO: Persistent cache is really difficult. maybe SQLite is needed. async fn new(path: impl AsRef<Path>, runtime_handle: &tokio::runtime::Handle) -> Self {
let path_buf = path.as_ref().to_path_buf();
if let Err(e) = fs::remove_dir_all(path).await { if let Err(e) = fs::remove_dir_all(path).await {
warn!("Failed to clean up the cache dir: {e}"); warn!("Failed to clean up the cache dir: {e}");
}; };
fs::create_dir_all(path).await.unwrap(); fs::create_dir_all(&path_buf).await.unwrap();
Self { Self {
cache_dir: path.clone(), cache_dir: path_buf.clone(),
cnt: 0, cnt: 0,
runtime_handle: runtime_handle.clone(), runtime_handle: runtime_handle.clone(),
} }
} }
async fn create(&mut self, cache_filename: &str, body_bytes: &Bytes, policy: &CachePolicy) -> Result<CacheObject> { /// Create a new temporary file cache
async fn create(&mut self, cache_filename: &str, body_bytes: &Bytes) -> Result<CacheFileOrOnMemory> {
let cache_filepath = self.cache_dir.join(cache_filename); let cache_filepath = self.cache_dir.join(cache_filename);
let Ok(mut file) = File::create(&cache_filepath).await else { let Ok(mut file) = File::create(&cache_filepath).await else {
return Err(RpxyError::Cache("Failed to create file")); return Err(RpxyError::Cache("Failed to create file"));
@ -68,12 +83,10 @@ impl CacheFileManager {
}; };
} }
self.cnt += 1; self.cnt += 1;
Ok(CacheObject { Ok(CacheFileOrOnMemory::File(cache_filepath))
policy: policy.clone(),
target: CacheFileOrOnMemory::File(cache_filepath),
})
} }
/// Retrieve a stored temporary file cache
async fn read(&self, path: impl AsRef<Path>) -> Result<Body> { async fn read(&self, path: impl AsRef<Path>) -> Result<Body> {
let Ok(mut file) = File::open(&path).await else { let Ok(mut file) = File::open(&path).await else {
warn!("Cache file object cannot be opened"); warn!("Cache file object cannot be opened");
@ -96,6 +109,7 @@ impl CacheFileManager {
Ok(res_body) Ok(res_body)
} }
/// Remove file
async fn remove(&mut self, path: impl AsRef<Path>) -> Result<()> { async fn remove(&mut self, path: impl AsRef<Path>) -> Result<()> {
fs::remove_file(path.as_ref()).await?; fs::remove_file(path.as_ref()).await?;
self.cnt -= 1; self.cnt -= 1;
@ -105,12 +119,103 @@ impl CacheFileManager {
} }
} }
#[derive(Debug, Clone)]
/// Cache file manager outer that is responsible to handle `RwLock`
struct CacheFileManager {
inner: Arc<RwLock<CacheFileManagerInner>>,
}
impl CacheFileManager {
/// Build manager
async fn new(path: impl AsRef<Path>, runtime_handle: &tokio::runtime::Handle) -> Self {
Self {
inner: Arc::new(RwLock::new(CacheFileManagerInner::new(path, runtime_handle).await)),
}
}
/// Evict a temporary file cache
async fn evict(&self, path: impl AsRef<Path>) {
// Acquire the write lock
let mut inner = self.inner.write().await;
if let Err(e) = inner.remove(path).await {
warn!("Eviction failed during file object removal: {:?}", e);
};
}
/// Read a temporary file cache
async fn read(&self, path: impl AsRef<Path>) -> Result<Body> {
let mgr = self.inner.read().await;
mgr.read(&path).await
}
/// Create a temporary file cache
async fn create(&mut self, cache_filename: &str, body_bytes: &Bytes) -> Result<CacheFileOrOnMemory> {
let mut mgr = self.inner.write().await;
mgr.create(cache_filename, body_bytes).await
}
async fn count(&self) -> usize {
let mgr = self.inner.read().await;
mgr.cnt
}
}
#[derive(Debug, Clone)]
/// Lru cache manager that is responsible to handle `Mutex` as an outer of `LruCache`
struct LruCacheManager {
inner: Arc<Mutex<LruCache<String, CacheObject>>>, // TODO: keyはstring urlでいいのか疑問。全requestに対してcheckすることになりそう
cnt: Arc<AtomicUsize>,
}
impl LruCacheManager {
/// Build LruCache
fn new(cache_max_entry: usize) -> Self {
Self {
inner: Arc::new(Mutex::new(LruCache::new(
std::num::NonZeroUsize::new(cache_max_entry).unwrap(),
))),
cnt: Arc::new(AtomicUsize::default()),
}
}
/// Count entries
fn count(&self) -> usize {
self.cnt.load(Ordering::Relaxed)
}
/// Evict an entry
fn evict(&self, cache_key: &str) -> Option<(String, CacheObject)> {
let Ok(mut lock) = self.inner.lock() else {
error!("Mutex can't be locked to evict a cache entry");
return None;
};
let res = lock.pop_entry(cache_key);
self.cnt.store(lock.len(), Ordering::Relaxed);
res
}
/// Get an entry
fn get(&self, cache_key: &str) -> Result<Option<CacheObject>> {
let Ok(mut lock) = self.inner.lock() else {
error!("Mutex can't be locked for checking cache entry");
return Err(RpxyError::Cache("Mutex can't be locked for checking cache entry"));
};
let Some(cached_object) = lock.get(cache_key) else {
return Ok(None);
};
Ok(Some(cached_object.clone()))
}
/// Push an entry
fn push(&self, cache_key: &str, cache_object: CacheObject) -> Result<Option<(String, CacheObject)>> {
let Ok(mut lock) = self.inner.lock() else {
error!("Failed to acquire mutex lock for writing cache entry");
return Err(RpxyError::Cache("Failed to acquire mutex lock for writing cache entry"));
};
let res = Ok(lock.push(cache_key.to_string(), cache_object));
self.cnt.store(lock.len(), Ordering::Relaxed);
res
}
}
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct RpxyCache { pub struct RpxyCache {
/// Managing cache file objects through RwLock's lock mechanism for file lock /// Managing cache file objects through RwLock's lock mechanism for file lock
cache_file_manager: Arc<RwLock<CacheFileManager>>, cache_file_manager: CacheFileManager,
/// Lru cache storing http message caching policy /// Lru cache storing http message caching policy
inner: Arc<Mutex<LruCache<String, CacheObject>>>, // TODO: keyはstring urlでいいのか疑問。全requestに対してcheckすることになりそう inner: LruCacheManager,
/// Async runtime /// Async runtime
runtime_handle: tokio::runtime::Handle, runtime_handle: tokio::runtime::Handle,
/// Maximum size of each cache file object /// Maximum size of each cache file object
@ -127,10 +232,8 @@ impl RpxyCache {
} }
let path = globals.proxy_config.cache_dir.as_ref().unwrap(); let path = globals.proxy_config.cache_dir.as_ref().unwrap();
let cache_file_manager = Arc::new(RwLock::new(CacheFileManager::new(path, &globals.runtime_handle).await)); let cache_file_manager = CacheFileManager::new(path, &globals.runtime_handle).await;
let inner = Arc::new(Mutex::new(LruCache::new( let inner = LruCacheManager::new(globals.proxy_config.cache_max_entry);
std::num::NonZeroUsize::new(globals.proxy_config.cache_max_entry).unwrap(),
)));
let max_each_size = globals.proxy_config.cache_max_each_size; let max_each_size = globals.proxy_config.cache_max_each_size;
let mut max_each_size_on_memory = globals.proxy_config.cache_max_each_size_on_memory; let mut max_each_size_on_memory = globals.proxy_config.cache_max_each_size_on_memory;
@ -150,37 +253,25 @@ impl RpxyCache {
}) })
} }
fn evict_cache_entry(&self, cache_key: &str) -> Option<(String, CacheObject)> { /// Count cache entries
let Ok(mut lock) = self.inner.lock() else { pub async fn count(&self) -> (usize, usize, usize) {
error!("Mutex can't be locked to evict a cache entry"); let total = self.inner.count();
return None; let file = self.cache_file_manager.count().await;
}; let on_memory = total - file;
lock.pop_entry(cache_key) (total, on_memory, file)
}
async fn evict_cache_file(&self, filepath: impl AsRef<Path>) {
// Acquire the write lock
let mut mgr = self.cache_file_manager.write().await;
if let Err(e) = mgr.remove(filepath).await {
warn!("Eviction failed during file object removal: {:?}", e);
};
} }
/// Get cached response /// Get cached response
pub async fn get<R>(&self, req: &Request<R>) -> Option<Response<Body>> { pub async fn get<R>(&self, req: &Request<R>) -> Option<Response<Body>> {
debug!("Current cache entries: {:?}", self.inner); debug!(
"Current cache status: (total, on-memory, file) = {:?}",
self.count().await
);
let cache_key = req.uri().to_string(); let cache_key = req.uri().to_string();
// First check cache chance // First check cache chance
let cached_object = { let Ok(Some(cached_object)) = self.inner.get(&cache_key) else {
let Ok(mut lock) = self.inner.lock() else { return None;
error!("Mutex can't be locked for checking cache entry");
return None;
};
let Some(cached_object) = lock.get(&cache_key) else {
return None;
};
cached_object.clone()
}; };
// Secondly check the cache freshness as an HTTP message // Secondly check the cache freshness as an HTTP message
@ -190,11 +281,11 @@ impl RpxyCache {
// This might be okay to keep as is since it would be updated later. // This might be okay to keep as is since it would be updated later.
// However, there is no guarantee that newly got objects will be still cacheable. // However, there is no guarantee that newly got objects will be still cacheable.
// So, we have to evict stale cache entries and cache file objects if found. // So, we have to evict stale cache entries and cache file objects if found.
debug!("Stale cache entry and file object: {cache_key}"); debug!("Stale cache entry: {cache_key}");
let _evicted_entry = self.evict_cache_entry(&cache_key); let _evicted_entry = self.inner.evict(&cache_key);
// For cache file // For cache file
if let CacheFileOrOnMemory::File(path) = cached_object.target { if let CacheFileOrOnMemory::File(path) = &cached_object.target {
self.evict_cache_file(&path).await; self.cache_file_manager.evict(&path).await;
} }
return None; return None;
}; };
@ -202,13 +293,12 @@ impl RpxyCache {
// Finally retrieve the file/on-memory object // Finally retrieve the file/on-memory object
match cached_object.target { match cached_object.target {
CacheFileOrOnMemory::File(path) => { CacheFileOrOnMemory::File(path) => {
let mgr = self.cache_file_manager.read().await; let res_body = match self.cache_file_manager.read(&path).await {
let res_body = match mgr.read(&path).await {
Ok(res_body) => res_body, Ok(res_body) => res_body,
Err(e) => { Err(e) => {
warn!("Failed to read from file cache: {e}"); warn!("Failed to read from file cache: {e}");
let _evicted_entry = self.evict_cache_entry(&cache_key); let _evicted_entry = self.inner.evict(&cache_key);
self.evict_cache_file(&path).await; self.cache_file_manager.evict(&path).await;
return None; return None;
} }
}; };
@ -223,9 +313,10 @@ impl RpxyCache {
} }
} }
/// Put response into the cache
pub async fn put(&self, uri: &hyper::Uri, body_bytes: &Bytes, policy: &CachePolicy) -> Result<()> { pub async fn put(&self, uri: &hyper::Uri, body_bytes: &Bytes, policy: &CachePolicy) -> Result<()> {
let my_cache = self.inner.clone(); let my_cache = self.inner.clone();
let mgr = self.cache_file_manager.clone(); let mut mgr = self.cache_file_manager.clone();
let uri = uri.clone(); let uri = uri.clone();
let bytes_clone = body_bytes.clone(); let bytes_clone = body_bytes.clone();
let policy_clone = policy.clone(); let policy_clone = policy.clone();
@ -238,18 +329,17 @@ impl RpxyCache {
return Err(RpxyError::Cache("Too large to cache")); return Err(RpxyError::Cache("Too large to cache"));
} }
let cache_key = derive_cache_key_from_uri(&uri); let cache_key = derive_cache_key_from_uri(&uri);
let cache_filename = derive_filename_from_uri(&uri);
debug!("Cache file of {:?} bytes to be written", bytes_clone.len()); debug!("Object of size {:?} bytes to be cached", bytes_clone.len());
let cache_object = if bytes_clone.len() > max_each_size_on_memory { let cache_object = if bytes_clone.len() > max_each_size_on_memory {
let mut mgr = mgr.write().await; let cache_filename = derive_filename_from_uri(&uri);
let Ok(cache_object) = mgr.create(&cache_filename, &bytes_clone, &policy_clone).await else { let target = mgr.create(&cache_filename, &bytes_clone).await?;
error!("Failed to put the body into the file object or cache entry"); debug!("Cached a new cache file: {} - {}", cache_key, cache_filename);
return Err(RpxyError::Cache("Failed to put the body into the file object or cache entry")); CacheObject {
}; policy: policy_clone,
debug!("Cached a new file: {} - {}", cache_key, cache_filename); target,
cache_object }
} else { } else {
debug!("Cached a new object on memory: {}", cache_key); debug!("Cached a new object on memory: {}", cache_key);
CacheObject { CacheObject {
@ -257,21 +347,12 @@ impl RpxyCache {
target: CacheFileOrOnMemory::OnMemory(bytes_clone.to_vec()), target: CacheFileOrOnMemory::OnMemory(bytes_clone.to_vec()),
} }
}; };
let push_opt = {
let Ok(mut lock) = my_cache.lock() else { if let Some((k, v)) = my_cache.push(&cache_key, cache_object)? {
error!("Failed to acquire mutex lock for writing cache entry");
return Err(RpxyError::Cache("Failed to acquire mutex lock for writing cache entry"));
};
lock.push(cache_key.clone(), cache_object)
};
if let Some((k, v)) = push_opt {
if k != cache_key { if k != cache_key {
info!("Over the cache capacity. Evict least recent used entry"); info!("Over the cache capacity. Evict least recent used entry");
if let CacheFileOrOnMemory::File(path) = v.target { if let CacheFileOrOnMemory::File(path) = v.target {
let mut mgr = mgr.write().await; mgr.evict(&path).await;
if let Err(e) = mgr.remove(&path).await {
warn!("Eviction failed during file object removal over the capacity: {:?}", e);
};
} }
} }
} }
@ -304,7 +385,7 @@ where
let new_policy = CachePolicy::new(req, res); let new_policy = CachePolicy::new(req, res);
if new_policy.is_storable() { if new_policy.is_storable() {
debug!("Response is cacheable: {:?}\n{:?}", req, res.headers()); // debug!("Response is cacheable: {:?}\n{:?}", req, res.headers());
Ok(Some(new_policy)) Ok(Some(new_policy))
} else { } else {
Ok(None) Ok(None)

1
rpxy-lib/src/handler/forwarder.rs
View file

@ -87,7 +87,6 @@ where
return res; return res;
}; };
let (parts, body) = res.unwrap().into_parts(); let (parts, body) = res.unwrap().into_parts();
// TODO: Inefficient?
let Ok(mut bytes) = hyper::body::aggregate(body).await else { let Ok(mut bytes) = hyper::body::aggregate(body).await else {
return Err(RpxyError::Cache("Failed to write byte buffer")); return Err(RpxyError::Cache("Failed to write byte buffer"));
}; };