hedl_mcp/resource_limits/
memory_cache.rs

1// Dweve HEDL - Hierarchical Entity Data Language
2//
3// Copyright (c) 2025 Dweve IP B.V. and individual contributors.
4//
5// SPDX-License-Identifier: Apache-2.0
6//
7// Licensed under the Apache License, Version 2.0 (the "License");
8// you may not use this file except in compliance with the License.
9// You may obtain a copy of the License in the LICENSE file at the
10// root of this repository or at: http://www.apache.org/licenses/LICENSE-2.0
11//
12// Unless required by applicable law or agreed to in writing, software
13// distributed under the License is distributed on an "AS IS" BASIS,
14// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15// See the License for the specific language governing permissions and
16// limitations under the License.
17
18//! Memory-aware cache with usage tracking.
19
20use super::error::ResourceLimitError;
21use dashmap::DashMap;
22use serde_json::Value;
23use std::sync::atomic::{AtomicUsize, Ordering};
24use tracing::debug;
25
26/// Memory-aware cache that tracks actual memory usage.
27///
28/// Unlike the basic cache which only tracks entry count, this estimates
29/// and enforces memory limits to prevent unbounded growth.
30#[derive(Debug)]
31pub struct MemoryAwareCache {
32    /// Entry size tracking (key -> size in bytes).
33    entry_sizes: DashMap<String, usize>,
34
35    /// Total memory usage in bytes.
36    total_size: AtomicUsize,
37
38    /// Maximum memory budget in bytes.
39    max_size: usize,
40}
41
42impl MemoryAwareCache {
43    /// Create a new memory-aware cache.
44    ///
45    /// # Arguments
46    ///
47    /// * `max_size` - Maximum memory budget in bytes
48    #[must_use]
49    pub fn new(max_size: usize) -> Self {
50        Self {
51            entry_sizes: DashMap::new(),
52            total_size: AtomicUsize::new(0),
53            max_size,
54        }
55    }
56
57    /// Insert a value with memory tracking.
58    ///
59    /// # Arguments
60    ///
61    /// * `key` - Cache key
62    /// * `value` - JSON value to cache
63    ///
64    /// # Returns
65    ///
66    /// `Ok(())` if inserted, `Err` if would exceed memory limit.
67    pub fn insert(&self, key: String, value: Value) -> Result<(), ResourceLimitError> {
68        let value_size = estimate_json_size(&value);
69
70        // Check memory limit
71        let current = self.total_size.load(Ordering::Relaxed);
72        let new_total = current.saturating_add(value_size);
73
74        if new_total > self.max_size {
75            return Err(ResourceLimitError::CacheMemoryExceeded {
76                current,
77                limit: self.max_size,
78                needed: value_size,
79            });
80        }
81
82        // Track size
83        self.entry_sizes.insert(key.clone(), value_size);
84        self.total_size.fetch_add(value_size, Ordering::Relaxed);
85
86        debug!(
87            "Cache insert: key={}, size={}, total={}",
88            key, value_size, new_total
89        );
90
91        Ok(())
92    }
93
94    /// Remove an entry and update memory tracking.
95    ///
96    /// # Arguments
97    ///
98    /// * `key` - Cache key to remove
99    pub fn remove(&self, key: &str) {
100        if let Some((_, size)) = self.entry_sizes.remove(key) {
101            self.total_size.fetch_sub(size, Ordering::Relaxed);
102            debug!("Cache remove: key={}, size={}", key, size);
103        }
104    }
105
106    /// Get current memory usage in bytes.
107    pub fn current_usage(&self) -> usize {
108        self.total_size.load(Ordering::Relaxed)
109    }
110
111    /// Get maximum memory budget in bytes.
112    pub fn max_size(&self) -> usize {
113        self.max_size
114    }
115
116    /// Get the number of cached entries.
117    pub fn entry_count(&self) -> usize {
118        self.entry_sizes.len()
119    }
120
121    /// Clear all entries and reset memory tracking.
122    pub fn clear(&self) {
123        self.entry_sizes.clear();
124        self.total_size.store(0, Ordering::Relaxed);
125    }
126}
127
128/// Estimate the memory size of a JSON value.
129///
130/// Provides a rough estimate of memory usage for cache entries.
131/// This is an approximation, not an exact measurement.
132pub(crate) fn estimate_json_size(value: &Value) -> usize {
133    match value {
134        Value::Null => 8,
135        Value::Bool(_) => 1,
136        Value::Number(_) => 8,
137        Value::String(s) => s.len() + 24, // String overhead
138        Value::Array(arr) => {
139            24 + arr.iter().map(estimate_json_size).sum::<usize>() // Array overhead
140        }
141        Value::Object(obj) => {
142            24 + obj
143                .iter()
144                .map(|(k, v)| k.len() + estimate_json_size(v))
145                .sum::<usize>() // Object overhead
146        }
147    }
148}
hedl_mcp/resource_limits/memory_cache.rs

hedl_mcp/resource_limits/
memory_cache.rs
