Skip to main content

ipfrs_storage/
storage_quota_manager.rs

1//! Per-namespace storage quota enforcement with usage tracking and eviction triggers.
2//!
3//! `StorageQuotaManager` tracks allocated objects per namespace, enforces configurable
4//! hard/soft byte and object-count limits, and provides multiple eviction strategies
5//! (Oldest, LRU, LFU, SizeDescending) for reclaiming space.
6
7use std::collections::HashMap;
8
9// ---------------------------------------------------------------------------
10// Public types
11// ---------------------------------------------------------------------------
12
13/// Newtype wrapper for namespace identifiers (e.g. user ID, app ID).
14#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
15pub struct QuotaNamespace(pub String);
16
17impl QuotaNamespace {
18    /// Create a new `QuotaNamespace` from any string-like value.
19    pub fn new(s: impl Into<String>) -> Self {
20        Self(s.into())
21    }
22
23    /// Return the inner string slice.
24    pub fn as_str(&self) -> &str {
25        &self.0
26    }
27}
28
29impl std::fmt::Display for QuotaNamespace {
30    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
31        self.0.fmt(f)
32    }
33}
34
35// ---------------------------------------------------------------------------
36// Eviction strategy
37// ---------------------------------------------------------------------------
38
39/// Selects which objects are evicted first when a namespace is over quota.
40#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
41pub enum SqmEvictionStrategy {
42    /// Evict the object with the smallest `created_at` timestamp first.
43    Oldest,
44    /// Least Recently Used — evict the object with the smallest `last_accessed` timestamp.
45    Lru,
46    /// Least Frequently Used — evict the object with the smallest `access_count`;
47    /// ties broken by smallest `last_accessed`.
48    Lfu,
49    /// Evict the largest object (by `size_bytes`) first; ties broken by object_id.
50    SizeDescending,
51}
52
53// ---------------------------------------------------------------------------
54// QuotaPolicy
55// ---------------------------------------------------------------------------
56
57/// Configuration for a single namespace's quota.
58#[derive(Clone, Debug)]
59pub struct QuotaPolicy {
60    /// Hard byte ceiling; allocations that would exceed this are rejected.
61    pub max_bytes: u64,
62    /// Hard object-count ceiling.
63    pub max_objects: u64,
64    /// Fraction of `max_bytes` at which a `SoftLimitWarning` is emitted.
65    /// Must be in `(0.0, 1.0]`; defaults to `0.8`.
66    pub soft_limit_fraction: f64,
67    /// Strategy used by `evict_candidates`.
68    pub eviction_strategy: SqmEvictionStrategy,
69}
70
71impl QuotaPolicy {
72    /// Create a policy with default soft-limit fraction (0.8) and LRU eviction.
73    pub fn new(max_bytes: u64, max_objects: u64) -> Self {
74        Self {
75            max_bytes,
76            max_objects,
77            soft_limit_fraction: 0.8,
78            eviction_strategy: SqmEvictionStrategy::Lru,
79        }
80    }
81
82    /// Builder: set the soft-limit fraction.
83    pub fn with_soft_limit_fraction(mut self, fraction: f64) -> Self {
84        self.soft_limit_fraction = fraction;
85        self
86    }
87
88    /// Builder: set the eviction strategy.
89    pub fn with_eviction_strategy(mut self, strategy: SqmEvictionStrategy) -> Self {
90        self.eviction_strategy = strategy;
91        self
92    }
93}
94
95// ---------------------------------------------------------------------------
96// QuotaEntry
97// ---------------------------------------------------------------------------
98
99/// Live accounting entry for a registered namespace.
100#[derive(Clone, Debug)]
101pub struct SqmQuotaEntry {
102    /// Namespace this entry belongs to.
103    pub namespace: QuotaNamespace,
104    /// Current total bytes allocated across all objects.
105    pub bytes_used: u64,
106    /// Current number of allocated objects.
107    pub object_count: u64,
108    /// Unix-epoch timestamp (seconds) of the most recent successful `allocate`.
109    pub last_write: u64,
110    /// The policy governing this namespace.
111    pub policy: QuotaPolicy,
112}
113
114// ---------------------------------------------------------------------------
115// ObjectRecord
116// ---------------------------------------------------------------------------
117
118/// Per-object tracking record.
119#[derive(Clone, Debug)]
120pub struct ObjectRecord {
121    /// Unique identifier for this object.
122    pub object_id: String,
123    /// Namespace to which this object belongs.
124    pub namespace: QuotaNamespace,
125    /// Size of the object in bytes.
126    pub size_bytes: u64,
127    /// Unix-epoch timestamp when the object was created.
128    pub created_at: u64,
129    /// Unix-epoch timestamp of the most recent `access_object` call.
130    pub last_accessed: u64,
131    /// Total number of `access_object` calls.
132    pub access_count: u64,
133}
134
135// ---------------------------------------------------------------------------
136// QuotaViolation
137// ---------------------------------------------------------------------------
138
139/// Describes a quota-related event returned from `allocate`.
140#[derive(Clone, Debug, PartialEq)]
141pub enum SqmQuotaViolation {
142    /// The namespace's hard byte limit would be exceeded.
143    HardLimitExceeded {
144        /// Namespace name string.
145        namespace: String,
146        /// Current usage before this allocation.
147        used: u64,
148        /// Configured hard limit.
149        limit: u64,
150    },
151    /// The namespace's hard object-count limit would be exceeded.
152    ObjectLimitExceeded {
153        /// Namespace name string.
154        namespace: String,
155        /// Current object count before this allocation.
156        count: u64,
157        /// Configured hard object limit.
158        limit: u64,
159    },
160    /// Usage has crossed the soft-limit threshold (warning, not an error).
161    SoftLimitWarning {
162        /// Namespace name string.
163        namespace: String,
164        /// Current fraction (bytes_used / max_bytes) after allocation.
165        fraction: f64,
166    },
167}
168
169// ---------------------------------------------------------------------------
170// QuotaError
171// ---------------------------------------------------------------------------
172
173/// Errors returned by `StorageQuotaManager` operations.
174#[derive(Clone, Debug, PartialEq, thiserror::Error)]
175pub enum QuotaError {
176    /// A namespace with this identifier is already registered.
177    #[error("namespace already exists: {0}")]
178    NamespaceAlreadyExists(String),
179    /// No namespace with this identifier is registered.
180    #[error("namespace not found: {0}")]
181    NamespaceNotFound(String),
182    /// No object with this identifier is tracked.
183    #[error("object not found: {0}")]
184    ObjectNotFound(String),
185    /// An object with this identifier is already tracked.
186    #[error("object already exists: {0}")]
187    ObjectAlreadyExists(String),
188    /// This allocation would exceed the global byte ceiling.
189    #[error("global limit exceeded: used={used}, limit={limit}")]
190    GlobalLimitExceeded {
191        /// Bytes that would be used after this allocation.
192        used: u64,
193        /// Configured global limit.
194        limit: u64,
195    },
196    /// Namespace-level hard limit exceeded (byte or object count).
197    #[error("namespace hard limit exceeded in {namespace}")]
198    HardLimitExceeded {
199        /// Namespace name.
200        namespace: String,
201    },
202}
203
204// ---------------------------------------------------------------------------
205// QuotaStats
206// ---------------------------------------------------------------------------
207
208/// Aggregate statistics snapshot for the whole manager.
209#[derive(Clone, Debug)]
210pub struct QuotaStats {
211    /// Number of registered namespaces.
212    pub namespace_count: usize,
213    /// Sum of `bytes_used` across all namespaces.
214    pub total_bytes_used: u64,
215    /// Sum of `object_count` across all namespaces.
216    pub total_objects: usize,
217    /// `total_bytes_used / global_max_bytes` (0.0 when global_max_bytes == 0).
218    pub global_utilization: f64,
219    /// Namespaces where `usage_fraction ≥ soft_limit_fraction`.
220    pub namespaces_at_soft_limit: usize,
221    /// Namespaces where `usage_fraction ≥ 1.0`.
222    pub namespaces_at_hard_limit: usize,
223}
224
225// ---------------------------------------------------------------------------
226// StorageQuotaManager
227// ---------------------------------------------------------------------------
228
229/// Per-namespace storage quota enforcement manager.
230///
231/// Maintains a registry of namespaces with their policies and accounting entries,
232/// plus a flat index of all tracked objects for O(1) lookup.
233pub struct StorageQuotaManager {
234    /// Namespace entries keyed by namespace identifier.
235    pub entries: HashMap<QuotaNamespace, SqmQuotaEntry>,
236    /// Object records keyed by object_id.
237    pub objects: HashMap<String, ObjectRecord>,
238    /// Global byte ceiling across all namespaces combined.
239    pub global_max_bytes: u64,
240}
241
242impl StorageQuotaManager {
243    // -----------------------------------------------------------------------
244    // Construction
245    // -----------------------------------------------------------------------
246
247    /// Create a new manager with the given global byte ceiling.
248    ///
249    /// The global ceiling is an additional guard on top of per-namespace limits.
250    /// Pass `u64::MAX` to effectively disable it.
251    pub fn new(global_max_bytes: u64) -> Self {
252        Self {
253            entries: HashMap::new(),
254            objects: HashMap::new(),
255            global_max_bytes,
256        }
257    }
258
259    // -----------------------------------------------------------------------
260    // Namespace management
261    // -----------------------------------------------------------------------
262
263    /// Register a new namespace with the given policy.
264    ///
265    /// Returns [`QuotaError::NamespaceAlreadyExists`] if `ns` is already registered.
266    pub fn register_namespace(
267        &mut self,
268        ns: QuotaNamespace,
269        policy: QuotaPolicy,
270    ) -> Result<(), QuotaError> {
271        if self.entries.contains_key(&ns) {
272            return Err(QuotaError::NamespaceAlreadyExists(ns.0));
273        }
274        let entry = SqmQuotaEntry {
275            namespace: ns.clone(),
276            bytes_used: 0,
277            object_count: 0,
278            last_write: 0,
279            policy,
280        };
281        self.entries.insert(ns, entry);
282        Ok(())
283    }
284
285    /// Unregister a namespace and remove all objects that belong to it.
286    ///
287    /// Returns [`QuotaError::NamespaceNotFound`] if `ns` is not registered.
288    pub fn unregister_namespace(&mut self, ns: &QuotaNamespace) -> Result<(), QuotaError> {
289        if !self.entries.contains_key(ns) {
290            return Err(QuotaError::NamespaceNotFound(ns.0.clone()));
291        }
292        // Collect and remove all objects belonging to this namespace.
293        let to_remove: Vec<String> = self
294            .objects
295            .values()
296            .filter(|o| &o.namespace == ns)
297            .map(|o| o.object_id.clone())
298            .collect();
299        for oid in to_remove {
300            self.objects.remove(&oid);
301        }
302        self.entries.remove(ns);
303        Ok(())
304    }
305
306    // -----------------------------------------------------------------------
307    // Object lifecycle
308    // -----------------------------------------------------------------------
309
310    /// Allocate a new object in the given namespace.
311    ///
312    /// Checks the global limit and the namespace's hard limits before inserting.
313    /// Returns a (possibly empty) vector of warnings (e.g. soft-limit crossed).
314    /// Returns `Err` if any hard limit or the global limit would be exceeded, or
315    /// if the object already exists.
316    pub fn allocate(
317        &mut self,
318        object_id: String,
319        ns: &QuotaNamespace,
320        size_bytes: u64,
321        now: u64,
322    ) -> Result<Vec<SqmQuotaViolation>, QuotaError> {
323        // Existence check
324        if self.objects.contains_key(&object_id) {
325            return Err(QuotaError::ObjectAlreadyExists(object_id));
326        }
327
328        // Namespace must exist
329        let entry = self
330            .entries
331            .get(ns)
332            .ok_or_else(|| QuotaError::NamespaceNotFound(ns.0.clone()))?;
333
334        // Global limit check
335        let current_global = self.total_bytes_used();
336        let new_global = current_global.saturating_add(size_bytes);
337        if new_global > self.global_max_bytes {
338            return Err(QuotaError::GlobalLimitExceeded {
339                used: new_global,
340                limit: self.global_max_bytes,
341            });
342        }
343
344        // Namespace byte hard limit
345        let new_bytes = entry.bytes_used.saturating_add(size_bytes);
346        if new_bytes > entry.policy.max_bytes {
347            return Err(QuotaError::HardLimitExceeded {
348                namespace: ns.0.clone(),
349            });
350        }
351
352        // Namespace object count hard limit
353        let new_count = entry.object_count.saturating_add(1);
354        if new_count > entry.policy.max_objects {
355            return Err(QuotaError::HardLimitExceeded {
356                namespace: ns.0.clone(),
357            });
358        }
359
360        // Insert the object record
361        let record = ObjectRecord {
362            object_id: object_id.clone(),
363            namespace: ns.clone(),
364            size_bytes,
365            created_at: now,
366            last_accessed: now,
367            access_count: 0,
368        };
369        self.objects.insert(object_id, record);
370
371        // Update namespace accounting — re-borrow mutably
372        let entry_mut = self
373            .entries
374            .get_mut(ns)
375            .ok_or_else(|| QuotaError::NamespaceNotFound(ns.0.clone()))?;
376        entry_mut.bytes_used = new_bytes;
377        entry_mut.object_count = new_count;
378        entry_mut.last_write = now;
379
380        // Emit any warnings
381        let mut warnings = Vec::new();
382        let fraction = if entry_mut.policy.max_bytes > 0 {
383            entry_mut.bytes_used as f64 / entry_mut.policy.max_bytes as f64
384        } else {
385            0.0
386        };
387        if fraction >= entry_mut.policy.soft_limit_fraction {
388            warnings.push(SqmQuotaViolation::SoftLimitWarning {
389                namespace: ns.0.clone(),
390                fraction,
391            });
392        }
393
394        Ok(warnings)
395    }
396
397    /// Remove an object and return the number of bytes freed.
398    ///
399    /// Updates the namespace accounting entry accordingly.
400    pub fn deallocate(&mut self, object_id: &str, _now: u64) -> Result<u64, QuotaError> {
401        let record = self
402            .objects
403            .remove(object_id)
404            .ok_or_else(|| QuotaError::ObjectNotFound(object_id.to_string()))?;
405
406        let freed = record.size_bytes;
407        let ns = &record.namespace.clone();
408
409        if let Some(entry) = self.entries.get_mut(ns) {
410            entry.bytes_used = entry.bytes_used.saturating_sub(freed);
411            entry.object_count = entry.object_count.saturating_sub(1);
412        }
413
414        Ok(freed)
415    }
416
417    /// Update `last_accessed` and increment `access_count` for LRU/LFU eviction.
418    pub fn access_object(&mut self, object_id: &str, now: u64) -> Result<(), QuotaError> {
419        let record = self
420            .objects
421            .get_mut(object_id)
422            .ok_or_else(|| QuotaError::ObjectNotFound(object_id.to_string()))?;
423        record.last_accessed = now;
424        record.access_count = record.access_count.saturating_add(1);
425        Ok(())
426    }
427
428    // -----------------------------------------------------------------------
429    // Usage queries
430    // -----------------------------------------------------------------------
431
432    /// Return `bytes_used / max_bytes` for the namespace, or `None` if not registered.
433    pub fn usage_fraction(&self, ns: &QuotaNamespace) -> Option<f64> {
434        let entry = self.entries.get(ns)?;
435        if entry.policy.max_bytes == 0 {
436            Some(0.0)
437        } else {
438            Some(entry.bytes_used as f64 / entry.policy.max_bytes as f64)
439        }
440    }
441
442    /// Return `true` if the namespace's usage fraction is ≥ 1.0 (at or over hard limit).
443    pub fn needs_eviction(&self, ns: &QuotaNamespace) -> bool {
444        self.usage_fraction(ns).is_some_and(|f| f >= 1.0)
445    }
446
447    /// Return a reference to the accounting entry for the namespace.
448    pub fn namespace_usage(&self, ns: &QuotaNamespace) -> Option<&SqmQuotaEntry> {
449        self.entries.get(ns)
450    }
451
452    /// Return all registered namespace identifiers, sorted alphabetically.
453    pub fn all_namespaces(&self) -> Vec<&QuotaNamespace> {
454        let mut keys: Vec<&QuotaNamespace> = self.entries.keys().collect();
455        keys.sort();
456        keys
457    }
458
459    /// Sum of `bytes_used` across all namespaces.
460    pub fn total_bytes_used(&self) -> u64 {
461        self.entries.values().map(|e| e.bytes_used).sum()
462    }
463
464    /// Sum of `object_count` across all namespaces.
465    pub fn total_objects(&self) -> usize {
466        self.entries.values().map(|e| e.object_count as usize).sum()
467    }
468
469    // -----------------------------------------------------------------------
470    // Eviction
471    // -----------------------------------------------------------------------
472
473    /// Return a list of object IDs to evict according to the namespace's strategy.
474    ///
475    /// Objects are selected until removing them would free at least `target_free_bytes`.
476    /// The returned list is never longer than the number of objects in the namespace.
477    /// Objects are **not** removed by this call; use `force_evict` or call `deallocate`
478    /// for each candidate.
479    pub fn evict_candidates(&self, ns: &QuotaNamespace, target_free_bytes: u64) -> Vec<String> {
480        let entry = match self.entries.get(ns) {
481            Some(e) => e,
482            None => return Vec::new(),
483        };
484
485        // Collect objects in this namespace
486        let mut candidates: Vec<&ObjectRecord> = self
487            .objects
488            .values()
489            .filter(|o| &o.namespace == ns)
490            .collect();
491
492        if candidates.is_empty() {
493            return Vec::new();
494        }
495
496        // Sort according to the eviction strategy (ascending = evicted first)
497        match entry.policy.eviction_strategy {
498            SqmEvictionStrategy::Oldest => {
499                candidates.sort_by_key(|o| (o.created_at, o.object_id.as_str().to_string()));
500            }
501            SqmEvictionStrategy::Lru => {
502                candidates.sort_by_key(|o| (o.last_accessed, o.object_id.as_str().to_string()));
503            }
504            SqmEvictionStrategy::Lfu => {
505                candidates.sort_by_key(|o| {
506                    (
507                        o.access_count,
508                        o.last_accessed,
509                        o.object_id.as_str().to_string(),
510                    )
511                });
512            }
513            SqmEvictionStrategy::SizeDescending => {
514                // Largest first → sort descending by size, then by object_id for stability
515                candidates.sort_by(|a, b| {
516                    b.size_bytes
517                        .cmp(&a.size_bytes)
518                        .then_with(|| a.object_id.cmp(&b.object_id))
519                });
520            }
521        }
522
523        // Greedily pick until we have covered `target_free_bytes`
524        let mut freed: u64 = 0;
525        let mut result = Vec::new();
526        for obj in candidates {
527            if freed >= target_free_bytes {
528                break;
529            }
530            freed = freed.saturating_add(obj.size_bytes);
531            result.push(obj.object_id.clone());
532        }
533
534        result
535    }
536
537    /// Evict objects from the namespace until at least `target_free_bytes` are freed.
538    ///
539    /// Calls `evict_candidates` then `deallocate` on each candidate.
540    /// Returns the total number of bytes actually freed.
541    pub fn force_evict(
542        &mut self,
543        ns: &QuotaNamespace,
544        target_free_bytes: u64,
545        now: u64,
546    ) -> Result<u64, QuotaError> {
547        // We must not hold a borrow when calling deallocate, so collect first.
548        let candidates = self.evict_candidates(ns, target_free_bytes);
549
550        if candidates.is_empty() {
551            // Return 0 if namespace exists but has no objects; error if unknown.
552            if self.entries.contains_key(ns) {
553                return Ok(0);
554            } else {
555                return Err(QuotaError::NamespaceNotFound(ns.0.clone()));
556            }
557        }
558
559        let mut total_freed: u64 = 0;
560        for oid in candidates {
561            match self.deallocate(&oid, now) {
562                Ok(freed) => total_freed = total_freed.saturating_add(freed),
563                Err(QuotaError::ObjectNotFound(_)) => {
564                    // Concurrent removal is acceptable; skip.
565                }
566                Err(e) => return Err(e),
567            }
568        }
569
570        Ok(total_freed)
571    }
572
573    // -----------------------------------------------------------------------
574    // Aggregate stats
575    // -----------------------------------------------------------------------
576
577    /// Return an aggregate statistics snapshot.
578    pub fn stats(&self) -> QuotaStats {
579        let namespace_count = self.entries.len();
580        let total_bytes_used = self.total_bytes_used();
581        let total_objects = self.total_objects();
582
583        let global_utilization = if self.global_max_bytes > 0 {
584            total_bytes_used as f64 / self.global_max_bytes as f64
585        } else {
586            0.0
587        };
588
589        let mut namespaces_at_soft_limit = 0usize;
590        let mut namespaces_at_hard_limit = 0usize;
591
592        for entry in self.entries.values() {
593            if entry.policy.max_bytes == 0 {
594                continue;
595            }
596            let frac = entry.bytes_used as f64 / entry.policy.max_bytes as f64;
597            if frac >= 1.0 {
598                namespaces_at_hard_limit += 1;
599                // A namespace at the hard limit is also at the soft limit
600                namespaces_at_soft_limit += 1;
601            } else if frac >= entry.policy.soft_limit_fraction {
602                namespaces_at_soft_limit += 1;
603            }
604        }
605
606        QuotaStats {
607            namespace_count,
608            total_bytes_used,
609            total_objects,
610            global_utilization,
611            namespaces_at_soft_limit,
612            namespaces_at_hard_limit,
613        }
614    }
615}
616
617// ---------------------------------------------------------------------------
618// Tests
619// ---------------------------------------------------------------------------
620
621#[cfg(test)]
622mod tests {
623    use crate::storage_quota_manager::{
624        QuotaError, QuotaNamespace, QuotaPolicy, QuotaStats, SqmEvictionStrategy,
625        SqmQuotaViolation, StorageQuotaManager,
626    };
627
628    fn make_ns(s: &str) -> QuotaNamespace {
629        QuotaNamespace::new(s)
630    }
631
632    fn make_policy(max_bytes: u64, max_objects: u64) -> QuotaPolicy {
633        QuotaPolicy::new(max_bytes, max_objects)
634    }
635
636    // --- Construction ---
637
638    #[test]
639    fn test_new_manager_empty() {
640        let mgr = StorageQuotaManager::new(1024 * 1024);
641        assert_eq!(mgr.total_bytes_used(), 0);
642        assert_eq!(mgr.total_objects(), 0);
643        assert!(mgr.all_namespaces().is_empty());
644    }
645
646    // --- Namespace registration ---
647
648    #[test]
649    fn test_register_namespace_ok() {
650        let mut mgr = StorageQuotaManager::new(u64::MAX);
651        let ns = make_ns("alice");
652        let policy = make_policy(1000, 10);
653        assert!(mgr.register_namespace(ns.clone(), policy).is_ok());
654        assert!(mgr.namespace_usage(&ns).is_some());
655    }
656
657    #[test]
658    fn test_register_namespace_duplicate_error() {
659        let mut mgr = StorageQuotaManager::new(u64::MAX);
660        let ns = make_ns("alice");
661        mgr.register_namespace(ns.clone(), make_policy(1000, 10))
662            .unwrap();
663        let result = mgr.register_namespace(ns.clone(), make_policy(2000, 20));
664        assert!(matches!(result, Err(QuotaError::NamespaceAlreadyExists(_))));
665    }
666
667    #[test]
668    fn test_unregister_namespace_ok() {
669        let mut mgr = StorageQuotaManager::new(u64::MAX);
670        let ns = make_ns("bob");
671        mgr.register_namespace(ns.clone(), make_policy(1000, 10))
672            .unwrap();
673        assert!(mgr.unregister_namespace(&ns).is_ok());
674        assert!(mgr.namespace_usage(&ns).is_none());
675    }
676
677    #[test]
678    fn test_unregister_namespace_not_found() {
679        let mut mgr = StorageQuotaManager::new(u64::MAX);
680        let ns = make_ns("ghost");
681        let result = mgr.unregister_namespace(&ns);
682        assert!(matches!(result, Err(QuotaError::NamespaceNotFound(_))));
683    }
684
685    #[test]
686    fn test_unregister_removes_objects() {
687        let mut mgr = StorageQuotaManager::new(u64::MAX);
688        let ns = make_ns("carol");
689        mgr.register_namespace(ns.clone(), make_policy(10_000, 100))
690            .unwrap();
691        mgr.allocate("obj-1".to_string(), &ns, 100, 1).unwrap();
692        mgr.allocate("obj-2".to_string(), &ns, 200, 2).unwrap();
693        assert_eq!(mgr.total_objects(), 2);
694        mgr.unregister_namespace(&ns).unwrap();
695        assert_eq!(mgr.total_objects(), 0);
696        assert!(!mgr.objects.contains_key("obj-1"));
697        assert!(!mgr.objects.contains_key("obj-2"));
698    }
699
700    // --- Allocation ---
701
702    #[test]
703    fn test_allocate_basic() {
704        let mut mgr = StorageQuotaManager::new(u64::MAX);
705        let ns = make_ns("dave");
706        mgr.register_namespace(ns.clone(), make_policy(1000, 10))
707            .unwrap();
708        let warnings = mgr.allocate("obj-a".to_string(), &ns, 100, 10).unwrap();
709        // 100/1000 = 10%, below soft limit 0.8
710        assert!(warnings.is_empty());
711        let entry = mgr.namespace_usage(&ns).unwrap();
712        assert_eq!(entry.bytes_used, 100);
713        assert_eq!(entry.object_count, 1);
714        assert_eq!(entry.last_write, 10);
715    }
716
717    #[test]
718    fn test_allocate_soft_limit_warning() {
719        let mut mgr = StorageQuotaManager::new(u64::MAX);
720        let ns = make_ns("eve");
721        // soft_limit_fraction = 0.8, max_bytes = 100
722        let policy = make_policy(100, 100).with_soft_limit_fraction(0.8);
723        mgr.register_namespace(ns.clone(), policy).unwrap();
724        // Allocate 85 bytes → 85%
725        let warnings = mgr.allocate("obj".to_string(), &ns, 85, 1).unwrap();
726        assert!(!warnings.is_empty());
727        assert!(warnings
728            .iter()
729            .any(|w| matches!(w, SqmQuotaViolation::SoftLimitWarning { .. })));
730    }
731
732    #[test]
733    fn test_allocate_hard_byte_limit_error() {
734        let mut mgr = StorageQuotaManager::new(u64::MAX);
735        let ns = make_ns("frank");
736        mgr.register_namespace(ns.clone(), make_policy(100, 100))
737            .unwrap();
738        let result = mgr.allocate("big".to_string(), &ns, 200, 1);
739        assert!(matches!(result, Err(QuotaError::HardLimitExceeded { .. })));
740    }
741
742    #[test]
743    fn test_allocate_object_count_limit_error() {
744        let mut mgr = StorageQuotaManager::new(u64::MAX);
745        let ns = make_ns("grace");
746        // max 2 objects, plenty of bytes
747        mgr.register_namespace(ns.clone(), make_policy(100_000, 2))
748            .unwrap();
749        mgr.allocate("o1".to_string(), &ns, 10, 1).unwrap();
750        mgr.allocate("o2".to_string(), &ns, 10, 2).unwrap();
751        let result = mgr.allocate("o3".to_string(), &ns, 10, 3);
752        assert!(matches!(result, Err(QuotaError::HardLimitExceeded { .. })));
753    }
754
755    #[test]
756    fn test_allocate_duplicate_object_error() {
757        let mut mgr = StorageQuotaManager::new(u64::MAX);
758        let ns = make_ns("henry");
759        mgr.register_namespace(ns.clone(), make_policy(10_000, 100))
760            .unwrap();
761        mgr.allocate("dup".to_string(), &ns, 10, 1).unwrap();
762        let result = mgr.allocate("dup".to_string(), &ns, 10, 2);
763        assert!(matches!(result, Err(QuotaError::ObjectAlreadyExists(_))));
764    }
765
766    #[test]
767    fn test_allocate_unknown_namespace_error() {
768        let mut mgr = StorageQuotaManager::new(u64::MAX);
769        let ns = make_ns("unknown");
770        let result = mgr.allocate("x".to_string(), &ns, 10, 1);
771        assert!(matches!(result, Err(QuotaError::NamespaceNotFound(_))));
772    }
773
774    #[test]
775    fn test_allocate_global_limit_error() {
776        let mut mgr = StorageQuotaManager::new(50);
777        let ns = make_ns("ivy");
778        mgr.register_namespace(ns.clone(), make_policy(1_000_000, 1_000_000))
779            .unwrap();
780        let result = mgr.allocate("big".to_string(), &ns, 100, 1);
781        assert!(matches!(
782            result,
783            Err(QuotaError::GlobalLimitExceeded { .. })
784        ));
785    }
786
787    // --- Deallocation ---
788
789    #[test]
790    fn test_deallocate_ok() {
791        let mut mgr = StorageQuotaManager::new(u64::MAX);
792        let ns = make_ns("jack");
793        mgr.register_namespace(ns.clone(), make_policy(10_000, 100))
794            .unwrap();
795        mgr.allocate("o1".to_string(), &ns, 300, 1).unwrap();
796        let freed = mgr.deallocate("o1", 2).unwrap();
797        assert_eq!(freed, 300);
798        assert_eq!(mgr.namespace_usage(&ns).unwrap().bytes_used, 0);
799        assert_eq!(mgr.namespace_usage(&ns).unwrap().object_count, 0);
800    }
801
802    #[test]
803    fn test_deallocate_not_found() {
804        let mut mgr = StorageQuotaManager::new(u64::MAX);
805        assert!(matches!(
806            mgr.deallocate("nope", 1),
807            Err(QuotaError::ObjectNotFound(_))
808        ));
809    }
810
811    #[test]
812    fn test_deallocate_updates_total() {
813        let mut mgr = StorageQuotaManager::new(u64::MAX);
814        let ns = make_ns("kate");
815        mgr.register_namespace(ns.clone(), make_policy(10_000, 100))
816            .unwrap();
817        mgr.allocate("a".to_string(), &ns, 100, 1).unwrap();
818        mgr.allocate("b".to_string(), &ns, 200, 2).unwrap();
819        assert_eq!(mgr.total_bytes_used(), 300);
820        mgr.deallocate("a", 3).unwrap();
821        assert_eq!(mgr.total_bytes_used(), 200);
822    }
823
824    // --- Access tracking ---
825
826    #[test]
827    fn test_access_object_updates_fields() {
828        let mut mgr = StorageQuotaManager::new(u64::MAX);
829        let ns = make_ns("leo");
830        mgr.register_namespace(ns.clone(), make_policy(10_000, 100))
831            .unwrap();
832        mgr.allocate("obj".to_string(), &ns, 50, 1).unwrap();
833        mgr.access_object("obj", 10).unwrap();
834        mgr.access_object("obj", 20).unwrap();
835        let rec = mgr.objects.get("obj").unwrap();
836        assert_eq!(rec.last_accessed, 20);
837        assert_eq!(rec.access_count, 2);
838    }
839
840    #[test]
841    fn test_access_object_not_found() {
842        let mut mgr = StorageQuotaManager::new(u64::MAX);
843        assert!(matches!(
844            mgr.access_object("missing", 1),
845            Err(QuotaError::ObjectNotFound(_))
846        ));
847    }
848
849    // --- Usage fraction ---
850
851    #[test]
852    fn test_usage_fraction_zero_when_empty() {
853        let mut mgr = StorageQuotaManager::new(u64::MAX);
854        let ns = make_ns("mia");
855        mgr.register_namespace(ns.clone(), make_policy(1000, 10))
856            .unwrap();
857        assert_eq!(mgr.usage_fraction(&ns), Some(0.0));
858    }
859
860    #[test]
861    fn test_usage_fraction_correct() {
862        let mut mgr = StorageQuotaManager::new(u64::MAX);
863        let ns = make_ns("ned");
864        mgr.register_namespace(ns.clone(), make_policy(1000, 10))
865            .unwrap();
866        mgr.allocate("o".to_string(), &ns, 500, 1).unwrap();
867        assert!((mgr.usage_fraction(&ns).unwrap() - 0.5).abs() < 1e-9);
868    }
869
870    #[test]
871    fn test_usage_fraction_none_for_unknown() {
872        let mgr = StorageQuotaManager::new(u64::MAX);
873        let ns = make_ns("nobody");
874        assert_eq!(mgr.usage_fraction(&ns), None);
875    }
876
877    // --- needs_eviction ---
878
879    #[test]
880    fn test_needs_eviction_false_when_under_limit() {
881        let mut mgr = StorageQuotaManager::new(u64::MAX);
882        let ns = make_ns("olivia");
883        mgr.register_namespace(ns.clone(), make_policy(1000, 10))
884            .unwrap();
885        mgr.allocate("o".to_string(), &ns, 500, 1).unwrap();
886        assert!(!mgr.needs_eviction(&ns));
887    }
888
889    #[test]
890    fn test_needs_eviction_true_when_at_limit() {
891        let mut mgr = StorageQuotaManager::new(u64::MAX);
892        let ns = make_ns("peter");
893        mgr.register_namespace(ns.clone(), make_policy(100, 100))
894            .unwrap();
895        mgr.allocate("o".to_string(), &ns, 100, 1).unwrap();
896        assert!(mgr.needs_eviction(&ns));
897    }
898
899    // --- all_namespaces sorted ---
900
901    #[test]
902    fn test_all_namespaces_sorted() {
903        let mut mgr = StorageQuotaManager::new(u64::MAX);
904        mgr.register_namespace(make_ns("z-ns"), make_policy(100, 10))
905            .unwrap();
906        mgr.register_namespace(make_ns("a-ns"), make_policy(100, 10))
907            .unwrap();
908        mgr.register_namespace(make_ns("m-ns"), make_policy(100, 10))
909            .unwrap();
910        let namespaces: Vec<String> = mgr.all_namespaces().iter().map(|n| n.0.clone()).collect();
911        assert_eq!(namespaces, vec!["a-ns", "m-ns", "z-ns"]);
912    }
913
914    // --- Eviction candidates ---
915
916    #[test]
917    fn test_evict_candidates_oldest() {
918        let mut mgr = StorageQuotaManager::new(u64::MAX);
919        let ns = make_ns("quinn");
920        let policy = make_policy(10_000, 100).with_eviction_strategy(SqmEvictionStrategy::Oldest);
921        mgr.register_namespace(ns.clone(), policy).unwrap();
922        mgr.allocate("old".to_string(), &ns, 100, 1).unwrap();
923        mgr.allocate("new".to_string(), &ns, 100, 100).unwrap();
924        let candidates = mgr.evict_candidates(&ns, 100);
925        assert_eq!(candidates.first().map(String::as_str), Some("old"));
926    }
927
928    #[test]
929    fn test_evict_candidates_lru() {
930        let mut mgr = StorageQuotaManager::new(u64::MAX);
931        let ns = make_ns("rose");
932        let policy = make_policy(10_000, 100).with_eviction_strategy(SqmEvictionStrategy::Lru);
933        mgr.register_namespace(ns.clone(), policy).unwrap();
934        mgr.allocate("a".to_string(), &ns, 100, 1).unwrap();
935        mgr.allocate("b".to_string(), &ns, 100, 1).unwrap();
936        // Access "a" more recently
937        mgr.access_object("a", 100).unwrap();
938        let candidates = mgr.evict_candidates(&ns, 100);
939        // "b" was last accessed earlier
940        assert_eq!(candidates.first().map(String::as_str), Some("b"));
941    }
942
943    #[test]
944    fn test_evict_candidates_lfu() {
945        let mut mgr = StorageQuotaManager::new(u64::MAX);
946        let ns = make_ns("sam");
947        let policy = make_policy(10_000, 100).with_eviction_strategy(SqmEvictionStrategy::Lfu);
948        mgr.register_namespace(ns.clone(), policy).unwrap();
949        mgr.allocate("freq".to_string(), &ns, 100, 1).unwrap();
950        mgr.allocate("rare".to_string(), &ns, 100, 1).unwrap();
951        // Access "freq" many times
952        for t in 2..12_u64 {
953            mgr.access_object("freq", t).unwrap();
954        }
955        let candidates = mgr.evict_candidates(&ns, 100);
956        // "rare" has access_count=0 → evicted first
957        assert_eq!(candidates.first().map(String::as_str), Some("rare"));
958    }
959
960    #[test]
961    fn test_evict_candidates_size_descending() {
962        let mut mgr = StorageQuotaManager::new(u64::MAX);
963        let ns = make_ns("tara");
964        let policy =
965            make_policy(10_000, 100).with_eviction_strategy(SqmEvictionStrategy::SizeDescending);
966        mgr.register_namespace(ns.clone(), policy).unwrap();
967        mgr.allocate("small".to_string(), &ns, 50, 1).unwrap();
968        mgr.allocate("large".to_string(), &ns, 5000, 1).unwrap();
969        mgr.allocate("medium".to_string(), &ns, 500, 1).unwrap();
970        let candidates = mgr.evict_candidates(&ns, 1);
971        assert_eq!(candidates.first().map(String::as_str), Some("large"));
972    }
973
974    #[test]
975    fn test_evict_candidates_covers_target() {
976        let mut mgr = StorageQuotaManager::new(u64::MAX);
977        let ns = make_ns("ulrich");
978        let policy = make_policy(10_000, 100).with_eviction_strategy(SqmEvictionStrategy::Oldest);
979        mgr.register_namespace(ns.clone(), policy).unwrap();
980        for i in 0..10_u64 {
981            mgr.allocate(format!("obj-{i}"), &ns, 100, i).unwrap();
982        }
983        // Need to free 350 bytes → 4 objects (each 100 bytes covers after 4: 400 ≥ 350)
984        let candidates = mgr.evict_candidates(&ns, 350);
985        let total: u64 = candidates
986            .iter()
987            .map(|id| mgr.objects.get(id).map_or(0, |o| o.size_bytes))
988            .sum();
989        assert!(total >= 350, "Should cover at least 350 bytes, got {total}");
990    }
991
992    #[test]
993    fn test_evict_candidates_empty_namespace() {
994        let mut mgr = StorageQuotaManager::new(u64::MAX);
995        let ns = make_ns("vera");
996        mgr.register_namespace(ns.clone(), make_policy(1000, 10))
997            .unwrap();
998        assert!(mgr.evict_candidates(&ns, 100).is_empty());
999    }
1000
1001    #[test]
1002    fn test_evict_candidates_unknown_namespace() {
1003        let mgr = StorageQuotaManager::new(u64::MAX);
1004        let ns = make_ns("nobody");
1005        assert!(mgr.evict_candidates(&ns, 100).is_empty());
1006    }
1007
1008    // --- force_evict ---
1009
1010    #[test]
1011    fn test_force_evict_frees_bytes() {
1012        let mut mgr = StorageQuotaManager::new(u64::MAX);
1013        let ns = make_ns("will");
1014        let policy = make_policy(10_000, 100).with_eviction_strategy(SqmEvictionStrategy::Oldest);
1015        mgr.register_namespace(ns.clone(), policy).unwrap();
1016        mgr.allocate("a".to_string(), &ns, 200, 1).unwrap();
1017        mgr.allocate("b".to_string(), &ns, 200, 2).unwrap();
1018        mgr.allocate("c".to_string(), &ns, 200, 3).unwrap();
1019        let freed = mgr.force_evict(&ns, 300, 10).unwrap();
1020        assert!(freed >= 300, "freed={freed}");
1021    }
1022
1023    #[test]
1024    fn test_force_evict_updates_accounting() {
1025        let mut mgr = StorageQuotaManager::new(u64::MAX);
1026        let ns = make_ns("xena");
1027        let policy =
1028            make_policy(10_000, 100).with_eviction_strategy(SqmEvictionStrategy::SizeDescending);
1029        mgr.register_namespace(ns.clone(), policy).unwrap();
1030        mgr.allocate("big".to_string(), &ns, 1000, 1).unwrap();
1031        mgr.allocate("small".to_string(), &ns, 50, 2).unwrap();
1032        let before = mgr.namespace_usage(&ns).unwrap().bytes_used;
1033        let freed = mgr.force_evict(&ns, 500, 10).unwrap();
1034        let after = mgr.namespace_usage(&ns).unwrap().bytes_used;
1035        assert_eq!(before - after, freed);
1036    }
1037
1038    #[test]
1039    fn test_force_evict_unknown_namespace_error() {
1040        let mut mgr = StorageQuotaManager::new(u64::MAX);
1041        let ns = make_ns("nobody");
1042        let result = mgr.force_evict(&ns, 100, 1);
1043        assert!(matches!(result, Err(QuotaError::NamespaceNotFound(_))));
1044    }
1045
1046    // --- stats ---
1047
1048    #[test]
1049    fn test_stats_empty() {
1050        let mgr = StorageQuotaManager::new(1024);
1051        let s: QuotaStats = mgr.stats();
1052        assert_eq!(s.namespace_count, 0);
1053        assert_eq!(s.total_bytes_used, 0);
1054        assert_eq!(s.total_objects, 0);
1055        assert_eq!(s.global_utilization, 0.0);
1056        assert_eq!(s.namespaces_at_soft_limit, 0);
1057        assert_eq!(s.namespaces_at_hard_limit, 0);
1058    }
1059
1060    #[test]
1061    fn test_stats_counts_correctly() {
1062        let mut mgr = StorageQuotaManager::new(10_000);
1063        let ns1 = make_ns("y-ns");
1064        let ns2 = make_ns("z-ns");
1065        mgr.register_namespace(ns1.clone(), make_policy(500, 10))
1066            .unwrap();
1067        mgr.register_namespace(ns2.clone(), make_policy(500, 10))
1068            .unwrap();
1069        // Bring ns1 to 90% (above soft limit 0.8)
1070        mgr.allocate("obj".to_string(), &ns1, 450, 1).unwrap();
1071        let s = mgr.stats();
1072        assert_eq!(s.namespace_count, 2);
1073        assert_eq!(s.total_bytes_used, 450);
1074        assert_eq!(s.total_objects, 1);
1075        assert_eq!(s.namespaces_at_soft_limit, 1);
1076        assert_eq!(s.namespaces_at_hard_limit, 0);
1077    }
1078
1079    #[test]
1080    fn test_stats_hard_limit_count() {
1081        let mut mgr = StorageQuotaManager::new(u64::MAX);
1082        let ns = make_ns("zara");
1083        mgr.register_namespace(ns.clone(), make_policy(100, 100))
1084            .unwrap();
1085        mgr.allocate("full".to_string(), &ns, 100, 1).unwrap();
1086        let s = mgr.stats();
1087        assert_eq!(s.namespaces_at_hard_limit, 1);
1088        assert_eq!(s.namespaces_at_soft_limit, 1);
1089    }
1090
1091    // --- QuotaNamespace helpers ---
1092
1093    #[test]
1094    fn test_quota_namespace_as_str() {
1095        let ns = QuotaNamespace::new("test-ns");
1096        assert_eq!(ns.as_str(), "test-ns");
1097    }
1098
1099    #[test]
1100    fn test_quota_namespace_display() {
1101        let ns = QuotaNamespace::new("display-me");
1102        assert_eq!(format!("{ns}"), "display-me");
1103    }
1104
1105    // --- QuotaPolicy builder ---
1106
1107    #[test]
1108    fn test_policy_builder_defaults() {
1109        let p = make_policy(1000, 20);
1110        assert!((p.soft_limit_fraction - 0.8).abs() < 1e-9);
1111        assert_eq!(p.eviction_strategy, SqmEvictionStrategy::Lru);
1112    }
1113
1114    #[test]
1115    fn test_policy_builder_custom() {
1116        let p = make_policy(1000, 20)
1117            .with_soft_limit_fraction(0.5)
1118            .with_eviction_strategy(SqmEvictionStrategy::Oldest);
1119        assert!((p.soft_limit_fraction - 0.5).abs() < 1e-9);
1120        assert_eq!(p.eviction_strategy, SqmEvictionStrategy::Oldest);
1121    }
1122
1123    // --- Global utilization ---
1124
1125    #[test]
1126    fn test_global_utilization_in_stats() {
1127        let mut mgr = StorageQuotaManager::new(1000);
1128        let ns = make_ns("util-test");
1129        mgr.register_namespace(ns.clone(), make_policy(1000, 100))
1130            .unwrap();
1131        mgr.allocate("o".to_string(), &ns, 250, 1).unwrap();
1132        let s = mgr.stats();
1133        assert!((s.global_utilization - 0.25).abs() < 1e-9);
1134    }
1135
1136    // --- Multiple namespaces isolation ---
1137
1138    #[test]
1139    fn test_namespaces_are_isolated() {
1140        let mut mgr = StorageQuotaManager::new(u64::MAX);
1141        let ns_a = make_ns("iso-a");
1142        let ns_b = make_ns("iso-b");
1143        mgr.register_namespace(ns_a.clone(), make_policy(500, 10))
1144            .unwrap();
1145        mgr.register_namespace(ns_b.clone(), make_policy(500, 10))
1146            .unwrap();
1147        mgr.allocate("a1".to_string(), &ns_a, 300, 1).unwrap();
1148        mgr.allocate("b1".to_string(), &ns_b, 200, 1).unwrap();
1149        assert_eq!(mgr.namespace_usage(&ns_a).unwrap().bytes_used, 300);
1150        assert_eq!(mgr.namespace_usage(&ns_b).unwrap().bytes_used, 200);
1151    }
1152
1153    // --- Re-use object id after deallocate ---
1154
1155    #[test]
1156    fn test_reallocate_after_deallocate() {
1157        let mut mgr = StorageQuotaManager::new(u64::MAX);
1158        let ns = make_ns("reuse");
1159        mgr.register_namespace(ns.clone(), make_policy(10_000, 100))
1160            .unwrap();
1161        mgr.allocate("slot".to_string(), &ns, 100, 1).unwrap();
1162        mgr.deallocate("slot", 2).unwrap();
1163        // Should succeed now that "slot" is gone
1164        assert!(mgr.allocate("slot".to_string(), &ns, 100, 3).is_ok());
1165    }
1166
1167    // --- Soft limit exactly at boundary ---
1168
1169    #[test]
1170    fn test_soft_limit_exactly_at_boundary() {
1171        let mut mgr = StorageQuotaManager::new(u64::MAX);
1172        let ns = make_ns("boundary");
1173        let policy = make_policy(100, 100).with_soft_limit_fraction(0.9);
1174        mgr.register_namespace(ns.clone(), policy).unwrap();
1175        // 90 bytes = exactly at soft limit
1176        let warnings = mgr.allocate("o".to_string(), &ns, 90, 1).unwrap();
1177        assert!(!warnings.is_empty());
1178    }
1179
1180    // --- total_objects accuracy ---
1181
1182    #[test]
1183    fn test_total_objects_across_namespaces() {
1184        let mut mgr = StorageQuotaManager::new(u64::MAX);
1185        let ns1 = make_ns("t1");
1186        let ns2 = make_ns("t2");
1187        mgr.register_namespace(ns1.clone(), make_policy(10_000, 100))
1188            .unwrap();
1189        mgr.register_namespace(ns2.clone(), make_policy(10_000, 100))
1190            .unwrap();
1191        mgr.allocate("a".to_string(), &ns1, 10, 1).unwrap();
1192        mgr.allocate("b".to_string(), &ns1, 10, 2).unwrap();
1193        mgr.allocate("c".to_string(), &ns2, 10, 3).unwrap();
1194        assert_eq!(mgr.total_objects(), 3);
1195        mgr.deallocate("b", 4).unwrap();
1196        assert_eq!(mgr.total_objects(), 2);
1197    }
1198}