From e7c01ef5396930d99ea2305ae39caa06eed76047 Mon Sep 17 00:00:00 2001
From: Christoph Herzog <chris@theduke.at>
Date: Sat, 11 Jan 2025 09:30:44 +0100
Subject: [PATCH] feat(cache): Allow using in-memory compression dicts

This PR allows setting the compression dictionary used for headers in
`pingora-cache` from in-memory bytes instead of an external file.

This is useful if the dict is already embedded in the binary, or is
retrieved from remote storage dynamically.

Also cleans up the code a bit by making the constant private, moving the
initialisers to the same file and re-exporting them to keep the logic
contained, and by slightly simplifying the file loading code.
---
 pingora-cache/src/lib.rs  | 10 +----
 pingora-cache/src/meta.rs | 89 +++++++++++++++++++--------------------
 2 files changed, 45 insertions(+), 54 deletions(-)

diff --git a/pingora-cache/src/lib.rs b/pingora-cache/src/lib.rs
index bc370b7f3..0464297d0 100644
--- a/pingora-cache/src/lib.rs
+++ b/pingora-cache/src/lib.rs
@@ -46,6 +46,7 @@ use crate::max_file_size::MaxFileSizeMissHandler;
 pub use key::CacheKey;
 use lock::{CacheLock, LockStatus, Locked};
 pub use memory::MemCache;
+pub use meta::{set_compression_dict_content, set_compression_dict_path};
 pub use meta::{CacheMeta, CacheMetaDefaults};
 pub use storage::{HitHandler, MissHandler, PurgeType, Storage};
 pub use variance::VarianceBuilder;
@@ -1237,12 +1238,3 @@ impl HttpCache {
             .set_tag(|| Tag::new("is_subrequest", true))
     }
 }
-
-/// Set the header compression dictionary, which helps serialize http header.
-///
-/// Return false if it is already set.
-pub fn set_compression_dict_path(path: &str) -> bool {
-    crate::meta::COMPRESSION_DICT_PATH
-        .set(path.to_string())
-        .is_ok()
-}
diff --git a/pingora-cache/src/meta.rs b/pingora-cache/src/meta.rs
index fd608d4e1..75a1a0dda 100644
--- a/pingora-cache/src/meta.rs
+++ b/pingora-cache/src/meta.rs
@@ -15,9 +15,13 @@
 //! Metadata for caching
 
 pub use http::Extensions;
+use log::warn;
+use once_cell::sync::{Lazy, OnceCell};
 use pingora_error::{Error, ErrorType::*, OrErr, Result};
+use pingora_header_serde::HeaderSerde;
 use pingora_http::{HMap, ResponseHeader};
 use serde::{Deserialize, Serialize};
+use std::borrow::Cow;
 use std::time::{Duration, SystemTime};
 
 use crate::key::HashBinary;
@@ -561,53 +565,20 @@ impl CacheMetaDefaults {
     }
 }
 
-use log::warn;
-use once_cell::sync::{Lazy, OnceCell};
-use pingora_header_serde::HeaderSerde;
-use std::fs::File;
-use std::io::Read;
-
-/* load header compression engine and its dictionary globally */
-pub(crate) static COMPRESSION_DICT_PATH: OnceCell<String> = OnceCell::new();
-
-fn load_file(path: &String) -> Option<Vec<u8>> {
-    let mut file = File::open(path)
-        .map_err(|e| {
-            warn!(
-                "failed to open header compress dictionary file at {}, {:?}",
-                path, e
-            );
-            e
-        })
-        .ok()?;
-    let mut dict = Vec::new();
-    file.read_to_end(&mut dict)
-        .map_err(|e| {
-            warn!(
-                "failed to read header compress dictionary file at {}, {:?}",
-                path, e
-            );
-            e
-        })
-        .ok()?;
-
-    Some(dict)
-}
+/// The dictionary content for header compression.
+///
+/// Used during initialization of [`HEADER_SERDE`].
+static COMPRESSION_DICT_CONTENT: OnceCell<Cow<'static, [u8]>> = OnceCell::new();
 
 static HEADER_SERDE: Lazy<HeaderSerde> = Lazy::new(|| {
-    let dict_path_opt = COMPRESSION_DICT_PATH.get();
-
-    if dict_path_opt.is_none() {
-        warn!("COMPRESSION_DICT_PATH is not set");
-    }
-
-    let result = dict_path_opt.and_then(load_file);
-
-    if result.is_none() {
-        warn!("HeaderSerde not loaded from file");
-    }
-
-    HeaderSerde::new(result)
+    let dict_opt = if let Some(dict_content) = COMPRESSION_DICT_CONTENT.get() {
+        Some(dict_content.to_vec())
+    } else {
+        warn!("no header compression dictionary loaded - use set_compression_dict_content() or set_compression_dict_path() to set one");
+        None
+    };
+
+    HeaderSerde::new(dict_opt)
 });
 
 pub(crate) fn header_serialize(header: &ResponseHeader) -> Result<Vec<u8>> {
@@ -617,3 +588,31 @@ pub(crate) fn header_serialize(header: &ResponseHeader) -> Result<Vec<u8>> {
 pub(crate) fn header_deserialize<T: AsRef<[u8]>>(buf: T) -> Result<ResponseHeader> {
     HEADER_SERDE.deserialize(buf.as_ref())
 }
+
+/// Load the header compression dictionary from a file, which helps serialize http header.
+///
+/// Returns false if it is already set or if the file could not be read.
+///
+/// Use [`set_compression_dict_content`] to set the dictionary from memory instead.
+pub fn set_compression_dict_path(path: &str) -> bool {
+    match std::fs::read(path) {
+        Ok(dict) => COMPRESSION_DICT_CONTENT.set(dict.into()).is_ok(),
+        Err(e) => {
+            warn!(
+                "failed to read header compress dictionary file at {}, {:?}",
+                path, e
+            );
+            false
+        }
+    }
+}
+
+/// Set the header compression dictionary content, which helps serialize http header.
+///
+/// Returns false if it is already set.
+///
+/// This is an alernative to [`set_compression_dict_path`], allowing use of
+/// a dictionary without an external file.
+pub fn set_compression_dict_content(content: Cow<'static, [u8]>) -> bool {
+    COMPRESSION_DICT_CONTENT.set(content).is_ok()
+}