diff --git a/Cargo.toml b/Cargo.toml index 38669ce..808f158 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,6 +34,10 @@ harness = false name = "nuget_cache" harness = false +[[example]] +name = "deviceatlas" +harness = false + [dependencies] hyper = { version = "1.2", features = ["full"] } hyper-tls = { version = "0.6.0" } @@ -53,6 +57,8 @@ futures-core = "0.3.30" prometheus = "0.13" serde-inline-default = "0.2.3" enum_dispatch = "0.3.13" +foyer = "0.17" +postcard = "1.1.3" [dev-dependencies] tonic = "0.11" diff --git a/examples/deviceatlas/config.yaml b/examples/deviceatlas/config.yaml new file mode 100644 index 0000000..c0a9531 --- /dev/null +++ b/examples/deviceatlas/config.yaml @@ -0,0 +1,28 @@ +listening_port: 8080 +prometheus_port: 8081 +heathcheck_port: 8082 +minimum_log_level: info + +middlewares: + - log: + level: info + value: 'Cache status: $cache_status, Forward target: $forward_target' + - add_response_header: + name: 'X-Cache-Status' + value: '$cache_status' + - use_cache: + cache_name: files + when: + status_code_is: 200 + - forward_request: + scheme: https + target_host: deviceatlas.com + +caches: + - hybrid: + name: files + path: /tmp/cacheus/deviceatlas + ttl_seconds: 10 + hash_path: false + hash_query: true + hash_body: false \ No newline at end of file diff --git a/examples/deviceatlas/main.rs b/examples/deviceatlas/main.rs new file mode 100644 index 0000000..60463cb --- /dev/null +++ b/examples/deviceatlas/main.rs @@ -0,0 +1,17 @@ +use cacheus::{self, CacheusServer}; +use futures::join; +use log::info; +use tokio::signal; + +#[tokio::main] +async fn main() +{ + let cacheus_fut = CacheusServer::start_from_config_file("examples/deviceatlas/config.yaml"); + + join!(cacheus_fut); + + // Wait for SIGTERM or SIGINT + signal::ctrl_c().await.expect("failed to listen for ctrl_c"); + + info!("Shutting down Cacheus server..."); +} diff --git a/examples/download/config.yaml b/examples/download/config.yaml index 3d9fe6a..086027a 100644 --- a/examples/download/config.yaml +++ b/examples/download/config.yaml @@ -1,7 +1,5 @@ listening_port: 3001 prometheus_port: 8081 -https: false -http2: false minimum_log_level: debug middlewares: diff --git a/examples/hello/config.yaml b/examples/hello/config.yaml index bff5f64..77be924 100644 --- a/examples/hello/config.yaml +++ b/examples/hello/config.yaml @@ -1,4 +1,2 @@ default_target_host: 127.0.0.1:3002 -listening_port: 3001 -https: false -http2: false \ No newline at end of file +listening_port: 3001 \ No newline at end of file diff --git a/examples/hello_grpc/config.yaml b/examples/hello_grpc/config.yaml index fc8397a..8957fa0 100644 --- a/examples/hello_grpc/config.yaml +++ b/examples/hello_grpc/config.yaml @@ -1,3 +1,2 @@ listening_port: 3001 -https: false -http2: true \ No newline at end of file +http2_only: true \ No newline at end of file diff --git a/examples/nuget_cache/config.yaml b/examples/nuget_cache/config.yaml index b096845..c49e781 100644 --- a/examples/nuget_cache/config.yaml +++ b/examples/nuget_cache/config.yaml @@ -1,8 +1,6 @@ --- listening_port: 4002 prometheus_port: 4004 -https: true -http2: false minimum_log_level: 'debug' middlewares: diff --git a/src/buffered_body.rs b/src/buffered_body.rs index 5b60410..7086887 100644 --- a/src/buffered_body.rs +++ b/src/buffered_body.rs @@ -9,6 +9,7 @@ use futures::Future; use hyper::body::{Body, Frame}; use hyper::HeaderMap; use pin_project_lite::pin_project; +use serde::{Deserialize, Serialize}; pin_project! { /// Future that resolves into a [`Collected`]. @@ -56,6 +57,71 @@ pub struct BufferedBody trailers: Option, } +impl BufferedBody { + pub fn to_bytes(&self) -> Vec { + let mut v = Vec::new(); + + // Serialize bufs length + data + let data = self.bufs.as_ref(); + let len = data.len() as u64; + v.extend_from_slice(&len.to_le_bytes()); + v.extend_from_slice(data); + + // Serialize trailers as optional key/value map + if let Some(t) = &self.trailers { + v.push(1); // present flag + let map: HashMap> = t + .iter() + .map(|(k, v)| (k.to_string(), v.to_str().map(|s| vec![s.to_string()]).unwrap_or_default())) + .collect(); + let mut buffer = Vec::new(); + let encoded = postcard::to_slice(&map, &mut buffer).unwrap(); + v.extend_from_slice(&(encoded.len() as u64).to_le_bytes()); + v.extend_from_slice(&encoded); + } else { + v.push(0); // not present + } + + v + } + + pub fn from_bytes(bytes: &[u8]) -> Self { + use std::convert::TryInto; + let mut pos = 0; + + // bufs + let len = u64::from_le_bytes(bytes[pos..pos + 8].try_into().unwrap()) as usize; + pos += 8; + let buf = BytesMut::from(&bytes[pos..pos + len]); + pos += len; + + // trailers + let has_trailers = bytes[pos] == 1; + pos += 1; + + let trailers = if has_trailers { + let map_len = u64::from_le_bytes(bytes[pos..pos + 8].try_into().unwrap()) as usize; + pos += 8; + let map_bytes = &bytes[pos..pos + map_len]; + let map: HashMap> = postcard::from_bytes(map_bytes).unwrap(); + let mut hm = HeaderMap::new(); + for (k, vs) in map { + for v in vs { + hm.append( + k.parse::().unwrap(), + v.parse::().unwrap(), + ); + } + } + Some(hm) + } else { + None + }; + + BufferedBody { bufs: buf, trailers } + } +} + impl BufferedBody { pub fn replace_strings(&mut self, find: &String, replacement: &String) -> usize @@ -114,7 +180,7 @@ impl BufferedBody } } - pub fn from_bytes(b: &[u8]) -> BufferedBody + pub fn from_body(b: &[u8]) -> BufferedBody { let mut bufs = BytesMut::new(); bufs.extend(b); diff --git a/src/config/cache.rs b/src/config/cache.rs index 40a983b..c7d5242 100644 --- a/src/config/cache.rs +++ b/src/config/cache.rs @@ -1,12 +1,15 @@ -use std::collections::HashMap; +use std::io::{Read, Write}; +use std::path::Path; +use std::sync::Arc; +use std::{collections::HashMap}; use std::hash::Hash; -use std::time::Duration; -use gxhash::GxHasher; +use gxhash::{GxBuildHasher, GxHasher}; use serde::{Serialize, Deserialize}; use serde_inline_default::serde_inline_default; use crate::buffered_body::BufferedBody; -use crate::{lru, ShardedCache}; -use hyper::{Request, Response}; +// use crate::{lru, ShardedCache}; +use foyer::{*}; +use hyper::{HeaderMap, Request, Response, StatusCode}; /// Cache configuration types #[derive(Debug, Serialize, Deserialize, Clone)] @@ -15,6 +18,9 @@ pub enum CacheConfig { InMemory { in_memory: MemoryCache, }, + Hybrid { + hybrid: HybridCacheConfig, + }, } /// In-memory cache configuration @@ -26,11 +32,37 @@ pub struct MemoryCache { #[serde_inline_default(8)] pub shards: usize, - #[serde_inline_default(0)] - pub probatory_size: usize, + #[serde_inline_default(128 * 1024 * 1024)] + pub memory_size_bytes: usize, + + #[serde_inline_default(3600)] + pub ttl_seconds: u64, + + #[serde_inline_default(true)] + pub hash_path: bool, + + #[serde_inline_default(true)] + pub hash_query: bool, - #[serde_inline_default(100_000)] - pub resident_size: usize, + #[serde_inline_default(true)] + pub hash_body: bool, +} + +#[serde_inline_default] +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct HybridCacheConfig { + pub name: String, + + pub path: String, + + #[serde_inline_default(128 * 1024 * 1024)] + pub memory_size_bytes: usize, + + #[serde_inline_default(128 * 1024 * 1024)] + pub disk_size_bytes: usize, + + #[serde_inline_default(8)] + pub shards: usize, #[serde_inline_default(3600)] pub ttl_seconds: u64, @@ -45,6 +77,96 @@ pub struct MemoryCache { pub hash_body: bool, } +#[derive(Debug, Clone)] +pub struct CachedResponse { + pub insertion_epoch: u64, + pub response: Response, +} + +impl foyer::Code for CachedResponse { + fn encode(&self, writer: &mut impl Write) -> std::result::Result<(), foyer::CodeError> { + let status = self.response.status().as_u16(); + writer.write_all(&status.to_be_bytes()).unwrap(); + + // Serialize headers + let headers = self.response.headers(); + let header_count = headers.len() as u32; + writer.write_all(&header_count.to_be_bytes()).unwrap(); + for (name, value) in headers.iter() { + let name_bytes = name.as_str().as_bytes(); + let value_bytes = value.as_bytes(); + + writer.write_all(&(name_bytes.len() as u32).to_be_bytes()).unwrap(); + writer.write_all(name_bytes).unwrap(); + writer.write_all(&(value_bytes.len() as u32).to_be_bytes()).unwrap(); + writer.write_all(value_bytes).unwrap(); + } + + // Serialize body (assuming BufferedBody can provide bytes) + let body_bytes = self.response.body().to_bytes(); + writer.write_all(&(body_bytes.len() as u64).to_be_bytes()).unwrap(); + writer.write_all(&body_bytes).unwrap(); + + writer.write_all(&self.insertion_epoch.to_le_bytes()).unwrap(); + + Ok(()) + } + + fn decode(reader: &mut impl Read) -> std::result::Result { + let mut buf2 = [0u8; 2]; + reader.read_exact(&mut buf2).unwrap(); + let status = StatusCode::from_u16(u16::from_be_bytes(buf2)).map_err(|_| foyer::CodeError::SizeLimit).unwrap(); + + let mut buf4 = [0u8; 4]; + reader.read_exact(&mut buf4).unwrap(); + let header_count = u32::from_be_bytes(buf4); + + let mut headers = HeaderMap::new(); + for _ in 0..header_count { + reader.read_exact(&mut buf4).unwrap(); + let name_len = u32::from_be_bytes(buf4) as usize; + let mut name_buf = vec![0u8; name_len]; + reader.read_exact(&mut name_buf).unwrap(); + let name = String::from_utf8(name_buf).map_err(|_| foyer::CodeError::SizeLimit).unwrap(); + + reader.read_exact(&mut buf4).unwrap(); + let value_len = u32::from_be_bytes(buf4) as usize; + let mut value_buf = vec![0u8; value_len]; + reader.read_exact(&mut value_buf).unwrap(); + headers.insert( + hyper::header::HeaderName::try_from(name).map_err(|_| foyer::CodeError::SizeLimit).unwrap(), + hyper::header::HeaderValue::from_bytes(&value_buf).map_err(|_| foyer::CodeError::SizeLimit).unwrap(), + ); + } + + let mut buf8 = [0u8; 8]; + reader.read_exact(&mut buf8)?; + let body_len = u64::from_be_bytes(buf8); + let mut body_buf = vec![0u8; body_len as usize]; + reader.read_exact(&mut body_buf)?; + + let buffered = BufferedBody::from_body(&body_buf); + let mut response = Response::builder() + .status(status) + .body(buffered) + .map_err(|_| foyer::CodeError::SizeLimit).unwrap(); + + *response.headers_mut() = headers; + + let mut insertion_epoch_buf = [0u8; size_of::()]; + reader.read_exact(&mut insertion_epoch_buf).unwrap(); + let insertion_epoch = u64::from_le_bytes(insertion_epoch_buf); + + Ok(Self { insertion_epoch, response }) + } + + fn estimated_size(&self) -> usize { + let body_len = self.response.body().to_bytes().len(); + let headers_len: usize = self.response.headers().iter().map(|(k, v)| k.as_str().len() + v.as_bytes().len()).sum(); + 2 + 4 + headers_len + 8 + body_len + } +} + impl CacheConfig { /// Create a cache key from a request pub fn create_key(&self, request: &Request) -> u128 { @@ -65,22 +187,64 @@ impl CacheConfig { } hasher.finish_u128() }, + CacheConfig::Hybrid { hybrid: hybrid_config } => { + let mut hasher = GxHasher::with_seed(123); + // Different path/query means different key + if hybrid_config.hash_path { + request.uri().path().hash(&mut hasher); + } + if hybrid_config.hash_query { + if let Some(query) = request.uri().query() { + query.hash(&mut hasher); + } + } + if hybrid_config.hash_body { + request.body().hash(&mut hasher); + } + hasher.finish_u128() + } } } /// Add cache to the provided cache map - pub fn add_cache(&self, caches: &mut HashMap>)>) { + pub async fn add_cache(&self, caches: &mut HashMap)>) { + match self { CacheConfig::InMemory { in_memory } => { - let cache = ShardedCache::>::new( - in_memory.shards, - in_memory.probatory_size, - in_memory.resident_size, - Duration::from_secs(in_memory.ttl_seconds), - lru::ExpirationType::Absolute, - ); + let cache = HybridCacheBuilder::new() + .with_name(in_memory.name.clone()) + .with_flush_on_close(true) + .with_policy(foyer::HybridCachePolicy::WriteOnInsertion) + .memory(in_memory.memory_size_bytes) + .with_hash_builder(gxhash::GxBuildHasher::with_seed(123)) + .storage(Engine::Large) + .with_admission_picker(Arc::new(foyer::RejectAllPicker::default())) + .build() + .await + .unwrap(); + caches.insert(in_memory.name.clone(), (self.clone(), cache)); }, + CacheConfig::Hybrid { hybrid: hybrid_config } => { + let hybrid = HybridCacheBuilder::new() + .with_name(hybrid_config.name.clone()) + .with_flush_on_close(true) + .with_policy(foyer::HybridCachePolicy::WriteOnInsertion) + .memory(hybrid_config.memory_size_bytes) + .with_hash_builder(gxhash::GxBuildHasher::with_seed(123)) + .storage(Engine::Large) + .with_admission_picker(Arc::new(foyer::AdmitAllPicker::default())) + .with_device_options( + DirectFsDeviceOptions::new(hybrid_config.path.clone()).with_capacity(hybrid_config.disk_size_bytes), + ) + .with_flush(true) + .with_recover_mode(foyer::RecoverMode::Strict) + .build() + .await + .unwrap(); + + caches.insert(hybrid_config.name.clone(), (self.clone(), hybrid)); + } } } } diff --git a/src/config/conditions.rs b/src/config/conditions.rs index 340e652..b1cb1fd 100644 --- a/src/config/conditions.rs +++ b/src/config/conditions.rs @@ -25,11 +25,11 @@ impl Condition { } } - pub fn evaluate_with_response(&self, request: &Request, response: &Response) -> bool { + pub fn evaluate_with_response(&self, request: &Request, response: &Option>) -> bool { match self { - Condition::HeaderExists { header_exists } => request.headers().contains_key(header_exists), + Condition::HeaderExists { header_exists } => request.headers().contains_key(header_exists), // Response headers to consider? Condition::PathContains { path_contains } => contains_wildcard(&request.uri().path_and_query().unwrap().as_str(), path_contains), - Condition::StatusCodeIs { status_code_is } => response.status().as_u16() == *status_code_is, + Condition::StatusCodeIs { status_code_is } => response.as_ref().map_or(false, |resp| resp.status().as_u16() == *status_code_is), Condition::All { all } => all.iter().all(|c| c.evaluate_with_response(request, response)), Condition::Any { any } => any.iter().any(|c| c.evaluate_with_response(request, response)), Condition::Not { not } => !not.evaluate_with_response(request, response), diff --git a/src/config/configuration.rs b/src/config/configuration.rs index 1ca1412..18ab4f1 100644 --- a/src/config/configuration.rs +++ b/src/config/configuration.rs @@ -12,8 +12,8 @@ pub struct Configuration { pub prometheus_port: u16, #[serde_inline_default(8001)] pub healthcheck_port: u16, - #[serde_inline_default(true)] - pub http2: bool, + #[serde_inline_default(false)] + pub http2_only: bool, #[serde_inline_default("info".to_string())] pub minimum_log_level: String, #[serde(default)] diff --git a/src/config/middleware.rs b/src/config/middleware.rs index 1290c85..184d5f7 100644 --- a/src/config/middleware.rs +++ b/src/config/middleware.rs @@ -1,20 +1,21 @@ use std::{str::FromStr, sync::Arc}; use hyper::{header::{HeaderName, HeaderValue}, Request, Response, Uri}; +use postcard::fixint::le; use serde::{Deserialize, Serialize}; use serde_inline_default::serde_inline_default; -use crate::{buffered_body::BufferedBody, config::conditions::Condition, status::Status, CacheusServer, CallContext}; +use crate::{buffered_body::BufferedBody, config::{cache::CachedResponse, conditions::Condition, CacheConfig}, status::Status, CacheusServer, CallContext}; trait Middleware { - async fn on_request(&self, context: &mut CallContext, service: Arc, request: &mut Request) -> Option<(Arc>, Status)>; - async fn on_response(&self, context: &mut CallContext, service: Arc, request: &Request, response: &mut Response); + async fn on_request(&self, context: &mut CallContext, service: Arc, request: &mut Request) -> Option>; + async fn on_response(&self, context: &mut CallContext, service: Arc, request: &Request, response: &mut Option>); } #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(untagged)] pub enum MiddlewareEnum { Block { block_request: Block }, - Cache { use_cache: Cache }, // TODO: Support HTTP cache headers, validation, etc. + Cache { use_cache: CacheMiddleware }, // TODO: Support HTTP cache headers, validation, etc. Forward { forward_request: Forward }, // TODO: Support multiple targets, DNS load balancing, redirects, etc. AddHeader { add_response_header: AddHeader }, // TODO: If before forward: add request header, else add response header? Or separate types? ReplaceResponseBody { replace_response_body: ReplaceResponseBody }, @@ -23,7 +24,7 @@ pub enum MiddlewareEnum { } impl MiddlewareEnum { - pub async fn on_request(&self, context: &mut CallContext, service: Arc, request: &mut Request) -> Option<(Arc>, Status)> { + pub async fn on_request(&self, context: &mut CallContext, service: Arc, request: &mut Request) -> Option> { debug!("Processing middleware request: {:?}", self); match self { MiddlewareEnum::Block { block_request } => block_request.on_request(context, service, request).await, @@ -35,7 +36,7 @@ impl MiddlewareEnum { } } - pub async fn on_response(&self, context: &mut CallContext, service: Arc, request: &Request, response: &mut Response) { + pub async fn on_response(&self, context: &mut CallContext, service: Arc, request: &Request, response: &mut Option>) { debug!("Processing middleware response: {:?}", self); match self { MiddlewareEnum::Block { block_request } => block_request.on_response(context, service, request, response).await, @@ -58,48 +59,85 @@ pub struct Block { } impl Middleware for Block { - async fn on_request(&self, _context: &mut CallContext, _service: Arc, request: &mut Request) -> Option<(Arc>, Status)> { + async fn on_request(&self, _context: &mut CallContext, _service: Arc, request: &mut Request) -> Option> { if !self.when.as_ref().is_none_or(|w| w.evaluate(&request)) { return None; } - return Some((Arc::new(Response::builder().status(self.status_code).body(BufferedBody::from_bytes(b"")).unwrap()), Status::Block)); + return Some(Response::builder().status(self.status_code).body(BufferedBody::from_body(b"Blocked by cacheus")).unwrap()); } - async fn on_response(&self, _context: &mut CallContext, _service: Arc, _request: &Request, _response: &mut Response) { + async fn on_response(&self, _context: &mut CallContext, _service: Arc, _request: &Request, _response: &mut Option>) { } } #[derive(Debug, Serialize, Deserialize, Clone)] -pub struct Cache { +pub struct CacheMiddleware { pub cache_name: String, #[serde(default)] pub when: Option, } -impl Middleware for Cache { - async fn on_request(&self, _context: &mut CallContext, service: Arc, request: &mut Request) -> Option<(Arc>, Status)> { +// - Stale +// [Cache::on_request] lookup: expired -> [Forward::on_request] failed -> [Forward::on_response] -> [Cache::on_response] no response: lookup: expired: mark stale +// - Miss +// [Cache::on_request] lookup: not found -> [Forward::on_request] -> [Forward::on_response] -> [Cache::on_response] response: insert +// - Hit +// [Cache::on_request] lookup: found -> [Cache::on_response] response + hit context: no insert +impl Middleware for CacheMiddleware { + async fn on_request(&self, context: &mut CallContext, service: Arc, request: &mut Request) -> Option> { if !self.when.as_ref().is_none_or(|w| w.evaluate(&request)) { return None; } if let Some((cache_config, cache_instance)) = service.caches.get(&self.cache_name) { let key = cache_config.create_key(&request); - if let Some(cached_response) = cache_instance.try_get_locked(&key) { - _context.variables.insert("$cache_status".to_string(), "hit".to_string()); - return Some((Arc::new(cached_response.as_ref().clone()), Status::Hit)); + if let Some(cached_response) = cache_instance.get(&key).await.ok().flatten() { + let current_epoch: u64 = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs(); + let ttl_seconds = match cache_config { + CacheConfig::InMemory { in_memory } => in_memory.ttl_seconds, + CacheConfig::Hybrid { hybrid } => hybrid.ttl_seconds + }; + if current_epoch - cached_response.value().insertion_epoch > ttl_seconds { + // Stale entry! We keep it in case we need it later + context.stale_response.replace(cached_response.value().response.clone()); + } + else + { + // Cache hit! + context.variables.insert("$cache_status".to_string(), "hit".to_string()); + return Some(cached_response.value().response.clone()); + } } } - _context.variables.insert("$cache_status".to_string(), "miss".to_string()); + context.variables.insert("$cache_status".to_string(), "miss".to_string()); return None; } - async fn on_response(&self, context: &mut CallContext, service: Arc, request: &Request, response: &mut Response) { - if !self.when.as_ref().is_none_or(|w| w.evaluate_with_response(&request, &response)) { + async fn on_response(&self, context: &mut CallContext, service: Arc, request: &Request, response: &mut Option>) { + // Here we have a problem: imagine we only want to cache 200 OK responses, so we add when: status_code_is: 200. + // If the response is 404, we may want to return a stale instead of a miss. But since the when condition fails, we won't even look into the cache. + // To handle such case, we handle the stale logic before the when condition. + let has_response = response.is_some(); + let when_applies = self.when.as_ref().is_none_or(|w| w.evaluate_with_response(&request, &response)); + if (!has_response || !when_applies) && context.stale_response.is_some() { + context.variables.insert("$cache_status".to_string(), "stale".to_string()); + response.replace(context.stale_response.take().unwrap()); + } + if !when_applies { return; } - if let Some((cache_config, cache_instance)) = service.caches.get(&self.cache_name) { - let key = cache_config.create_key(&request); - cache_instance.try_add_arc_locked(key, response.clone()); + if has_response + { + // Response: either miss (fetched from subsequent middleware) or hit (taken from cache) + // In case of miss, we insert the response into the cache + if context.variables.get("$cache_status").map_or(false, |v| v == "hit") { + return; + } + if let Some((cache_config, cache_instance)) = service.caches.get(&self.cache_name) { + let key = cache_config.create_key(&request); + let insertion_epoch: u64 = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs(); + cache_instance.insert(key, CachedResponse { insertion_epoch: insertion_epoch, response: response.clone().unwrap() }); + } } } } @@ -112,7 +150,7 @@ pub struct Forward { } impl Middleware for Forward { - async fn on_request(&self, context: &mut CallContext, service: Arc, request: &mut Request) -> Option<(Arc>, Status)> { + async fn on_request(&self, context: &mut CallContext, service: Arc, request: &mut Request) -> Option> { let target_host = match request.headers().get("x-target-host") { Some(value) => value.to_str().unwrap().to_string(), None => self.target_host.clone(), @@ -145,10 +183,10 @@ impl Middleware for Forward { let res = service.client.request(forwarded_req).await.expect("Failed to send request"); let (parts, body) = res.into_parts(); let buffered_body = BufferedBody::collect_buffered(body).await.unwrap(); - return Some((Arc::new(Response::from_parts(parts, buffered_body)), Status::Miss)); + return Some(Response::from_parts(parts, buffered_body)); } - async fn on_response(&self, _context: &mut CallContext, _service: Arc, _request: &Request, _response: &mut Response) { + async fn on_response(&self, _context: &mut CallContext, _service: Arc, _request: &Request, _response: &mut Option>) { } } @@ -162,20 +200,22 @@ pub struct AddHeader { } impl Middleware for AddHeader { - async fn on_request(&self, _context: &mut CallContext, _service: Arc, _request: &mut Request) -> Option<(Arc>, Status)> { + async fn on_request(&self, _context: &mut CallContext, _service: Arc, _request: &mut Request) -> Option> { None } - async fn on_response(&self, _context: &mut CallContext, _service: Arc, request: &Request, response: &mut Response) { + async fn on_response(&self, _context: &mut CallContext, _service: Arc, request: &Request, response: &mut Option>) { if !self.when.as_ref().is_none_or(|w| w.evaluate_with_response(&request, &response)) { return; } - // Replace variables in header value, if any - let value = _context.variables.iter().fold(self.value.clone(), |acc, (k, v)| acc.replace(k.as_str(), v.as_str())); - // Insert the header - response.headers_mut().insert( - HeaderName::from_str(self.name.as_str()).unwrap(), // Clearly not optimal - HeaderValue::from_str(value.as_str()).unwrap()); + if let Some(response) = response { + // Replace variables in header value, if any + let value = _context.variables.iter().fold(self.value.clone(), |acc, (k, v)| acc.replace(k.as_str(), v.as_str())); + // Insert the header + response.headers_mut().insert( + HeaderName::from_str(self.name.as_str()).unwrap(), // Clearly not optimal + HeaderValue::from_str(value.as_str()).unwrap()); + } } } @@ -188,19 +228,21 @@ pub struct ReplaceResponseBody { } impl Middleware for ReplaceResponseBody { - async fn on_request(&self, _context: &mut CallContext, _service: Arc, _request: &mut Request) -> Option<(Arc>, Status)> { + async fn on_request(&self, _context: &mut CallContext, _service: Arc, _request: &mut Request) -> Option> { None } - async fn on_response(&self, _context: &mut CallContext, _service: Arc, request: &Request, response: &mut Response) { + async fn on_response(&self, _context: &mut CallContext, _service: Arc, request: &Request, response: &mut Option>) { if !self.when.as_ref().is_none_or(|w| w.evaluate_with_response(&request, &response)) { return; } - if let Some(content_type) = response.headers().get("content-type") { - if content_type.to_str().unwrap().contains("json") { - let content_length = response.body_mut().replace_strings(&self.find, &self.replacement); - // Response length may have changed, so we need to update the content-length header - response.headers_mut().insert("content-length", content_length.to_string().parse().unwrap()); + if let Some(response) = response { + if let Some(content_type) = response.headers().get("content-type") { + if content_type.to_str().unwrap().contains("json") { + let content_length = response.body_mut().replace_strings(&self.find, &self.replacement); + // Response length may have changed, so we need to update the content-length header + response.headers_mut().insert("content-length", content_length.to_string().parse().unwrap()); + } } } } @@ -215,11 +257,11 @@ pub struct Log { } impl Middleware for Log { - async fn on_request(&self, _context: &mut CallContext, _service: Arc, _request: &mut Request) -> Option<(Arc>, Status)> { + async fn on_request(&self, _context: &mut CallContext, _service: Arc, _request: &mut Request) -> Option> { None } - async fn on_response(&self, _context: &mut CallContext, _service: Arc, request: &Request, response: &mut Response) { + async fn on_response(&self, _context: &mut CallContext, _service: Arc, request: &Request, response: &mut Option>) { if !self.when.as_ref().is_none_or(|w| w.evaluate_with_response(&request, &response)) { return; } diff --git a/src/config/mod.rs b/src/config/mod.rs index 862a83c..f718f02 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -8,4 +8,4 @@ pub use configuration::Configuration; pub use cache::{CacheConfig, MemoryCache}; pub use conditions::Condition; -pub use middleware::{MiddlewareEnum, Block, Cache, Forward, AddHeader}; +pub use middleware::{MiddlewareEnum, Block, CacheMiddleware, Forward, AddHeader}; diff --git a/src/lib.rs b/src/lib.rs index 32bc61a..ac6323c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,7 +2,7 @@ extern crate log; mod buffered_body; -mod caches; +// mod caches; mod collections; // pub mod config; pub mod config; @@ -18,12 +18,13 @@ use std::ops::Deref; use std::str::FromStr; use std::sync::Arc; -pub use caches::*; +// pub use caches::*; pub use collections::*; pub use config::{Configuration, MiddlewareEnum}; use buffered_body::BufferedBody; use executor::TokioExecutor; use futures::join; +use gxhash::GxBuildHasher; use hyper::body::{Body, Incoming}; use hyper::server::conn::{http1, http2}; use hyper::service::service_fn; @@ -38,12 +39,14 @@ use metrics::Metrics; use simplelog::{ColorChoice, CombinedLogger, Config, TermLogger, TerminalMode}; use status::Status; use tokio::net::TcpListener; +use foyer::HybridCache; +use crate::config::cache::CachedResponse; use crate::config::CacheConfig; pub struct CacheusServer { configuration: Configuration, - caches: HashMap>)>, + caches: HashMap)>, metrics: Metrics, client: Client, BufferedBody>, } @@ -86,7 +89,7 @@ impl CacheusServer // Initialize caches from configuration for cache in &configuration.caches { - cache.add_cache(&mut caches); + cache.add_cache(&mut caches).await; } let server = Arc::new(CacheusServer { @@ -94,7 +97,7 @@ impl CacheusServer caches: caches, metrics: Metrics::new(), client: Client::builder(TokioExecutor) - .http2_only(configuration.http2) + .http2_only(configuration.http2_only) // .pool_max_idle_per_host(configuration.max_idle_connections_per_host as usize) // .http2_max_send_buf_size(128_000_000) // .timer(hyper_util::rt::TokioTimer::new()) @@ -110,7 +113,7 @@ impl CacheusServer let service_address = SocketAddr::from(([0, 0, 0, 0], server.configuration.listening_port)); info!( "Service listening on http://{}, http2:{}", - service_address, configuration.http2 + service_address, configuration.http2_only ); let listener = TcpListener::bind(service_address).await.unwrap(); @@ -122,7 +125,7 @@ impl CacheusServer // `hyper::rt` IO traits. let io = TokioIo::new(stream); let server = server.clone(); - if configuration.http2 { + if configuration.http2_only { tokio::task::spawn(async move { trace!("Listening for http2 connections..."); let server_for_metrics = server.clone(); @@ -203,14 +206,14 @@ impl CacheusServer async fn healthcheck(_req: Request) -> Result, hyper::Error> { - Ok(Response::new(BufferedBody::from_bytes(b"Healthy"))) + Ok(Response::new(BufferedBody::from_body(b"Healthy"))) } async fn prometheus( server: Arc, _: Request, ) -> Result, hyper::Error> { - Ok(Response::new(BufferedBody::from_bytes(&server.metrics.encode()))) + Ok(Response::new(BufferedBody::from_body(&server.metrics.encode()))) } async fn call_async( @@ -221,28 +224,29 @@ impl CacheusServer let timestamp = std::time::Instant::now(); - let (response, status) = CacheusServer::call_internal_async(service.clone(), request).await; + let response = CacheusServer::call_internal_async(service.clone(), request).await; - let elapsed = timestamp.elapsed(); - let status_str = status.to_string(); - service.metrics.request_duration.with_label_values(&[&status_str]).observe(elapsed.as_secs_f64()); + //let elapsed = timestamp.elapsed(); + //let status_str = status.to_string(); + // service.metrics.request_duration.with_label_values(&[&status_str]).observe(elapsed.as_secs_f64()); //service.metrics.cache_entries.set(service.cache.len() as f64); match response { Ok(ok) => { - service.metrics.requests.with_label_values(&[&status_str, ok.status().as_str()]).inc(); + //service.metrics.requests.with_label_values(&[&status_str, ok.status().as_str()]).inc(); Ok(ok) } Err(e) => { - service.metrics.requests.with_label_values(&[&status_str, "error"]).inc(); + //service.metrics.requests.with_label_values(&[&status_str, "error"]).inc(); Err(e) } } } - async fn call_internal_async(service: Arc, request: Request,) -> (Result, CacheusError>, Status) + async fn call_internal_async(service: Arc, request: Request,) -> Result, CacheusError> { let mut context = CallContext { variables: HashMap::new(), + stale_response: None, }; // Request buffering @@ -250,7 +254,7 @@ impl CacheusServer let buffered_body = BufferedBody::collect_buffered(body).await.unwrap(); let mut buffered_request = Request::from_parts(parts, buffered_body); - let mut response: Option<(Arc>, Status)> = None; + let mut response: Option> = None; let mut i = 0; // Evaluate middlewares until one returns a response @@ -262,25 +266,25 @@ impl CacheusServer } } - if response.is_none() { - return (Err(CacheusError { message: "No response generated by middlewares".to_string() }), Status::Ignore); - } - - let response = response.unwrap(); - let status = response.1; - let mut response = response.0.deref().clone(); - // In reverse order, apply middlewares on the response, starting from the one that produced the response for j in (0..i).rev() { service.configuration.middlewares[j].on_response(&mut context, service.clone(), &buffered_request, &mut response).await; } - return (Ok(response), status); + if response.is_none() { + return Err(CacheusError { message: "No response generated by middlewares".to_string() }); + } + + let response = response.unwrap(); + //let mut response = response.0.deref().clone(); + + return Ok(response); } } pub struct CallContext { variables: HashMap, + stale_response: Option>, } struct CacheusError