Weekly Sync 2024-06-07

.bleep

@@ -1 +1 @@
-946e1733bc009bf3e73c635ad1b8bac0079dc680
+a73e00f51bc6643e93abedc2a763d38122ffb21d

pingora-cache/src/key.rs

@@ -98,6 +98,7 @@ pub struct CacheKey {
     // All strings for now. It can be more structural as long as it can hash
     namespace: String,
     primary: String,
+    primary_bin_override: Option<HashBinary>,
     variance: Option<HashBinary>,
     /// An extra tag for identifying users
     ///
@@ -120,6 +121,11 @@ impl CacheKey {
     pub fn remove_variance_key(&mut self) {
         self.variance = None
     }
+
+    /// Override the primary key hash
+    pub fn set_primary_bin_override(&mut self, key: HashBinary) {
+        self.primary_bin_override = Some(key)
+    }
 }
 
 /// Storage optimized cache key to keep in memory or in storage
@@ -186,6 +192,7 @@ impl CacheKey {
         CacheKey {
             namespace: "".into(),
             primary: format!("{}", req_header.uri),
+            primary_bin_override: None,
             variance: None,
             user_tag: "".into(),
         }
@@ -203,6 +210,7 @@ impl CacheKey {
         CacheKey {
             namespace: namespace.into(),
             primary: primary.into(),
+            primary_bin_override: None,
             variance: None,
             user_tag: user_tag.into(),
         }
@@ -231,7 +239,11 @@ impl CacheKey {
 
 impl CacheHashKey for CacheKey {
     fn primary_bin(&self) -> HashBinary {
-        self.primary_hasher().finalize().into()
+        if let Some(primary_bin_override) = self.primary_bin_override {
+            primary_bin_override
+        } else {
+            self.primary_hasher().finalize().into()
+        }
     }
 
     fn variance_bin(&self) -> Option<HashBinary> {
@@ -252,6 +264,7 @@ mod tests {
         let key = CacheKey {
             namespace: "".into(),
             primary: "aa".into(),
+            primary_bin_override: None,
             variance: None,
             user_tag: "1".into(),
         };
@@ -265,11 +278,64 @@ mod tests {
         assert_eq!(compact.combined(), hash);
     }
 
+    #[test]
+    fn test_cache_key_hash_override() {
+        let mut key = CacheKey {
+            namespace: "".into(),
+            primary: "aa".into(),
+            primary_bin_override: str2hex("27c35e6e9373877f29e562464e46497e"),
+            variance: None,
+            user_tag: "1".into(),
+        };
+        let hash = key.primary();
+        assert_eq!(hash, "27c35e6e9373877f29e562464e46497e");
+        assert!(key.variance().is_none());
+        assert_eq!(key.combined(), hash);
+        let compact = key.to_compact();
+        assert_eq!(compact.primary(), hash);
+        assert!(compact.variance().is_none());
+        assert_eq!(compact.combined(), hash);
+
+        // make sure set_primary_bin_override overrides the primary key hash correctly
+        key.set_primary_bin_override(str2hex("004174d3e75a811a5b44c46b3856f3ee").unwrap());
+        let hash = key.primary();
+        assert_eq!(hash, "004174d3e75a811a5b44c46b3856f3ee");
+        assert!(key.variance().is_none());
+        assert_eq!(key.combined(), hash);
+        let compact = key.to_compact();
+        assert_eq!(compact.primary(), hash);
+        assert!(compact.variance().is_none());
+        assert_eq!(compact.combined(), hash);
+    }
+
     #[test]
     fn test_cache_key_vary_hash() {
         let key = CacheKey {
             namespace: "".into(),
             primary: "aa".into(),
+            primary_bin_override: None,
+            variance: Some([0u8; 16]),
+            user_tag: "1".into(),
+        };
+        let hash = key.primary();
+        assert_eq!(hash, "ac10f2aef117729f8dad056b3059eb7e");
+        assert_eq!(key.variance().unwrap(), "00000000000000000000000000000000");
+        assert_eq!(key.combined(), "004174d3e75a811a5b44c46b3856f3ee");
+        let compact = key.to_compact();
+        assert_eq!(compact.primary(), "ac10f2aef117729f8dad056b3059eb7e");
+        assert_eq!(
+            compact.variance().unwrap(),
+            "00000000000000000000000000000000"
+        );
+        assert_eq!(compact.combined(), "004174d3e75a811a5b44c46b3856f3ee");
+    }
+
+    #[test]
+    fn test_cache_key_vary_hash_override() {
+        let key = CacheKey {
+            namespace: "".into(),
+            primary: "saaaad".into(),
+            primary_bin_override: str2hex("ac10f2aef117729f8dad056b3059eb7e"),
             variance: Some([0u8; 16]),
             user_tag: "1".into(),
         };

pingora-core/src/apps/http_app.rs

@@ -89,7 +89,7 @@ where
         }
         if !body.is_empty() {
             // TODO: check if chunked encoding is needed
-            match http.write_response_body(body.into()).await {
+            match http.write_response_body(body.into(), true).await {
                 Ok(_) => debug!("HTTP response written."),
                 Err(e) => error!(
                     "HTTP server fails to write to downstream: {e}, {}",

pingora-core/src/modules/http/compression.rs

@@ -16,10 +16,25 @@
 
 use super::*;
 use crate::protocols::http::compression::ResponseCompressionCtx;
+use std::ops::{Deref, DerefMut};
 
 /// HTTP response compression module
 pub struct ResponseCompression(ResponseCompressionCtx);
 
+impl Deref for ResponseCompression {
+    type Target = ResponseCompressionCtx;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for ResponseCompression {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
 #[async_trait]
 impl HttpModule for ResponseCompression {
     fn as_any(&self) -> &dyn std::any::Any {
@@ -52,7 +67,9 @@ impl HttpModule for ResponseCompression {
             return Ok(());
         }
         let compressed = self.0.response_body_filter(body.as_ref(), end_of_stream);
-        *body = compressed;
+        if compressed.is_some() {
+            *body = compressed;
+        }
         Ok(())
     }
 }

pingora-core/src/protocols/http/compression/mod.rs

@@ -200,6 +200,8 @@ impl ResponseCompressionCtx {
     }
 
     /// Stream the response body chunks into this ctx. The return value will be the compressed data
+    ///
+    /// Return None if the compressed is not enabled
     pub fn response_body_filter(&mut self, data: Option<&Bytes>, end: bool) -> Option<Bytes> {
         match &mut self.0 {
             CtxInner::HeaderPhase {

pingora-core/src/protocols/http/server.rs

@@ -131,13 +131,19 @@ impl Session {
     }
 
     /// Write the response body to client
-    pub async fn write_response_body(&mut self, data: Bytes) -> Result<()> {
+    pub async fn write_response_body(&mut self, data: Bytes, end: bool) -> Result<()> {
+        if data.is_empty() && !end {
+            // writing 0 byte to a chunked encoding h1 would finish the stream
+            // writing 0 bytes to h2 is noop
+            // we don't want to actually write in either cases
+            return Ok(());
+        }
         match self {
             Self::H1(s) => {
                 s.write_body(&data).await?;
                 Ok(())
             }
-            Self::H2(s) => s.write_body(data, false),
+            Self::H2(s) => s.write_body(data, end),
         }
     }
 
@@ -236,14 +242,18 @@ impl Session {
         }
     }
 
-    /// Send error response to client
-    pub async fn respond_error(&mut self, error: u16) {
-        let resp = match error {
+    pub fn generate_error(error: u16) -> ResponseHeader {
+        match error {
             /* common error responses are pre-generated */
             502 => error_resp::HTTP_502_RESPONSE.clone(),
             400 => error_resp::HTTP_400_RESPONSE.clone(),
             _ => error_resp::gen_error_response(error),
-        };
+        }
+    }
+
+    /// Send error response to client
+    pub async fn respond_error(&mut self, error: u16) {
+        let resp = Self::generate_error(error);
 
         // TODO: we shouldn't be closing downstream connections on internally generated errors
         // and possibly other upstream connect() errors (connection refused, timeout, etc)

pingora-core/src/protocols/http/v1/client.rs

@@ -427,7 +427,16 @@ impl HttpSession {
     }
 
     fn is_connection_keepalive(&self) -> Option<bool> {
-        is_buf_keepalive(self.get_header(header::CONNECTION).map(|v| v.as_bytes()))
+        let request_keepalive = self
+            .request_written
+            .as_ref()
+            .and_then(|req| is_buf_keepalive(req.headers.get(header::CONNECTION)));
+
+        match request_keepalive {
+            // ignore what the server sends if request disables keepalive explicitly
+            Some(false) => Some(false),
+            _ => is_buf_keepalive(self.get_header(header::CONNECTION)),
+        }
     }
 
     /// `Keep-Alive: timeout=5, max=1000` => 5, 1000

pingora-core/src/protocols/http/v1/common.rs

@@ -14,7 +14,7 @@
 
 //! Common functions and constants
 
-use http::header;
+use http::{header, HeaderValue};
 use log::warn;
 use pingora_http::{HMap, RequestHeader, ResponseHeader};
 use std::str;
@@ -202,9 +202,9 @@ pub(super) fn buf_to_content_length(header_value: Option<&[u8]>) -> Option<usize
 }
 
 #[inline]
-pub(super) fn is_buf_keepalive(header_value: Option<&[u8]>) -> Option<bool> {
+pub(super) fn is_buf_keepalive(header_value: Option<&HeaderValue>) -> Option<bool> {
     header_value.and_then(|value| {
-        let value = parse_connection_header(value);
+        let value = parse_connection_header(value.as_bytes());
         if value.keep_alive {
             Some(true)
         } else if value.close {

pingora-core/src/protocols/http/v1/server.rs

@@ -508,7 +508,7 @@ impl HttpSession {
     }
 
     fn is_connection_keepalive(&self) -> Option<bool> {
-        is_buf_keepalive(self.get_header(header::CONNECTION).map(|v| v.as_bytes()))
+        is_buf_keepalive(self.get_header(header::CONNECTION))
     }
 
     /// Apply keepalive settings according to the client

pingora-http/src/lib.rs

@@ -573,10 +573,13 @@ fn remove_header<'a, T, N: ?Sized>(
 where
     &'a N: 'a + AsHeaderName,
 {
-    if let Some(name_map) = name_map {
-        name_map.remove(name);
+    let removed = value_map.remove(name);
+    if removed.is_some() {
+        if let Some(name_map) = name_map {
+            name_map.remove(name);
+        }
     }
-    value_map.remove(name)
+    removed
 }
 
 #[inline]