use parking_lot::RwLock;
use std::alloc::{alloc, dealloc, Layout};
use std::sync::Arc;
const CACHE_LINE_SIZE: usize = 64;
const L1_CACHE_SIZE: usize = 32 * 1024;
const L2_CACHE_SIZE: usize = 256 * 1024;
const L3_CACHE_SIZE: usize = 8 * 1024 * 1024;
pub struct CacheHierarchy {
hot_storage: Arc<RwLock<HotStorage>>,
cold_storage: Arc<RwLock<ColdStorage>>,
access_tracker: Arc<RwLock<AccessTracker>>,
}
impl CacheHierarchy {
pub fn new(hot_capacity: usize, cold_capacity: usize) -> Self {
Self {
hot_storage: Arc::new(RwLock::new(HotStorage::new(hot_capacity))),
cold_storage: Arc::new(RwLock::new(ColdStorage::new(cold_capacity))),
access_tracker: Arc::new(RwLock::new(AccessTracker::new())),
}
}
pub fn get_node(&self, node_id: u64) -> Option<NodeData> {
self.access_tracker.write().record_access(node_id);
if let Some(data) = self.hot_storage.read().get(node_id) {
return Some(data);
}
if let Some(data) = self.cold_storage.read().get(node_id) {
if self.access_tracker.read().should_promote(node_id) {
self.promote_to_hot(node_id, data.clone());
}
return Some(data);
}
None
}
pub fn insert_node(&self, node_id: u64, data: NodeData) {
self.access_tracker.write().record_access(node_id);
if self.hot_storage.read().is_at_capacity() {
self.evict_one_to_cold(node_id); }
self.hot_storage.write().insert(node_id, data.clone());
}
fn promote_to_hot(&self, node_id: u64, data: NodeData) {
if self.hot_storage.read().is_full() {
self.evict_one_to_cold(node_id); }
self.hot_storage.write().insert(node_id, data);
self.cold_storage.write().remove(node_id);
}
fn evict_cold(&self) {
let tracker = self.access_tracker.read();
let lru_nodes = tracker.get_lru_nodes_by_frequency(10);
drop(tracker);
let mut hot = self.hot_storage.write();
let mut cold = self.cold_storage.write();
for node_id in lru_nodes {
if let Some(data) = hot.remove(node_id) {
cold.insert(node_id, data);
}
}
}
fn evict_one_to_cold(&self, protected_id: u64) {
let tracker = self.access_tracker.read();
let candidates = tracker.get_lru_nodes_by_frequency(5);
drop(tracker);
let mut hot = self.hot_storage.write();
let mut cold = self.cold_storage.write();
for node_id in candidates {
if node_id != protected_id {
if let Some(data) = hot.remove(node_id) {
cold.insert(node_id, data);
return;
}
}
}
}
pub fn prefetch_neighbors(&self, node_ids: &[u64]) {
for &node_id in node_ids {
#[cfg(target_arch = "x86_64")]
unsafe {
std::arch::x86_64::_mm_prefetch(
&node_id as *const u64 as *const i8,
std::arch::x86_64::_MM_HINT_T0,
);
}
}
}
}
#[repr(align(64))]
struct HotStorage {
entries: Vec<CacheLineEntry>,
capacity: usize,
size: usize,
}
impl HotStorage {
fn new(capacity: usize) -> Self {
Self {
entries: Vec::with_capacity(capacity),
capacity,
size: 0,
}
}
fn get(&self, node_id: u64) -> Option<NodeData> {
self.entries
.iter()
.find(|e| e.node_id == node_id)
.map(|e| e.data.clone())
}
fn insert(&mut self, node_id: u64, data: NodeData) {
self.entries.retain(|e| e.node_id != node_id);
if self.entries.len() >= self.capacity {
self.entries.remove(0); }
self.entries.push(CacheLineEntry { node_id, data });
self.size = self.entries.len();
}
fn remove(&mut self, node_id: u64) -> Option<NodeData> {
if let Some(pos) = self.entries.iter().position(|e| e.node_id == node_id) {
let entry = self.entries.remove(pos);
self.size = self.entries.len();
Some(entry.data)
} else {
None
}
}
fn is_full(&self) -> bool {
self.size >= self.capacity
}
fn is_at_capacity(&self) -> bool {
self.size >= self.capacity
}
}
#[repr(align(64))]
#[derive(Clone)]
struct CacheLineEntry {
node_id: u64,
data: NodeData,
}
struct ColdStorage {
entries: dashmap::DashMap<u64, Vec<u8>>,
capacity: usize,
}
impl ColdStorage {
fn new(capacity: usize) -> Self {
Self {
entries: dashmap::DashMap::new(),
capacity,
}
}
fn get(&self, node_id: u64) -> Option<NodeData> {
self.entries.get(&node_id).and_then(|compressed| {
bincode::decode_from_slice(&compressed, bincode::config::standard())
.ok()
.map(|(data, _)| data)
})
}
fn insert(&mut self, node_id: u64, data: NodeData) {
if let Ok(compressed) = bincode::encode_to_vec(&data, bincode::config::standard()) {
self.entries.insert(node_id, compressed);
}
}
fn remove(&mut self, node_id: u64) -> Option<NodeData> {
self.entries.remove(&node_id).and_then(|(_, compressed)| {
bincode::decode_from_slice(&compressed, bincode::config::standard())
.ok()
.map(|(data, _)| data)
})
}
}
struct AccessTracker {
access_counts: dashmap::DashMap<u64, u32>,
last_access: dashmap::DashMap<u64, u64>,
timestamp: u64,
}
impl AccessTracker {
fn new() -> Self {
Self {
access_counts: dashmap::DashMap::new(),
last_access: dashmap::DashMap::new(),
timestamp: 0,
}
}
fn record_access(&mut self, node_id: u64) {
self.timestamp += 1;
self.access_counts
.entry(node_id)
.and_modify(|count| *count += 1)
.or_insert(1);
self.last_access.insert(node_id, self.timestamp);
}
fn should_promote(&self, node_id: u64) -> bool {
self.access_counts
.get(&node_id)
.map(|count| *count > 5)
.unwrap_or(false)
}
fn get_lru_nodes(&self, count: usize) -> Vec<u64> {
let mut nodes: Vec<_> = self
.last_access
.iter()
.map(|entry| (*entry.key(), *entry.value()))
.collect();
nodes.sort_by_key(|(_, timestamp)| *timestamp);
nodes
.into_iter()
.take(count)
.map(|(node_id, _)| node_id)
.collect()
}
fn get_lru_nodes_by_frequency(&self, count: usize) -> Vec<u64> {
let mut nodes: Vec<_> = self
.access_counts
.iter()
.map(|entry| (*entry.key(), *entry.value()))
.collect();
nodes.sort_by_key(|(_, access_count)| *access_count);
nodes
.into_iter()
.take(count)
.map(|(node_id, _)| node_id)
.collect()
}
}
#[derive(Clone, serde::Serialize, serde::Deserialize, bincode::Encode, bincode::Decode)]
pub struct NodeData {
pub id: u64,
pub labels: Vec<String>,
pub properties: Vec<(String, CachePropertyValue)>,
}
#[derive(Clone, serde::Serialize, serde::Deserialize, bincode::Encode, bincode::Decode)]
pub enum CachePropertyValue {
String(String),
Integer(i64),
Float(f64),
Boolean(bool),
}
pub struct HotColdStorage {
cache_hierarchy: CacheHierarchy,
}
impl HotColdStorage {
pub fn new() -> Self {
Self {
cache_hierarchy: CacheHierarchy::new(1000, 10000),
}
}
pub fn get(&self, node_id: u64) -> Option<NodeData> {
self.cache_hierarchy.get_node(node_id)
}
pub fn insert(&self, node_id: u64, data: NodeData) {
self.cache_hierarchy.insert_node(node_id, data);
}
pub fn prefetch(&self, node_ids: &[u64]) {
self.cache_hierarchy.prefetch_neighbors(node_ids);
}
}
impl Default for HotColdStorage {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_cache_hierarchy() {
let cache = CacheHierarchy::new(10, 100);
let data = NodeData {
id: 1,
labels: vec!["Person".to_string()],
properties: vec![(
"name".to_string(),
CachePropertyValue::String("Alice".to_string()),
)],
};
cache.insert_node(1, data.clone());
let retrieved = cache.get_node(1);
assert!(retrieved.is_some());
}
#[test]
fn test_hot_cold_promotion() {
let cache = CacheHierarchy::new(2, 10);
for i in 1..=3 {
cache.insert_node(
i,
NodeData {
id: i,
labels: vec![],
properties: vec![],
},
);
}
for _ in 0..10 {
cache.get_node(1);
}
assert!(cache.get_node(1).is_some());
}
}