use crate::types::{GhostNode, ShardId};
use phago_core::types::NodeId;
use std::collections::HashMap;
pub struct GhostNodeCache {
cache: HashMap<NodeId, GhostNode>,
max_size: usize,
access_order: Vec<NodeId>,
}
impl GhostNodeCache {
pub fn new(max_size: usize) -> Self {
Self {
cache: HashMap::with_capacity(max_size),
max_size,
access_order: Vec::with_capacity(max_size),
}
}
pub fn get(&mut self, id: &NodeId) -> Option<&GhostNode> {
if self.cache.contains_key(id) {
self.access_order.retain(|x| x != id);
self.access_order.push(*id);
self.cache.get(id)
} else {
None
}
}
pub fn peek(&self, id: &NodeId) -> Option<&GhostNode> {
self.cache.get(id)
}
pub fn insert(&mut self, ghost: GhostNode) {
let id = ghost.node_id;
if self.cache.contains_key(&id) {
self.cache.insert(id, ghost);
self.access_order.retain(|x| *x != id);
self.access_order.push(id);
return;
}
while self.cache.len() >= self.max_size && !self.access_order.is_empty() {
let oldest = self.access_order.remove(0);
self.cache.remove(&oldest);
}
self.cache.insert(id, ghost);
self.access_order.push(id);
}
pub fn update_full_data(&mut self, id: &NodeId, data: phago_core::types::NodeData) {
if let Some(ghost) = self.cache.get_mut(id) {
ghost.full_data = Some(data);
}
}
pub fn contains(&self, id: &NodeId) -> bool {
self.cache.contains_key(id)
}
pub fn nodes_from_shard(&self, shard_id: ShardId) -> Vec<&GhostNode> {
self.cache
.values()
.filter(|g| g.shard_id == shard_id)
.collect()
}
pub fn all_nodes(&self) -> Vec<&GhostNode> {
self.cache.values().collect()
}
pub fn remove(&mut self, id: &NodeId) -> Option<GhostNode> {
self.access_order.retain(|x| x != id);
self.cache.remove(id)
}
pub fn clear(&mut self) {
self.cache.clear();
self.access_order.clear();
}
pub fn len(&self) -> usize {
self.cache.len()
}
pub fn is_empty(&self) -> bool {
self.cache.is_empty()
}
pub fn capacity(&self) -> usize {
self.max_size
}
pub fn invalidate_shard(&mut self, shard_id: ShardId) -> usize {
let to_remove: Vec<NodeId> = self
.cache
.iter()
.filter(|(_, g)| g.shard_id == shard_id)
.map(|(id, _)| *id)
.collect();
let count = to_remove.len();
for id in to_remove {
self.cache.remove(&id);
self.access_order.retain(|x| *x != id);
}
count
}
pub fn stats(&self) -> GhostCacheStats {
let mut nodes_by_shard: HashMap<ShardId, usize> = HashMap::new();
let mut with_full_data = 0;
for ghost in self.cache.values() {
*nodes_by_shard.entry(ghost.shard_id).or_insert(0) += 1;
if ghost.full_data.is_some() {
with_full_data += 1;
}
}
GhostCacheStats {
total_nodes: self.cache.len(),
max_capacity: self.max_size,
nodes_by_shard,
nodes_with_full_data: with_full_data,
}
}
}
#[derive(Debug, Clone)]
pub struct GhostCacheStats {
pub total_nodes: usize,
pub max_capacity: usize,
pub nodes_by_shard: HashMap<ShardId, usize>,
pub nodes_with_full_data: usize,
}
#[cfg(test)]
mod tests {
use super::*;
fn make_ghost(id: u64, shard: u32) -> GhostNode {
GhostNode::new(
NodeId::from_seed(id),
ShardId::new(shard),
format!("node_{}", id),
)
}
#[test]
fn test_insert_and_get() {
let mut cache = GhostNodeCache::new(10);
let ghost = make_ghost(1, 0);
let id = ghost.node_id;
cache.insert(ghost);
assert!(cache.contains(&id));
assert_eq!(cache.len(), 1);
let retrieved = cache.get(&id).unwrap();
assert_eq!(retrieved.label, "node_1");
}
#[test]
fn test_lru_eviction() {
let mut cache = GhostNodeCache::new(3);
cache.insert(make_ghost(1, 0));
cache.insert(make_ghost(2, 0));
cache.insert(make_ghost(3, 0));
let _ = cache.get(&NodeId::from_seed(1));
cache.insert(make_ghost(4, 0));
assert!(cache.contains(&NodeId::from_seed(1)));
assert!(!cache.contains(&NodeId::from_seed(2))); assert!(cache.contains(&NodeId::from_seed(3)));
assert!(cache.contains(&NodeId::from_seed(4)));
}
#[test]
fn test_nodes_from_shard() {
let mut cache = GhostNodeCache::new(10);
cache.insert(make_ghost(1, 0));
cache.insert(make_ghost(2, 1));
cache.insert(make_ghost(3, 0));
cache.insert(make_ghost(4, 2));
let shard0_nodes = cache.nodes_from_shard(ShardId::new(0));
assert_eq!(shard0_nodes.len(), 2);
let shard1_nodes = cache.nodes_from_shard(ShardId::new(1));
assert_eq!(shard1_nodes.len(), 1);
}
#[test]
fn test_invalidate_shard() {
let mut cache = GhostNodeCache::new(10);
cache.insert(make_ghost(1, 0));
cache.insert(make_ghost(2, 1));
cache.insert(make_ghost(3, 0));
let count = cache.invalidate_shard(ShardId::new(0));
assert_eq!(count, 2);
assert_eq!(cache.len(), 1);
assert!(cache.contains(&NodeId::from_seed(2)));
}
#[test]
fn test_clear() {
let mut cache = GhostNodeCache::new(10);
cache.insert(make_ghost(1, 0));
cache.insert(make_ghost(2, 0));
cache.clear();
assert!(cache.is_empty());
assert_eq!(cache.len(), 0);
}
#[test]
fn test_stats() {
let mut cache = GhostNodeCache::new(10);
cache.insert(make_ghost(1, 0));
cache.insert(make_ghost(2, 1));
cache.insert(make_ghost(3, 0));
let stats = cache.stats();
assert_eq!(stats.total_nodes, 3);
assert_eq!(stats.max_capacity, 10);
assert_eq!(*stats.nodes_by_shard.get(&ShardId::new(0)).unwrap(), 2);
assert_eq!(*stats.nodes_by_shard.get(&ShardId::new(1)).unwrap(), 1);
}
}