use crate::error::Result;
use crate::value::Value;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::fmt;
#[derive(Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct CacheKey(pub [u8; 32]);
impl CacheKey {
pub fn from_parts(parts: &[&[u8]]) -> Self {
let mut hasher = Sha256::new();
for part in parts {
hasher.update((part.len() as u64).to_le_bytes());
hasher.update(part);
}
Self(hasher.finalize().into())
}
pub fn for_state(config_hash: &CacheKey, data_hash: &CacheKey) -> Self {
Self::from_parts(&[&config_hash.0, &data_hash.0])
}
pub fn for_output(
config_hash: &CacheKey,
state_hash: &CacheKey,
input_hash: &CacheKey,
) -> Self {
Self::from_parts(&[&config_hash.0, &state_hash.0, &input_hash.0])
}
pub fn hash_data(data: &[u8]) -> Self {
Self::from_parts(&[data])
}
pub fn to_hex(&self) -> String {
self.0.iter().map(|b| format!("{b:02x}")).collect()
}
}
impl fmt::Debug for CacheKey {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "CacheKey({}...)", &self.to_hex()[..12])
}
}
impl fmt::Display for CacheKey {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", &self.to_hex()[..16])
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum CacheTier {
Memory,
Local,
Remote,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum Origin {
Computed {
node_id: String,
run_id: String,
},
Ingested {
source: String,
},
Streamed {
window_start: DateTime<Utc>,
window_end: DateTime<Utc>,
},
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EntryMeta {
pub key: CacheKey,
pub size_bytes: u64,
pub created_at: DateTime<Utc>,
pub last_accessed: DateTime<Utc>,
pub ttl: Option<std::time::Duration>,
pub origin: Origin,
}
pub trait CacheStore: Send + Sync {
fn get(&self, key: &CacheKey) -> Result<Option<Value>>;
fn put(&self, key: &CacheKey, value: &Value) -> Result<()>;
fn exists(&self, key: &CacheKey) -> Result<bool>;
fn remove(&self, key: &CacheKey) -> Result<()>;
fn metadata(&self, key: &CacheKey) -> Result<Option<EntryMeta>>;
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cache_key_deterministic() {
let k1 = CacheKey::from_parts(&[b"hello", b"world"]);
let k2 = CacheKey::from_parts(&[b"hello", b"world"]);
assert_eq!(k1, k2);
}
#[test]
fn cache_key_sensitive_to_content() {
let k1 = CacheKey::from_parts(&[b"hello", b"world"]);
let k2 = CacheKey::from_parts(&[b"hello", b"world!"]);
assert_ne!(k1, k2);
}
#[test]
fn cache_key_sensitive_to_part_boundaries() {
let k1 = CacheKey::from_parts(&[b"ab", b"c"]);
let k2 = CacheKey::from_parts(&[b"a", b"bc"]);
assert_ne!(k1, k2);
}
#[test]
fn cache_key_for_state() {
let config = CacheKey::hash_data(b"scaler_config");
let data = CacheKey::hash_data(b"training_data");
let state_key = CacheKey::for_state(&config, &data);
let state_key2 = CacheKey::for_state(&config, &data);
assert_eq!(state_key, state_key2);
let data2 = CacheKey::hash_data(b"different_data");
let state_key3 = CacheKey::for_state(&config, &data2);
assert_ne!(state_key, state_key3);
}
#[test]
fn cache_key_for_output() {
let config = CacheKey::hash_data(b"config");
let state = CacheKey::hash_data(b"state");
let input = CacheKey::hash_data(b"input");
let key = CacheKey::for_output(&config, &state, &input);
let state2 = CacheKey::hash_data(b"state2");
let key2 = CacheKey::for_output(&config, &state2, &input);
assert_ne!(key, key2);
}
#[test]
fn cache_key_hex_and_display() {
let key = CacheKey::hash_data(b"test");
let hex = key.to_hex();
assert_eq!(hex.len(), 64);
let display = format!("{key}");
assert_eq!(display.len(), 16);
let debug = format!("{key:?}");
assert!(debug.starts_with("CacheKey("));
}
#[test]
fn cache_key_serde_roundtrip() {
let key = CacheKey::hash_data(b"test_data");
let json = serde_json::to_string(&key).unwrap();
let deserialized: CacheKey = serde_json::from_str(&json).unwrap();
assert_eq!(key, deserialized);
}
}