wip: feat: define response body enum

rpxy-lib/src/forwarder/cache/cache_error.rs

@@ -0,0 +1,35 @@
+use thiserror::Error;
+
+pub type CacheResult<T> = std::result::Result<T, CacheError>;
+
+/// Describes things that can go wrong in the Rpxy
+#[derive(Debug, Error)]
+pub enum CacheError {
+  // Cache errors,
+  #[error("Invalid null request and/or response")]
+  NullRequestOrResponse,
+
+  #[error("Failed to write byte buffer")]
+  FailedToWriteByteBufferForCache,
+
+  #[error("Failed to acquire mutex lock for cache")]
+  FailedToAcquiredMutexLockForCache,
+
+  #[error("Failed to create file cache")]
+  FailedToCreateFileCache,
+
+  #[error("Failed to write file cache")]
+  FailedToWriteFileCache,
+
+  #[error("Failed to open cache file")]
+  FailedToOpenCacheFile,
+
+  #[error("Too large to cache")]
+  TooLargeToCache,
+
+  #[error("Failed to cache bytes: {0}")]
+  FailedToCacheBytes(String),
+
+  #[error("Failed to send frame to cache {0}")]
+  FailedToSendFrameToCache(String),
+}

rpxy-lib/src/forwarder/cache/cache_main.rs

