wip: feat: define response body enum

rpxy-lib/src/error.rs

@@ -84,35 +84,15 @@ pub enum RpxyError {
   #[error("Failed to fetch from upstream: {0}")]
   FailedToFetchFromUpstream(String),
 
-  // Cache errors,
-  #[cfg(feature = "cache")]
-  #[error("Invalid null request and/or response")]
-  NullRequestOrResponse,
-
-  #[cfg(feature = "cache")]
-  #[error("Failed to write byte buffer")]
-  FailedToWriteByteBufferForCache,
-
-  #[cfg(feature = "cache")]
-  #[error("Failed to acquire mutex lock for cache")]
-  FailedToAcquiredMutexLockForCache,
-
-  #[cfg(feature = "cache")]
-  #[error("Failed to create file cache")]
-  FailedToCreateFileCache,
-
-  #[cfg(feature = "cache")]
-  #[error("Failed to write file cache")]
-  FailedToWriteFileCache,
-
-  #[cfg(feature = "cache")]
-  #[error("Failed to open cache file")]
-  FailedToOpenCacheFile,
-
   // Upstream connection setting errors
   #[error("Unsupported upstream option")]
   UnsupportedUpstreamOption,
 
+  // Cache error map
+  #[cfg(feature = "cache")]
+  #[error("Cache error: {0}")]
+  CacheError(#[from] crate::forwarder::CacheError),
+
   // Others
   #[error("Infallible")]
   Infallible(#[from] std::convert::Infallible),

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

@@ -1,8 +1,11 @@
-use crate::{error::*, globals::Globals, log::*};
+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::StreamBody;
+use http_body_util::{BodyExt, StreamBody};
 use http_cache_semantics::CachePolicy;
+use hyper::body::{Body, Frame, Incoming};
 use lru::LruCache;
 use std::{
   convert::Infallible,
@@ -69,6 +72,73 @@ impl RpxyCache {
     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)
+  }
 }
 
 /* ---------------------------------------------- */
@@ -93,7 +163,7 @@ impl FileStore {
     inner.cnt
   }
   /// Create a temporary file cache
-  async fn create(&mut self, cache_filename: &str, body_bytes: &Bytes) -> RpxyResult<CacheFileOrOnMemory> {
+  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
   }
@@ -106,7 +176,7 @@ impl FileStore {
   //   };
   // }
   // /// Read a temporary file cache
-  // async fn read(&self, path: impl AsRef<Path>) -> RpxyResult<Bytes> {
+  // async fn read(&self, path: impl AsRef<Path>) -> CacheResult<Bytes> {
   //   let inner = self.inner.read().await;
   //   inner.read(&path).await
   // }
@@ -141,16 +211,16 @@ impl FileStoreInner {
   }
 
   /// Create a new temporary file cache
-  async fn create(&mut self, cache_filename: &str, body_bytes: &Bytes) -> RpxyResult<CacheFileOrOnMemory> {
+  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(RpxyError::FailedToCreateFileCache);
+      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(RpxyError::FailedToWriteFileCache);
+        return Err(CacheError::FailedToWriteFileCache);
       };
     }
     self.cnt += 1;
@@ -158,15 +228,14 @@ impl FileStoreInner {
   }
 
   /// Retrieve a stored temporary file cache
-  async fn read(&self, path: impl AsRef<Path>) -> RpxyResult<()> {
+  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(RpxyError::FailedToOpenCacheFile);
+      return Err(CacheError::FailedToOpenCacheFile);
     };
 
     /* ----------------------------- */
     // PoC for streaming body
-    use futures::channel::mpsc;
     let (tx, rx) = mpsc::unbounded::<Result<hyper::body::Frame<bytes::Bytes>, Infallible>>();
 
     // let (body_sender, res_body) = Body::channel();
@@ -263,10 +332,10 @@ impl LruCacheManager {
   }
 
   /// Push an entry
-  fn push(&self, cache_key: &str, cache_object: CacheObject) -> RpxyResult<Option<(String, CacheObject)>> {
+  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(RpxyError::FailedToAcquiredMutexLockForCache);
+      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
@@ -280,13 +349,13 @@ impl LruCacheManager {
 pub fn get_policy_if_cacheable<B1, B2>(
   req: Option<&Request<B1>>,
   res: Option<&Response<B2>>,
-) -> RpxyResult<Option<CachePolicy>>
+) -> 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(RpxyError::NullRequestOrResponse);
+    return Err(CacheError::NullRequestOrResponse);
   };
 
   let new_policy = CachePolicy::new(req, res);

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};

rpxy-lib/src/forwarder/client.rs

@@ -1,14 +1,10 @@
 use crate::{
   error::{RpxyError, RpxyResult},
   globals::Globals,
-  hyper_ext::{
+  hyper_ext::{body::ResponseBody, rt::LocalExecutor},
-    body::{wrap_incoming_body_response, BoxBody, IncomingOr},
-    rt::LocalExecutor,
-  },
   log::*,
 };
 use async_trait::async_trait;
-use chrono::Duration;
 use http::{Request, Response, Version};
 use hyper::body::{Body, Incoming};
 use hyper_util::client::legacy::{
@@ -19,10 +15,6 @@ use std::sync::Arc;
 
 #[cfg(feature = "cache")]
 use super::cache::{get_policy_if_cacheable, RpxyCache};
-#[cfg(feature = "cache")]
-use crate::hyper_ext::body::{full, wrap_synthetic_body_response};
-#[cfg(feature = "cache")]
-use http_body_util::BodyExt;
 
 #[async_trait]
 /// Definition of the forwarder that simply forward requests from downstream client to upstream app servers.
@@ -40,7 +32,7 @@ pub struct Forwarder<C, B> {
 }
 
 #[async_trait]
-impl<C, B1> ForwardRequest<B1, IncomingOr<BoxBody>> for Forwarder<C, B1>
+impl<C, B1> ForwardRequest<B1, ResponseBody> for Forwarder<C, B1>
 where
   C: Send + Sync + Connect + Clone + 'static,
   B1: Body + Send + Sync + Unpin + 'static,
@@ -49,7 +41,7 @@ where
 {
   type Error = RpxyError;
 
-  async fn request(&self, req: Request<B1>) -> Result<Response<IncomingOr<BoxBody>>, Self::Error> {
+  async fn request(&self, req: Request<B1>) -> Result<Response<ResponseBody>, Self::Error> {
     // TODO: cache handling
     #[cfg(feature = "cache")]
     {
@@ -67,38 +59,27 @@ where
       let res = self.request_directly(req).await;
 
       if self.cache.is_none() {
-        return res.map(wrap_incoming_body_response::<BoxBody>);
+        return res.map(|inner| inner.map(ResponseBody::Incoming));
       }
 
       // check cacheability and store it if cacheable
       let Ok(Some(cache_policy)) = get_policy_if_cacheable(synth_req.as_ref(), res.as_ref().ok()) else {
-        return res.map(wrap_incoming_body_response::<BoxBody>);
+        return res.map(|inner| inner.map(ResponseBody::Incoming));
       };
       let (parts, body) = res.unwrap().into_parts();
 
-      // TODO: This is inefficient since current strategy needs to copy the whole body onto memory to cache it.
+      // Get streamed body without waiting for the arrival of the body,
-      // This should be handled by copying buffer simultaneously while forwarding response to downstream.
+      // which is done simultaneously with caching.
-      let Ok(bytes) = body.collect().await.map(|v| v.to_bytes()) else {
+      let stream_body = self
-        return Err(RpxyError::FailedToWriteByteBufferForCache);
+        .cache
-      };
+        .as_ref()
-      let bytes_clone = bytes.clone();
+        .unwrap()
+        .put(synth_req.unwrap().uri(), body, &cache_policy)
+        .await?;
 
-      // TODO: this is inefficient. needs to be reconsidered to avoid unnecessary copy and should spawn async task to store cache.
+      // response with body being cached in background
-      // We may need to use the same logic as h3.
+      let new_res = Response::from_parts(parts, ResponseBody::Streamed(stream_body));
-      // Is bytes.clone() enough?
+      Ok(new_res)
-
-      // if let Err(cache_err) = self
-      //   .cache
-      //   .as_ref()
-      //   .unwrap()
-      //   .put(synth_req.unwrap().uri(), &bytes, &cache_policy)
-      //   .await
-      // {
-      //   error!("{:?}", cache_err);
-      // };
-
-      // response with cached body
-      Ok(wrap_synthetic_body_response(Response::from_parts(parts, full(bytes))))
     }
 
     // No cache handling
@@ -107,7 +88,7 @@ where
       self
         .request_directly(req)
         .await
-        .map(wrap_incoming_body_response::<BoxBody>)
+        .map(|inner| inner.map(ResponseBody::Incoming))
     }
   }
 }

rpxy-lib/src/forwarder/mod.rs

@@ -3,6 +3,9 @@ mod cache;
 mod client;
 
 use crate::hyper_ext::body::{IncomingLike, IncomingOr};
-pub type Forwarder<C> = client::Forwarder<C, IncomingOr<IncomingLike>>;
 
-pub use client::ForwardRequest;
+pub(crate) type Forwarder<C> = client::Forwarder<C, IncomingOr<IncomingLike>>;
+pub(crate) use client::ForwardRequest;
+
+#[cfg(feature = "cache")]
+pub(crate) use cache::CacheError;

rpxy-lib/src/hyper_ext/body_type.rs

@@ -1,24 +1,25 @@
-use http::Response;
+// use http::Response;
 use http_body_util::{combinators, BodyExt, Either, Empty, Full};
-use hyper::body::{Bytes, Incoming};
+use hyper::body::{Body, Bytes, Incoming};
+use std::pin::Pin;
 
 /// Type for synthetic boxed body
 pub(crate) type BoxBody = combinators::BoxBody<Bytes, hyper::Error>;
 /// Type for either passthrough body or given body type, specifically synthetic boxed body
 pub(crate) type IncomingOr<B> = Either<Incoming, B>;
 
-/// helper function to build http response with passthrough body
+// /// helper function to build http response with passthrough body
-pub(crate) fn wrap_incoming_body_response<B>(response: Response<Incoming>) -> Response<IncomingOr<B>>
+// pub(crate) fn wrap_incoming_body_response<B>(response: Response<Incoming>) -> Response<IncomingOr<B>>
-where
+// where
-  B: hyper::body::Body,
+//   B: hyper::body::Body,
-{
+// {
-  response.map(IncomingOr::Left)
+//   response.map(IncomingOr::Left)
-}
+// }
 
-/// helper function to build http response with synthetic body
+// /// helper function to build http response with synthetic body
-pub(crate) fn wrap_synthetic_body_response<B>(response: Response<B>) -> Response<IncomingOr<B>> {
+// pub(crate) fn wrap_synthetic_body_response<B>(response: Response<B>) -> Response<IncomingOr<B>> {
-  response.map(IncomingOr::Right)
+//   response.map(IncomingOr::Right)
-}
+// }
 
 /// helper function to build a empty body
 pub(crate) fn empty() -> BoxBody {
@@ -29,3 +30,43 @@ pub(crate) fn empty() -> BoxBody {
 pub(crate) fn full(body: Bytes) -> BoxBody {
   Full::new(body).map_err(|never| match never {}).boxed()
 }
+
+/* ------------------------------------ */
+#[cfg(feature = "cache")]
+use futures::channel::mpsc::UnboundedReceiver;
+#[cfg(feature = "cache")]
+use http_body_util::StreamBody;
+#[cfg(feature = "cache")]
+use hyper::body::Frame;
+
+#[cfg(feature = "cache")]
+pub(crate) type UnboundedStreamBody = StreamBody<UnboundedReceiver<Result<Frame<bytes::Bytes>, hyper::Error>>>;
+
+/// Response body use in this project
+/// - Incoming: just a type that only forwards the upstream response body to downstream.
+/// - BoxedCache: a type that is generated from cache, e.g.,, small byte object.
+/// - StreamedCache: another type that is generated from cache as stream, e.g., large byte object.
+pub(crate) enum ResponseBody {
+  Incoming(Incoming),
+  Boxed(BoxBody),
+  #[cfg(feature = "cache")]
+  Streamed(UnboundedStreamBody),
+}
+
+impl Body for ResponseBody {
+  type Data = bytes::Bytes;
+  type Error = hyper::Error;
+
+  fn poll_frame(
+    self: Pin<&mut Self>,
+    cx: &mut std::task::Context<'_>,
+  ) -> std::task::Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
+    match self.get_mut() {
+      ResponseBody::Incoming(incoming) => Pin::new(incoming).poll_frame(cx),
+      #[cfg(feature = "cache")]
+      ResponseBody::Boxed(boxed) => Pin::new(boxed).poll_frame(cx),
+      #[cfg(feature = "cache")]
+      ResponseBody::Streamed(streamed) => Pin::new(streamed).poll_frame(cx),
+    }
+  }
+}

rpxy-lib/src/hyper_ext/mod.rs

@@ -12,7 +12,5 @@ pub(crate) mod rt {
 #[allow(unused)]
 pub(crate) mod body {
   pub(crate) use super::body_incoming_like::IncomingLike;
-  pub(crate) use super::body_type::{
+  pub(crate) use super::body_type::{empty, full, BoxBody, IncomingOr, ResponseBody, UnboundedStreamBody};
-    empty, full, wrap_incoming_body_response, wrap_synthetic_body_response, BoxBody, IncomingOr,
-  };
 }

rpxy-lib/src/message_handler/handler_main.rs

@@ -11,7 +11,7 @@ use crate::{
   error::*,
   forwarder::{ForwardRequest, Forwarder},
   globals::Globals,
-  hyper_ext::body::{BoxBody, IncomingLike, IncomingOr},
+  hyper_ext::body::{IncomingLike, IncomingOr, ResponseBody},
   log::*,
   name_exp::ServerName,
 };
@@ -58,7 +58,7 @@ where
     listen_addr: SocketAddr,
     tls_enabled: bool,
     tls_server_name: Option<ServerName>,
-  ) -> RpxyResult<Response<IncomingOr<BoxBody>>> {
+  ) -> RpxyResult<Response<ResponseBody>> {
     // preparing log data
     let mut log_data = HttpMessageLog::from(&req);
     log_data.client_addr(&client_addr);
@@ -99,7 +99,7 @@ where
     listen_addr: SocketAddr,
     tls_enabled: bool,
     tls_server_name: Option<ServerName>,
-  ) -> HttpResult<Response<IncomingOr<BoxBody>>> {
+  ) -> HttpResult<Response<ResponseBody>> {
     // Here we start to inspect and parse with server_name
     let server_name = req
       .inspect_parse_host()

rpxy-lib/src/message_handler/synthetic_response.rs

@@ -1,16 +1,16 @@
 use super::http_result::{HttpError, HttpResult};
 use crate::{
   error::*,
-  hyper_ext::body::{empty, BoxBody, IncomingOr},
+  hyper_ext::body::{empty, ResponseBody},
   name_exp::ServerName,
 };
 use http::{Request, Response, StatusCode, Uri};
 
 /// build http response with status code of 4xx and 5xx
-pub(crate) fn synthetic_error_response(status_code: StatusCode) -> RpxyResult<Response<IncomingOr<BoxBody>>> {
+pub(crate) fn synthetic_error_response(status_code: StatusCode) -> RpxyResult<Response<ResponseBody>> {
   let res = Response::builder()
     .status(status_code)
-    .body(IncomingOr::Right(empty()))
+    .body(ResponseBody::Boxed(empty()))
     .unwrap();
   Ok(res)
 }
@@ -20,7 +20,7 @@ pub(super) fn secure_redirection_response<B>(
   server_name: &ServerName,
   tls_port: Option<u16>,
   req: &Request<B>,
-) -> HttpResult<Response<IncomingOr<BoxBody>>> {
+) -> HttpResult<Response<ResponseBody>> {
   let server_name: String = server_name.try_into().unwrap_or_default();
   let pq = match req.uri().path_and_query() {
     Some(x) => x.as_str(),
@@ -36,7 +36,7 @@ pub(super) fn secure_redirection_response<B>(
   let response = Response::builder()
     .status(StatusCode::MOVED_PERMANENTLY)
     .header("Location", dest_uri.to_string())
-    .body(IncomingOr::Right(empty()))
+    .body(ResponseBody::Boxed(empty()))
     .map_err(|e| HttpError::FailedToRedirect(e.to_string()))?;
   Ok(response)
 }

rpxy-lib/src/proxy/proxy_h3.rs

@@ -138,7 +138,6 @@ where
     });
 
     let new_req: Request<IncomingOr<IncomingLike>> = Request::from_parts(req_parts, IncomingOr::Right(req_body));
-    // Response<IncomingOr<BoxBody>> wrapped by RpxyResult
     let res = self
       .message_handler
       .handle_request(

rpxy-lib/src/proxy/proxy_main.rs

@@ -5,7 +5,7 @@ use crate::{
   error::*,
   globals::Globals,
   hyper_ext::{
-    body::{BoxBody, IncomingOr},
+    body::{IncomingOr, ResponseBody},
     rt::LocalExecutor,
   },
   log::*,
@@ -32,7 +32,7 @@ async fn serve_request<U, T>(
   listen_addr: SocketAddr,
   tls_enabled: bool,
   tls_server_name: Option<ServerName>,
-) -> RpxyResult<Response<IncomingOr<BoxBody>>>
+) -> RpxyResult<Response<ResponseBody>>
 where
   T: Send + Sync + Connect + Clone,
   U: CryptoSource + Clone,