use serde::{Deserialize, Serialize};
use std::cell::Cell;
use wasm_bindgen::prelude::*;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum MemoryPressureLevel {
Normal,
Warning,
Critical,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct MemoryPressure {
pub level: MemoryPressureLevel,
pub used_bytes: usize,
pub total_bytes: usize,
pub usage_percent: f64,
}
impl MemoryPressure {
#[must_use]
pub fn current() -> Self {
Self::current_with_thresholds(80.0, 95.0)
}
#[must_use]
pub fn current_with_thresholds(warning_threshold: f64, critical_threshold: f64) -> Self {
let memory = wasm_bindgen::memory();
let buffer = memory
.unchecked_ref::<js_sys::WebAssembly::Memory>()
.buffer();
let array_buffer = buffer.unchecked_ref::<js_sys::ArrayBuffer>();
let total_bytes = array_buffer.byte_length() as usize;
let used_bytes = ALLOCATION_ESTIMATE.with(Cell::get);
#[allow(clippy::cast_precision_loss)]
let usage_percent = if total_bytes > 0 {
(used_bytes as f64 / total_bytes as f64) * 100.0
} else {
0.0
};
let level = if usage_percent >= critical_threshold {
MemoryPressureLevel::Critical
} else if usage_percent >= warning_threshold {
MemoryPressureLevel::Warning
} else {
MemoryPressureLevel::Normal
};
Self {
level,
used_bytes,
total_bytes,
usage_percent,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct MemoryConfig {
#[serde(default = "default_warning_threshold")]
pub warning_threshold: f64,
#[serde(default = "default_critical_threshold")]
pub critical_threshold: f64,
#[serde(default)]
pub auto_compact_on_warning: bool,
#[serde(default = "default_block_inserts")]
pub block_inserts_on_critical: bool,
}
fn default_warning_threshold() -> f64 {
80.0
}
fn default_critical_threshold() -> f64 {
95.0
}
fn default_block_inserts() -> bool {
true
}
impl Default for MemoryConfig {
fn default() -> Self {
Self {
warning_threshold: 80.0,
critical_threshold: 95.0,
auto_compact_on_warning: false,
block_inserts_on_critical: true,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct MemoryRecommendation {
pub action: String,
pub message: String,
pub can_insert: bool,
pub suggest_compact: bool,
}
thread_local! {
static ALLOCATION_ESTIMATE: Cell<usize> = const { Cell::new(0) };
}
#[inline]
pub fn track_allocation(bytes: usize) {
ALLOCATION_ESTIMATE.with(|e| {
e.set(e.get().saturating_add(bytes));
});
}
#[inline]
#[allow(dead_code)]
pub fn track_deallocation(bytes: usize) {
ALLOCATION_ESTIMATE.with(|e| {
e.set(e.get().saturating_sub(bytes));
});
}
#[must_use]
#[inline]
#[allow(dead_code)]
pub fn get_allocation_estimate() -> usize {
ALLOCATION_ESTIMATE.with(Cell::get)
}
#[inline]
pub fn track_vector_insert(dimensions: u32) {
let bytes = (dimensions as usize) * std::mem::size_of::<f32>();
track_allocation(bytes);
}
#[inline]
pub fn track_batch_insert(count: usize, dimensions: u32) {
let bytes_per_vector = (dimensions as usize) * std::mem::size_of::<f32>();
track_allocation(count * bytes_per_vector);
}
#[cfg(test)]
pub fn reset_allocation_estimate() {
ALLOCATION_ESTIMATE.with(|e| e.set(0));
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_memory_pressure_level_serialize() {
let normal = MemoryPressureLevel::Normal;
let json = serde_json::to_string(&normal).unwrap();
assert_eq!(json, "\"normal\"");
let warning = MemoryPressureLevel::Warning;
let json = serde_json::to_string(&warning).unwrap();
assert_eq!(json, "\"warning\"");
let critical = MemoryPressureLevel::Critical;
let json = serde_json::to_string(&critical).unwrap();
assert_eq!(json, "\"critical\"");
}
#[test]
fn test_memory_config_default() {
let config = MemoryConfig::default();
assert!((config.warning_threshold - 80.0).abs() < f64::EPSILON);
assert!((config.critical_threshold - 95.0).abs() < f64::EPSILON);
assert!(!config.auto_compact_on_warning);
assert!(config.block_inserts_on_critical);
}
#[test]
fn test_memory_config_deserialize() {
let json = r#"{
"warningThreshold": 70,
"criticalThreshold": 90,
"autoCompactOnWarning": true,
"blockInsertsOnCritical": false
}"#;
let config: MemoryConfig = serde_json::from_str(json).unwrap();
assert!((config.warning_threshold - 70.0).abs() < f64::EPSILON);
assert!((config.critical_threshold - 90.0).abs() < f64::EPSILON);
assert!(config.auto_compact_on_warning);
assert!(!config.block_inserts_on_critical);
}
#[test]
fn test_allocation_tracking() {
reset_allocation_estimate();
assert_eq!(get_allocation_estimate(), 0);
track_allocation(1000);
assert_eq!(get_allocation_estimate(), 1000);
track_allocation(500);
assert_eq!(get_allocation_estimate(), 1500);
track_deallocation(300);
assert_eq!(get_allocation_estimate(), 1200);
track_deallocation(2000);
assert_eq!(get_allocation_estimate(), 0);
}
#[test]
fn test_memory_recommendation_serialize() {
let rec = MemoryRecommendation {
action: "compact".to_string(),
message: "Memory usage high".to_string(),
can_insert: true,
suggest_compact: true,
};
let json = serde_json::to_string(&rec).unwrap();
assert!(json.contains("\"action\":\"compact\""));
assert!(json.contains("\"canInsert\":true"));
assert!(json.contains("\"suggestCompact\":true"));
}
}