@@ -0,0 +1,368 @@
+use super::cache_error::*;
+use crate::{globals::Globals, hyper_ext::body::UnboundedStreamBody, log::*};
+use bytes::{Buf, Bytes, BytesMut};
+use futures::channel::mpsc;
+use http::{Request, Response};
+use http_body_util::{BodyExt, StreamBody};
+use http_cache_semantics::CachePolicy;
+use hyper::body::{Body, Frame, Incoming};
+use lru::LruCache;
+use std::{
+  convert::Infallible,
+  path::{Path, PathBuf},
+  sync::{
+    atomic::{AtomicUsize, Ordering},
+    Arc, Mutex,
+  },
+};
+use tokio::{
+  fs::{self, File},
+  io::{AsyncReadExt, AsyncWriteExt},
+  sync::RwLock,
+};
+
+/* ---------------------------------------------- */
+#[derive(Clone, Debug)]
+/// Cache main manager
+pub struct RpxyCache {
+  /// Inner lru cache manager storing http message caching policy
+  inner: LruCacheManager,
+  /// Managing cache file objects through RwLock's lock mechanism for file lock
+  file_store: FileStore,
+  /// Async runtime
+  runtime_handle: tokio::runtime::Handle,
+  /// Maximum size of each cache file object
+  max_each_size: usize,
+  /// Maximum size of cache object on memory
+  max_each_size_on_memory: usize,
+}
+
+impl RpxyCache {
+  /// Generate cache storage
+  pub async fn new(globals: &Globals) -> Option<Self> {
+    if !globals.proxy_config.cache_enabled {
+      return None;
+    }
+    let path = globals.proxy_config.cache_dir.as_ref().unwrap();
+    let file_store = FileStore::new(path, &globals.runtime_handle).await;
+    let inner = LruCacheManager::new(globals.proxy_config.cache_max_entry);
+
+    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;
+    if max_each_size < max_each_size_on_memory {
+      warn!(
+        "Maximum size of on memory cache per entry must be smaller than or equal to the maximum of each file cache"
+      );
+      max_each_size_on_memory = max_each_size;
+    }
+
+    Some(Self {
+      file_store,
+      inner,
+      runtime_handle: globals.runtime_handle.clone(),
+      max_each_size,
+      max_each_size_on_memory,
+    })
+  }
+
+  /// Count cache entries
+  pub async fn count(&self) -> (usize, usize, usize) {
+    let total = self.inner.count();
+    let file = self.file_store.count().await;
+    let on_memory = total - file;
+    (total, on_memory, file)
+  }
+
+  /// Put response into the cache
+  pub async fn put(
+    &self,
+    uri: &hyper::Uri,
+    mut body: Incoming,
+    policy: &CachePolicy,
+  ) -> CacheResult<UnboundedStreamBody> {
+    let my_cache = self.inner.clone();
+    let mut file_store = self.file_store.clone();
+    let uri = uri.clone();
+    let policy_clone = policy.clone();
+    let max_each_size = self.max_each_size;
+    let max_each_size_on_memory = self.max_each_size_on_memory;
+
+    let (body_tx, body_rx) = mpsc::unbounded::<Result<Frame<Bytes>, hyper::Error>>();
+
+    self.runtime_handle.spawn(async move {
+      let mut size = 0usize;
+      loop {
+        let frame = match body.frame().await {
+          Some(frame) => frame,
+          None => {
+            debug!("Response body finished");
+            break;
+          }
+        };
+        let frame_size = frame.as_ref().map(|f| {
+          if f.is_data() {
+            f.data_ref().map(|bytes| bytes.remaining()).unwrap_or_default()
+          } else {
+            0
+          }
+        });
+        size += frame_size.unwrap_or_default();
+
+        // check size
+        if size > max_each_size {
+          warn!("Too large to cache");
+          return Err(CacheError::TooLargeToCache);
+        }
+        frame
+          .as_ref()
+          .map(|f| {
+            if f.is_data() {
+              let data_bytes = f.data_ref().unwrap().clone();
+              println!("ddddde");
+              // TODO: cache data bytes as file or on memory
+              // fileにするかmemoryにするかの判断はある程度までバッファしてやってという手を使うことになる。途中までキャッシュしたやつはどうするかとかいう判断も必要。
+              // ファイルとObjectのbindをどうやってするか
+            }
+          })
+          .map_err(|e| CacheError::FailedToCacheBytes(e.to_string()))?;
+
+        // send data to use response downstream
+        body_tx
+          .unbounded_send(frame)
+          .map_err(|e| CacheError::FailedToSendFrameToCache(e.to_string()))?;
+      }
+
+      Ok(()) as CacheResult<()>
+    });
+
+    let stream_body = StreamBody::new(body_rx);
+
+    Ok(stream_body)
+  }
+}
+
+/* ---------------------------------------------- */
+#[derive(Debug, Clone)]
+/// Cache file manager outer that is responsible to handle `RwLock`
+struct FileStore {
+  inner: Arc<RwLock<FileStoreInner>>,
+}
+impl FileStore {
+  /// Build manager
+  async fn new(path: impl AsRef<Path>, runtime_handle: &tokio::runtime::Handle) -> Self {
+    Self {
+      inner: Arc::new(RwLock::new(FileStoreInner::new(path, runtime_handle).await)),
+    }
+  }
+}
+
+impl FileStore {
+  /// Count file cache entries
+  async fn count(&self) -> usize {
+    let inner = self.inner.read().await;
+    inner.cnt
+  }
+  /// Create a temporary file cache
+  async fn create(&mut self, cache_filename: &str, body_bytes: &Bytes) -> CacheResult<CacheFileOrOnMemory> {
+    let mut inner = self.inner.write().await;
+    inner.create(cache_filename, body_bytes).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>) -> CacheResult<Bytes> {
+  //   let inner = self.inner.read().await;
+  //   inner.read(&path).await
+  // }
+}
+
+#[derive(Debug)]
+/// Manager inner for cache on file system
+struct FileStoreInner {
+  /// Directory of temporary files
+  cache_dir: PathBuf,
+  /// Counter of current cached files
+  cnt: usize,
+  /// Async runtime
+  runtime_handle: tokio::runtime::Handle,
+}
+
+impl FileStoreInner {
+  /// Build new cache file manager.
+  /// This first creates cache file dir if not exists, and cleans up the file inside the directory.
+  /// TODO: Persistent cache is really difficult. `sqlite` or something like that 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 {
+      warn!("Failed to clean up the cache dir: {e}");
+    };
+    fs::create_dir_all(&path_buf).await.unwrap();
+    Self {
+      cache_dir: path_buf.clone(),
+      cnt: 0,
+      runtime_handle: runtime_handle.clone(),
+    }
+  }
+
+  /// Create a new temporary file cache
+  async fn create(&mut self, cache_filename: &str, body_bytes: &Bytes) -> CacheResult<CacheFileOrOnMemory> {
+    let cache_filepath = self.cache_dir.join(cache_filename);
+    let Ok(mut file) = File::create(&cache_filepath).await else {
+      return Err(CacheError::FailedToCreateFileCache);
+    };
+    let mut bytes_clone = body_bytes.clone();
+    while bytes_clone.has_remaining() {
+      if let Err(e) = file.write_buf(&mut bytes_clone).await {
+        error!("Failed to write file cache: {e}");
+        return Err(CacheError::FailedToWriteFileCache);
+      };
+    }
+    self.cnt += 1;
+    Ok(CacheFileOrOnMemory::File(cache_filepath))
+  }
+
+  /// Retrieve a stored temporary file cache
+  async fn read(&self, path: impl AsRef<Path>) -> CacheResult<()> {
+    let Ok(mut file) = File::open(&path).await else {
+      warn!("Cache file object cannot be opened");
+      return Err(CacheError::FailedToOpenCacheFile);
+    };
+
+    /* ----------------------------- */
+    // PoC for streaming body
+    let (tx, rx) = mpsc::unbounded::<Result<hyper::body::Frame<bytes::Bytes>, Infallible>>();
+
+    // let (body_sender, res_body) = Body::channel();
+    self.runtime_handle.spawn(async move {
+      //   let mut sender = body_sender;
+      let mut buf = BytesMut::new();
+      loop {
+        match file.read_buf(&mut buf).await {
+          Ok(0) => break,
+          Ok(_) => tx
+            .unbounded_send(Ok(hyper::body::Frame::data(buf.copy_to_bytes(buf.remaining()))))
+            .map_err(|e| anyhow::anyhow!("Failed to read cache file: {e}"))?,
+          //sender.send_data(buf.copy_to_bytes(buf.remaining())).await?,
+          Err(_) => break,
+        };
+      }
+      Ok(()) as anyhow::Result<()>
+    });
+
+    let mut rx = http_body_util::StreamBody::new(rx);
+    // TODO: 結局incominglikeなbodystreamを定義することになる。これだったらh3と合わせて自分で定義した方が良さそう。
+    // typeが長すぎるのでwrapperを作った方がいい。
+    // let response = Response::builder()
+    //   .status(200)
+    //   .header("content-type", "application/octet-stream")
+    //   .body(rx)
+    //   .unwrap();
+
+    todo!()
+    /* ----------------------------- */
+
+    // Ok(res_body)
+  }
+}
+
+/* ---------------------------------------------- */
+
+#[derive(Clone, Debug)]
+/// Cache target in hybrid manner of on-memory and file system
+pub enum CacheFileOrOnMemory {
+  /// Pointer to the temporary cache file
+  File(PathBuf),
+  /// Cached body itself
+  OnMemory(Vec<u8>),
+}
+
+#[derive(Clone, Debug)]
+/// Cache object definition
+struct CacheObject {
+  /// Cache policy to determine if the stored cache can be used as a response to a new incoming request
+  pub policy: CachePolicy,
+  /// Cache target: on-memory object or temporary file
+  pub target: CacheFileOrOnMemory,
+  /// SHA256 hash of target to strongly bind the cache metadata (this object) and file target
+  pub hash: Vec<u8>,
+}
+
+/* ---------------------------------------------- */
+#[derive(Debug, Clone)]
+/// Lru cache manager that is responsible to handle `Mutex` as an outer of `LruCache`
+struct LruCacheManager {
+  /// Inner lru cache manager main object
+  inner: Arc<Mutex<LruCache<String, CacheObject>>>, // TODO: keyはstring urlでいいのか疑問。全requestに対してcheckすることになりそう
+  /// Counter of current cached object (total)
+  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: Default::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);
+    // This may be inconsistent with the actual number of entries
+    self.cnt.store(lock.len(), Ordering::Relaxed);
+    res
+  }
+
+  /// Push an entry
+  fn push(&self, cache_key: &str, cache_object: CacheObject) -> CacheResult<Option<(String, CacheObject)>> {
+    let Ok(mut lock) = self.inner.lock() else {
+      error!("Failed to acquire mutex lock for writing cache entry");
+      return Err(CacheError::FailedToAcquiredMutexLockForCache);
+    };
+    let res = Ok(lock.push(cache_key.to_string(), cache_object));
+    // This may be inconsistent with the actual number of entries
+    self.cnt.store(lock.len(), Ordering::Relaxed);
+    res
+  }
+}
+
+/* ---------------------------------------------- */
+/// Generate cache policy if the response is cacheable
+pub fn get_policy_if_cacheable<B1, B2>(
+  req: Option<&Request<B1>>,
+  res: Option<&Response<B2>>,
+) -> CacheResult<Option<CachePolicy>>
+// where
+//   B1: core::fmt::Debug,
+{
+  // deduce cache policy from req and res
+  let (Some(req), Some(res)) = (req, res) else {
+    return Err(CacheError::NullRequestOrResponse);
+  };
+
+  let new_policy = CachePolicy::new(req, res);
+  if new_policy.is_storable() {
+    // debug!("Response is cacheable: {:?}\n{:?}", req, res.headers());
+    Ok(Some(new_policy))
+  } else {
+    Ok(None)
+  }
+}

rpxy-lib/src/forwarder/cache/mod.rs

@@ -0,0 +1,5 @@
+mod cache_error;
+mod cache_main;
+
+pub use cache_error::CacheError;
+pub use cache_main::{get_policy_if_cacheable, CacheFileOrOnMemory, RpxyCache};