Skip to main content

eredu_core/
residency.rs

1//! Backend-neutral policy, ownership, capacity, accounting, and telemetry for weight residency.
2//!
3//! Placement determines which tensors a rank owns. The types in this module
4//! describe a separate residency decision: the tier in which an owned logical
5//! unit is intended to reside and its lifetime policy. This module validates
6//! explicit plans and records observations without knowing how a backend
7//! materializes, transfers, or synchronizes resources.
8
9use serde::{de::Error as _, Deserialize, Deserializer, Serialize};
10use std::{
11    collections::{BTreeMap, BTreeSet},
12    fmt,
13    time::Duration,
14};
15
16mod prefetch;
17
18pub use prefetch::{
19    BackgroundPrefetchReport, PrefetchAdmission, PrefetchCompletion, PrefetchDemandObservation,
20    PrefetchDemandResolution, PrefetchExecutionState, PrefetchStateError, PrefetchWork,
21};
22
23/// Current serialized residency-plan schema.
24pub const OFFLOAD_PLAN_SCHEMA_VERSION: u32 = 1;
25
26/// A storage or execution-memory tier used by an offload plan.
27#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
28#[serde(rename_all = "snake_case")]
29pub enum MemoryTier {
30    /// Memory directly used for device execution.
31    Device,
32    /// Host-accessible memory.
33    Host,
34    /// Disk-backed storage.
35    Disk,
36}
37
38impl MemoryTier {
39    const fn index(self) -> usize {
40        match self {
41            Self::Device => 0,
42            Self::Host => 1,
43            Self::Disk => 2,
44        }
45    }
46}
47
48/// The intended lifetime behavior of an offload unit within a tier.
49#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
50#[serde(rename_all = "snake_case")]
51pub enum ResidencyPolicy {
52    /// Keep the unit resident for the lifetime of the residency manager.
53    Pinned,
54    /// Keep the unit resident only within a bounded execution window.
55    Windowed,
56    /// Allow the residency manager to retain or evict the unit as cache policy permits.
57    Cacheable,
58}
59
60/// Deterministic eviction ordering for cacheable residency units.
61#[derive(
62    Debug, Clone, Copy, Default, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize,
63)]
64#[serde(rename_all = "snake_case")]
65pub enum CacheEvictionPolicy {
66    /// Evict the least recently used cacheable copy first.
67    #[default]
68    LeastRecentlyUsed,
69    /// Evict the least frequently used copy, using recency and unit id as ties.
70    LeastFrequentlyUsed,
71}
72
73/// A stable logical identifier for one independently managed offload unit.
74#[derive(Debug, Clone, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
75pub struct OffloadUnitId(String);
76
77impl OffloadUnitId {
78    /// Creates an identifier from a non-empty string.
79    pub fn new(id: impl Into<String>) -> Result<Self, OffloadError> {
80        let id = id.into();
81        if id.trim().is_empty() {
82            Err(OffloadError::EmptyUnitId)
83        } else {
84            Ok(Self(id))
85        }
86    }
87
88    /// Returns the identifier as a string slice.
89    pub fn as_str(&self) -> &str {
90        &self.0
91    }
92}
93
94impl AsRef<str> for OffloadUnitId {
95    fn as_ref(&self) -> &str {
96        self.as_str()
97    }
98}
99
100impl fmt::Display for OffloadUnitId {
101    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
102        self.0.fmt(f)
103    }
104}
105
106impl<'de> Deserialize<'de> for OffloadUnitId {
107    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
108        Self::new(String::deserialize(deserializer)?).map_err(D::Error::custom)
109    }
110}
111
112/// Global limits and lookahead used when validating an explicit offload plan.
113#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize)]
114pub struct OffloadConfig {
115    device_budget_bytes: Option<u64>,
116    host_budget_bytes: Option<u64>,
117    prefetch_depth: usize,
118    eviction_policy: CacheEvictionPolicy,
119}
120
121impl OffloadConfig {
122    /// Creates a configuration with optional finite device and host budgets.
123    ///
124    /// A zero-byte budget is meaningful and forbids assigning non-empty units
125    /// to that tier. `prefetch_depth` must be nonzero.
126    pub fn new(
127        device_budget_bytes: Option<u64>,
128        host_budget_bytes: Option<u64>,
129        prefetch_depth: usize,
130    ) -> Result<Self, OffloadError> {
131        if prefetch_depth == 0 {
132            return Err(OffloadError::ZeroPrefetchDepth);
133        }
134        Ok(Self {
135            device_budget_bytes,
136            host_budget_bytes,
137            prefetch_depth,
138            eviction_policy: CacheEvictionPolicy::LeastRecentlyUsed,
139        })
140    }
141
142    /// Returns the finite device-tier budget, if configured.
143    pub const fn device_budget_bytes(self) -> Option<u64> {
144        self.device_budget_bytes
145    }
146
147    /// Returns the finite physical host-allocation budget, if configured.
148    pub const fn host_budget_bytes(self) -> Option<u64> {
149        self.host_budget_bytes
150    }
151
152    /// Returns the number of logical units the executor may prefetch ahead.
153    pub const fn prefetch_depth(self) -> usize {
154        self.prefetch_depth
155    }
156
157    /// Selects deterministic cache eviction without changing tier budgets.
158    pub const fn with_eviction_policy(mut self, policy: CacheEvictionPolicy) -> Self {
159        self.eviction_policy = policy;
160        self
161    }
162
163    /// Returns the configured cache eviction ordering.
164    pub const fn eviction_policy(self) -> CacheEvictionPolicy {
165        self.eviction_policy
166    }
167}
168
169impl Default for OffloadConfig {
170    fn default() -> Self {
171        Self {
172            device_budget_bytes: None,
173            host_budget_bytes: None,
174            prefetch_depth: 1,
175            eviction_policy: CacheEvictionPolicy::LeastRecentlyUsed,
176        }
177    }
178}
179
180#[derive(Deserialize)]
181struct SerializedOffloadConfig {
182    device_budget_bytes: Option<u64>,
183    host_budget_bytes: Option<u64>,
184    prefetch_depth: usize,
185    eviction_policy: CacheEvictionPolicy,
186}
187
188impl<'de> Deserialize<'de> for OffloadConfig {
189    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
190        let value = SerializedOffloadConfig::deserialize(deserializer)?;
191        Self::new(
192            value.device_budget_bytes,
193            value.host_budget_bytes,
194            value.prefetch_depth,
195        )
196        .map(|config| config.with_eviction_policy(value.eviction_policy))
197        .map_err(D::Error::custom)
198    }
199}
200
201/// One explicit logical unit assignment in an offload plan.
202#[derive(Debug, Clone, Eq, PartialEq, Serialize)]
203pub struct OffloadUnitSpec {
204    id: OffloadUnitId,
205    bytes: u64,
206    policy: ResidencyPolicy,
207    tier: MemoryTier,
208}
209
210impl OffloadUnitSpec {
211    /// Creates and validates one explicit assignment.
212    pub fn new(
213        id: OffloadUnitId,
214        bytes: u64,
215        policy: ResidencyPolicy,
216        tier: MemoryTier,
217    ) -> Result<Self, OffloadError> {
218        if bytes == 0 {
219            return Err(OffloadError::ZeroSizedUnit { id });
220        }
221        if policy == ResidencyPolicy::Pinned && tier == MemoryTier::Disk {
222            return Err(OffloadError::ContradictoryAssignment {
223                id,
224                policy,
225                tier,
226                reason: "pinned units must be assigned to a resident memory tier",
227            });
228        }
229        Ok(Self {
230            id,
231            bytes,
232            policy,
233            tier,
234        })
235    }
236
237    /// Returns the logical unit identifier.
238    pub fn id(&self) -> &OffloadUnitId {
239        &self.id
240    }
241
242    /// Returns the planned unit size in bytes.
243    pub const fn bytes(&self) -> u64 {
244        self.bytes
245    }
246
247    /// Returns the planned residency policy.
248    pub const fn policy(&self) -> ResidencyPolicy {
249        self.policy
250    }
251
252    /// Returns the explicitly assigned tier.
253    pub const fn tier(&self) -> MemoryTier {
254        self.tier
255    }
256}
257
258#[derive(Deserialize)]
259struct SerializedOffloadUnitSpec {
260    id: OffloadUnitId,
261    bytes: u64,
262    policy: ResidencyPolicy,
263    tier: MemoryTier,
264}
265
266impl<'de> Deserialize<'de> for OffloadUnitSpec {
267    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
268        let value = SerializedOffloadUnitSpec::deserialize(deserializer)?;
269        Self::new(value.id, value.bytes, value.policy, value.tier).map_err(D::Error::custom)
270    }
271}
272
273/// Byte totals indexed by memory tier.
274#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
275pub struct TierByteTotals {
276    device: u64,
277    host: u64,
278    disk: u64,
279}
280
281/// Current or peak logical resident-unit counts by tier.
282#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
283pub struct TierUnitTotals {
284    device: usize,
285    host: usize,
286    disk: usize,
287}
288
289impl TierUnitTotals {
290    /// Creates explicit device, host, and disk unit totals.
291    pub const fn new(device: usize, host: usize, disk: usize) -> Self {
292        Self { device, host, disk }
293    }
294
295    /// Returns the unit total for `tier`.
296    pub const fn get(self, tier: MemoryTier) -> usize {
297        match tier {
298            MemoryTier::Device => self.device,
299            MemoryTier::Host => self.host,
300            MemoryTier::Disk => self.disk,
301        }
302    }
303
304    fn set(&mut self, tier: MemoryTier, units: usize) {
305        match tier {
306            MemoryTier::Device => self.device = units,
307            MemoryTier::Host => self.host = units,
308            MemoryTier::Disk => self.disk = units,
309        }
310    }
311}
312
313impl TierByteTotals {
314    /// Creates explicit device, host, and disk byte totals.
315    pub const fn new(device: u64, host: u64, disk: u64) -> Self {
316        Self { device, host, disk }
317    }
318
319    /// Returns the byte total for `tier`.
320    pub const fn get(self, tier: MemoryTier) -> u64 {
321        match tier {
322            MemoryTier::Device => self.device,
323            MemoryTier::Host => self.host,
324            MemoryTier::Disk => self.disk,
325        }
326    }
327
328    fn set(&mut self, tier: MemoryTier, bytes: u64) {
329        match tier {
330            MemoryTier::Device => self.device = bytes,
331            MemoryTier::Host => self.host = bytes,
332            MemoryTier::Disk => self.disk = bytes,
333        }
334    }
335
336    fn checked_add(&mut self, tier: MemoryTier, bytes: u64) -> Result<(), OffloadError> {
337        let total = self
338            .get(tier)
339            .checked_add(bytes)
340            .ok_or(OffloadError::ByteTotalOverflow { tier })?;
341        self.set(tier, total);
342        Ok(())
343    }
344}
345
346/// A deterministic, validated explicit offload plan.
347#[derive(Debug, Clone, Eq, PartialEq, Serialize)]
348pub struct OffloadPlan {
349    schema_version: u32,
350    config: OffloadConfig,
351    units: Vec<OffloadUnitSpec>,
352    #[serde(skip)]
353    planned_bytes: TierByteTotals,
354}
355
356impl OffloadPlan {
357    /// Validates explicit assignments and sorts them by logical identifier.
358    ///
359    /// This constructor does not materialize tensors or choose assignments.
360    pub fn new(
361        config: OffloadConfig,
362        units: impl IntoIterator<Item = OffloadUnitSpec>,
363    ) -> Result<Self, OffloadError> {
364        let mut units = units.into_iter().collect::<Vec<_>>();
365        units.sort_by(|left, right| left.id.cmp(&right.id));
366
367        if let Some(pair) = units.windows(2).find(|pair| pair[0].id == pair[1].id) {
368            return Err(OffloadError::DuplicateUnitId {
369                id: pair[0].id.clone(),
370            });
371        }
372
373        let mut planned_bytes = TierByteTotals::default();
374        for unit in &units {
375            planned_bytes.checked_add(unit.tier, unit.bytes)?;
376        }
377
378        validate_budget(
379            MemoryTier::Device,
380            planned_bytes.device,
381            config.device_budget_bytes,
382        )?;
383        validate_budget(
384            MemoryTier::Host,
385            planned_bytes.host,
386            config.host_budget_bytes,
387        )?;
388
389        Ok(Self {
390            schema_version: OFFLOAD_PLAN_SCHEMA_VERSION,
391            config,
392            units,
393            planned_bytes,
394        })
395    }
396
397    /// Returns the stable serialized schema version.
398    pub const fn schema_version(&self) -> u32 {
399        self.schema_version
400    }
401
402    /// Revalidates every invariant represented by this plan.
403    pub fn validate(&self) -> Result<(), OffloadError> {
404        Self::new(self.config, self.units.clone()).map(|_| ())
405    }
406
407    /// Returns the configuration used to validate this plan.
408    pub const fn config(&self) -> OffloadConfig {
409        self.config
410    }
411
412    /// Returns assignments in stable logical-identifier order.
413    pub fn units(&self) -> &[OffloadUnitSpec] {
414        &self.units
415    }
416
417    /// Looks up a unit by its logical identifier.
418    pub fn unit(&self, id: &OffloadUnitId) -> Option<&OffloadUnitSpec> {
419        self.units
420            .binary_search_by(|unit| unit.id.cmp(id))
421            .ok()
422            .map(|index| &self.units[index])
423    }
424
425    /// Returns checked planned byte totals for every tier.
426    pub const fn planned_bytes(&self) -> TierByteTotals {
427        self.planned_bytes
428    }
429}
430
431#[derive(Deserialize)]
432struct SerializedOffloadPlan {
433    schema_version: u32,
434    config: OffloadConfig,
435    units: Vec<OffloadUnitSpec>,
436}
437
438impl<'de> Deserialize<'de> for OffloadPlan {
439    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
440        let value = SerializedOffloadPlan::deserialize(deserializer)?;
441        if value.schema_version != OFFLOAD_PLAN_SCHEMA_VERSION {
442            return Err(D::Error::custom(OffloadError::UnsupportedSchemaVersion(
443                value.schema_version,
444            )));
445        }
446        Self::new(value.config, value.units).map_err(D::Error::custom)
447    }
448}
449
450fn validate_budget(
451    tier: MemoryTier,
452    planned_bytes: u64,
453    budget_bytes: Option<u64>,
454) -> Result<(), OffloadError> {
455    if let Some(budget_bytes) = budget_bytes {
456        if planned_bytes > budget_bytes {
457            return Err(OffloadError::BudgetExceeded {
458                tier,
459                planned_bytes,
460                budget_bytes,
461            });
462        }
463    }
464    Ok(())
465}
466
467/// Structured validation failures for offload contracts.
468#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
469pub enum OffloadError {
470    /// A serialized plan used an unsupported schema version.
471    #[error("unsupported offload plan schema version {0}")]
472    UnsupportedSchemaVersion(u32),
473    /// A logical identifier was empty or whitespace-only.
474    #[error("offload unit identifiers must not be empty")]
475    EmptyUnitId,
476    /// A unit had no bytes to manage.
477    #[error("offload unit {id} must contain at least one byte")]
478    ZeroSizedUnit {
479        /// The invalid unit identifier.
480        id: OffloadUnitId,
481    },
482    /// More than one unit used the same stable identifier.
483    #[error("duplicate offload unit identifier: {id}")]
484    DuplicateUnitId {
485        /// The duplicated identifier.
486        id: OffloadUnitId,
487    },
488    /// Summing unit sizes overflowed the stable byte counter.
489    #[error("planned byte total overflowed for the {tier:?} tier")]
490    ByteTotalOverflow {
491        /// The tier whose total overflowed.
492        tier: MemoryTier,
493    },
494    /// Explicit assignments exceeded a configured finite budget.
495    #[error(
496        "planned {planned_bytes} bytes for the {tier:?} tier exceed its {budget_bytes}-byte budget"
497    )]
498    BudgetExceeded {
499        /// The over-budget tier.
500        tier: MemoryTier,
501        /// The checked planned total.
502        planned_bytes: u64,
503        /// The configured finite budget.
504        budget_bytes: u64,
505    },
506    /// A policy and tier assignment had incompatible meanings.
507    #[error("offload unit {id} has contradictory {policy:?}/{tier:?} assignment: {reason}")]
508    ContradictoryAssignment {
509        /// The invalid unit identifier.
510        id: OffloadUnitId,
511        /// The requested policy.
512        policy: ResidencyPolicy,
513        /// The requested tier.
514        tier: MemoryTier,
515        /// A stable explanation of the contradiction.
516        reason: &'static str,
517    },
518    /// Prefetching was configured with a meaningless zero-unit lookahead.
519    #[error("offload prefetch depth must be nonzero")]
520    ZeroPrefetchDepth,
521}
522
523/// A strongly typed transfer direction between two distinct tiers.
524#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
525#[serde(rename_all = "snake_case")]
526pub enum TransferDirection {
527    /// Device memory to host memory.
528    DeviceToHost,
529    /// Device memory to disk.
530    DeviceToDisk,
531    /// Host memory to device memory.
532    HostToDevice,
533    /// Host memory to disk.
534    HostToDisk,
535    /// Disk to device memory.
536    DiskToDevice,
537    /// Disk to host memory.
538    DiskToHost,
539}
540
541impl TransferDirection {
542    /// All directions in stable reporting order.
543    pub const ALL: [Self; 6] = [
544        Self::DeviceToHost,
545        Self::DeviceToDisk,
546        Self::HostToDevice,
547        Self::HostToDisk,
548        Self::DiskToDevice,
549        Self::DiskToHost,
550    ];
551
552    /// Returns the source tier.
553    pub const fn source(self) -> MemoryTier {
554        match self {
555            Self::DeviceToHost | Self::DeviceToDisk => MemoryTier::Device,
556            Self::HostToDevice | Self::HostToDisk => MemoryTier::Host,
557            Self::DiskToDevice | Self::DiskToHost => MemoryTier::Disk,
558        }
559    }
560
561    /// Returns the destination tier.
562    pub const fn destination(self) -> MemoryTier {
563        match self {
564            Self::DeviceToHost | Self::DiskToHost => MemoryTier::Host,
565            Self::DeviceToDisk | Self::HostToDisk => MemoryTier::Disk,
566            Self::HostToDevice | Self::DiskToDevice => MemoryTier::Device,
567        }
568    }
569
570    const fn index(self) -> usize {
571        match self {
572            Self::DeviceToHost => 0,
573            Self::DeviceToDisk => 1,
574            Self::HostToDevice => 2,
575            Self::HostToDisk => 3,
576            Self::DiskToDevice => 4,
577            Self::DiskToHost => 5,
578        }
579    }
580}
581
582/// Accumulated transfer observations for one direction.
583#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
584pub struct TransferMetrics {
585    count: u64,
586    bytes: u64,
587    duration: Duration,
588}
589
590impl TransferMetrics {
591    /// Returns the number of recorded transfers.
592    pub const fn count(self) -> u64 {
593        self.count
594    }
595
596    /// Returns the number of recorded bytes.
597    pub const fn bytes(self) -> u64 {
598        self.bytes
599    }
600
601    /// Returns the accumulated transfer duration.
602    pub const fn duration(self) -> Duration {
603        self.duration
604    }
605
606    fn record(&mut self, bytes: u64, duration: Duration) {
607        self.count = self.count.saturating_add(1);
608        self.bytes = self.bytes.saturating_add(bytes);
609        self.duration = self.duration.saturating_add(duration);
610    }
611}
612
613/// The result of one completed prefetch request.
614#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
615#[serde(rename_all = "snake_case")]
616pub enum PrefetchOutcome {
617    /// The requested unit was already available at the required tier.
618    Hit,
619    /// The request required a transfer or load.
620    Miss,
621}
622
623/// Accumulated prefetch observations.
624#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
625pub struct PrefetchMetrics {
626    requests: u64,
627    hits: u64,
628    misses: u64,
629    stalls: u64,
630    stall_duration: Duration,
631}
632
633impl PrefetchMetrics {
634    /// Returns the number of completed prefetch requests.
635    pub const fn requests(self) -> u64 {
636        self.requests
637    }
638
639    /// Returns the number of prefetch hits.
640    pub const fn hits(self) -> u64 {
641        self.hits
642    }
643
644    /// Returns the number of prefetch misses.
645    pub const fn misses(self) -> u64 {
646        self.misses
647    }
648
649    /// Returns the number of demand waits attributed to prefetching.
650    pub const fn stalls(self) -> u64 {
651        self.stalls
652    }
653
654    /// Returns the accumulated prefetch stall duration.
655    pub const fn stall_duration(self) -> Duration {
656        self.stall_duration
657    }
658}
659
660/// Accumulated eviction observations.
661#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
662pub struct EvictionMetrics {
663    count: u64,
664    bytes: u64,
665}
666
667impl EvictionMetrics {
668    /// Returns the number of recorded evictions.
669    pub const fn count(self) -> u64 {
670        self.count
671    }
672
673    /// Returns the number of recorded evicted bytes.
674    pub const fn bytes(self) -> u64 {
675        self.bytes
676    }
677}
678
679/// A point-in-time sample of backend-managed allocator memory.
680#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
681pub struct AllocatorMemoryMetrics {
682    active_bytes: u64,
683    cached_bytes: u64,
684    peak_bytes: u64,
685}
686
687impl AllocatorMemoryMetrics {
688    /// Creates an explicit backend allocator sample.
689    pub const fn new(active_bytes: u64, cached_bytes: u64, peak_bytes: u64) -> Self {
690        Self {
691            active_bytes,
692            cached_bytes,
693            peak_bytes,
694        }
695    }
696
697    /// Returns active backend-managed bytes.
698    pub const fn active_bytes(self) -> u64 {
699        self.active_bytes
700    }
701
702    /// Returns bytes retained by the backend allocator cache.
703    pub const fn cached_bytes(self) -> u64 {
704        self.cached_bytes
705    }
706
707    /// Returns peak active backend-managed bytes.
708    pub const fn peak_bytes(self) -> u64 {
709        self.peak_bytes
710    }
711}
712
713/// Optional process-level memory and page-fault observations.
714///
715/// Individual values are absent when they cannot be obtained safely on the
716/// current platform. The built-in sampler currently reads Linux `/proc`; it
717/// makes no availability guarantee on Apple or Windows targets.
718#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
719pub struct ProcessMetrics {
720    rss_bytes: Option<u64>,
721    minor_page_faults: Option<u64>,
722    major_page_faults: Option<u64>,
723}
724
725impl ProcessMetrics {
726    /// Creates an explicit process sample.
727    pub const fn new(
728        rss_bytes: Option<u64>,
729        minor_page_faults: Option<u64>,
730        major_page_faults: Option<u64>,
731    ) -> Self {
732        Self {
733            rss_bytes,
734            minor_page_faults,
735            major_page_faults,
736        }
737    }
738
739    /// Returns resident-set bytes when available.
740    pub const fn rss_bytes(self) -> Option<u64> {
741        self.rss_bytes
742    }
743
744    /// Returns minor page faults when available.
745    pub const fn minor_page_faults(self) -> Option<u64> {
746        self.minor_page_faults
747    }
748
749    /// Returns major page faults when available.
750    pub const fn major_page_faults(self) -> Option<u64> {
751        self.major_page_faults
752    }
753}
754
755/// Samples optional process metrics without adding a platform runtime dependency.
756pub fn sample_process_metrics() -> ProcessMetrics {
757    platform_process_metrics()
758}
759
760#[cfg(target_os = "linux")]
761fn platform_process_metrics() -> ProcessMetrics {
762    let rss_bytes = std::fs::read_to_string("/proc/self/status")
763        .ok()
764        .and_then(|status| {
765            status.lines().find_map(|line| {
766                let value = line.strip_prefix("VmRSS:")?.trim();
767                let kibibytes = value.strip_suffix("kB")?.trim().parse::<u64>().ok()?;
768                kibibytes.checked_mul(1024)
769            })
770        });
771
772    let faults = std::fs::read_to_string("/proc/self/stat")
773        .ok()
774        .and_then(|stat| {
775            // The parenthesized command name may contain spaces. Fields after
776            // its final ')' begin at process-stat field 3 (state).
777            let fields = stat.get(stat.rfind(')')? + 1..)?.split_whitespace();
778            let fields = fields.collect::<Vec<_>>();
779            Some((fields.get(7)?.parse().ok()?, fields.get(9)?.parse().ok()?))
780        });
781
782    ProcessMetrics::new(
783        rss_bytes,
784        faults.map(|value| value.0),
785        faults.map(|value| value.1),
786    )
787}
788
789#[cfg(not(target_os = "linux"))]
790fn platform_process_metrics() -> ProcessMetrics {
791    ProcessMetrics::default()
792}
793
794/// Mutable, single-threaded offload telemetry collector.
795///
796/// Updates use saturating arithmetic for monotonic counters and durations.
797/// Resident bytes are set explicitly, and setting a new value updates the
798/// corresponding peak. Wrap this value in a mutex if multiple threads need to
799/// record into one collector.
800#[derive(Debug, Default, Clone, Serialize, Deserialize)]
801pub struct OffloadTelemetry {
802    planned_bytes: TierByteTotals,
803    resident_bytes: TierByteTotals,
804    peak_resident_bytes: TierByteTotals,
805    resident_units: TierUnitTotals,
806    peak_resident_units: TierUnitTotals,
807    transfers: [TransferMetrics; 6],
808    prefetch: PrefetchMetrics,
809    tier_prefetch: [PrefetchMetrics; 3],
810    evictions: EvictionMetrics,
811    tier_evictions: [EvictionMetrics; 3],
812    allocator_memory: Option<AllocatorMemoryMetrics>,
813    process: ProcessMetrics,
814    process_sampled: bool,
815}
816
817impl OffloadTelemetry {
818    /// Creates a collector initialized with a validated plan's byte totals.
819    pub fn from_plan(plan: &OffloadPlan) -> Self {
820        Self {
821            planned_bytes: plan.planned_bytes,
822            ..Self::default()
823        }
824    }
825
826    /// Replaces the planned byte totals recorded by this collector.
827    pub fn set_planned_bytes(&mut self, planned_bytes: TierByteTotals) {
828        self.planned_bytes = planned_bytes;
829    }
830
831    /// Sets current resident bytes and updates the peak for `tier`.
832    pub fn set_resident_bytes(&mut self, tier: MemoryTier, bytes: u64) {
833        self.resident_bytes.set(tier, bytes);
834        if bytes > self.peak_resident_bytes.get(tier) {
835            self.peak_resident_bytes.set(tier, bytes);
836        }
837    }
838
839    /// Sets current resident units and updates the peak for `tier`.
840    pub fn set_resident_units(&mut self, tier: MemoryTier, units: usize) {
841        self.resident_units.set(tier, units);
842        if units > self.peak_resident_units.get(tier) {
843            self.peak_resident_units.set(tier, units);
844        }
845    }
846
847    /// Records one completed transfer using saturating counter updates.
848    pub fn record_transfer(
849        &mut self,
850        direction: TransferDirection,
851        bytes: u64,
852        duration: Duration,
853    ) {
854        self.transfers[direction.index()].record(bytes, duration);
855    }
856
857    /// Records one completed prefetch request and its outcome.
858    pub fn record_prefetch(&mut self, outcome: PrefetchOutcome) {
859        self.prefetch.requests = self.prefetch.requests.saturating_add(1);
860        match outcome {
861            PrefetchOutcome::Hit => {
862                self.prefetch.hits = self.prefetch.hits.saturating_add(1);
863            }
864            PrefetchOutcome::Miss => {
865                self.prefetch.misses = self.prefetch.misses.saturating_add(1);
866            }
867        }
868    }
869
870    /// Records a cache request both globally and for its target tier.
871    pub fn record_tier_prefetch(&mut self, tier: MemoryTier, outcome: PrefetchOutcome) {
872        self.record_prefetch(outcome);
873        let metrics = &mut self.tier_prefetch[tier.index()];
874        metrics.requests = metrics.requests.saturating_add(1);
875        match outcome {
876            PrefetchOutcome::Hit => metrics.hits = metrics.hits.saturating_add(1),
877            PrefetchOutcome::Miss => metrics.misses = metrics.misses.saturating_add(1),
878        }
879    }
880
881    /// Records a demand stall while waiting for a prefetched unit.
882    pub fn record_prefetch_stall(&mut self, duration: Duration) {
883        self.prefetch.stalls = self.prefetch.stalls.saturating_add(1);
884        self.prefetch.stall_duration = self.prefetch.stall_duration.saturating_add(duration);
885    }
886
887    /// Records one eviction using saturating counter updates.
888    pub fn record_eviction(&mut self, bytes: u64) {
889        self.evictions.count = self.evictions.count.saturating_add(1);
890        self.evictions.bytes = self.evictions.bytes.saturating_add(bytes);
891    }
892
893    /// Records an eviction both globally and for its source tier.
894    pub fn record_tier_eviction(&mut self, tier: MemoryTier, bytes: u64) {
895        self.record_eviction(bytes);
896        let metrics = &mut self.tier_evictions[tier.index()];
897        metrics.count = metrics.count.saturating_add(1);
898        metrics.bytes = metrics.bytes.saturating_add(bytes);
899    }
900
901    /// Records an allocator sample supplied by the selected backend.
902    pub fn record_allocator_memory(&mut self, metrics: AllocatorMemoryMetrics) {
903        self.allocator_memory = Some(metrics);
904    }
905
906    /// Records an externally obtained process sample.
907    pub fn record_process_metrics(&mut self, metrics: ProcessMetrics) {
908        self.process = metrics;
909        self.process_sampled = true;
910    }
911
912    /// Updates process observations using the built-in optional sampler.
913    pub fn sample_process_metrics(&mut self) {
914        self.process = sample_process_metrics();
915        self.process_sampled = true;
916    }
917
918    /// Returns an immutable point-in-time report.
919    pub fn snapshot(&self) -> OffloadReport {
920        OffloadReport {
921            planned_bytes: self.planned_bytes,
922            resident_bytes: self.resident_bytes,
923            peak_resident_bytes: self.peak_resident_bytes,
924            resident_units: self.resident_units,
925            peak_resident_units: self.peak_resident_units,
926            transfers: self.transfers,
927            prefetch: self.prefetch,
928            tier_prefetch: self.tier_prefetch,
929            evictions: self.evictions,
930            tier_evictions: self.tier_evictions,
931            allocator_memory: self.allocator_memory,
932            process: self.process,
933            process_sampled: self.process_sampled,
934        }
935    }
936
937    /// Clears all configuration and observations, including resident peaks.
938    pub fn reset(&mut self) {
939        *self = Self::default();
940    }
941}
942
943/// Immutable point-in-time offload telemetry report.
944#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
945pub struct OffloadReport {
946    planned_bytes: TierByteTotals,
947    resident_bytes: TierByteTotals,
948    peak_resident_bytes: TierByteTotals,
949    resident_units: TierUnitTotals,
950    peak_resident_units: TierUnitTotals,
951    transfers: [TransferMetrics; 6],
952    prefetch: PrefetchMetrics,
953    tier_prefetch: [PrefetchMetrics; 3],
954    evictions: EvictionMetrics,
955    tier_evictions: [EvictionMetrics; 3],
956    allocator_memory: Option<AllocatorMemoryMetrics>,
957    process: ProcessMetrics,
958    process_sampled: bool,
959}
960
961impl OffloadReport {
962    /// Returns planned bytes per tier.
963    pub const fn planned_bytes(&self) -> TierByteTotals {
964        self.planned_bytes
965    }
966
967    /// Returns current resident bytes per tier.
968    pub const fn resident_bytes(&self) -> TierByteTotals {
969        self.resident_bytes
970    }
971
972    /// Returns peak resident bytes per tier.
973    pub const fn peak_resident_bytes(&self) -> TierByteTotals {
974        self.peak_resident_bytes
975    }
976
977    /// Returns current resident-unit counts per tier.
978    pub const fn resident_units(&self) -> TierUnitTotals {
979        self.resident_units
980    }
981
982    /// Returns peak resident-unit counts per tier.
983    pub const fn peak_resident_units(&self) -> TierUnitTotals {
984        self.peak_resident_units
985    }
986
987    /// Returns accumulated metrics for one transfer direction.
988    pub const fn transfer(&self, direction: TransferDirection) -> TransferMetrics {
989        self.transfers[direction.index()]
990    }
991
992    /// Returns accumulated prefetch metrics.
993    pub const fn prefetch(&self) -> PrefetchMetrics {
994        self.prefetch
995    }
996
997    /// Returns cache request metrics for one target tier.
998    pub const fn tier_prefetch(&self, tier: MemoryTier) -> PrefetchMetrics {
999        self.tier_prefetch[tier.index()]
1000    }
1001
1002    /// Returns accumulated eviction metrics.
1003    pub const fn evictions(&self) -> EvictionMetrics {
1004        self.evictions
1005    }
1006
1007    /// Returns eviction metrics for one source tier.
1008    pub const fn tier_evictions(&self, tier: MemoryTier) -> EvictionMetrics {
1009        self.tier_evictions[tier.index()]
1010    }
1011
1012    /// Returns the latest backend allocator sample, if one was recorded.
1013    pub const fn allocator_memory(&self) -> Option<AllocatorMemoryMetrics> {
1014        self.allocator_memory
1015    }
1016
1017    /// Returns the latest optional process sample.
1018    pub const fn process_metrics(&self) -> ProcessMetrics {
1019        self.process
1020    }
1021
1022    /// Returns whether process sampling was requested, including unsupported platforms.
1023    pub const fn process_sampled(&self) -> bool {
1024        self.process_sampled
1025    }
1026}
1027
1028/// Logical state of one materialized tier copy.
1029#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
1030pub struct ResidentCopyStatus {
1031    bytes: u64,
1032    pins: u64,
1033    in_flight: Option<u64>,
1034}
1035
1036impl ResidentCopyStatus {
1037    /// Charged physical bytes.
1038    pub const fn bytes(self) -> u64 {
1039        self.bytes
1040    }
1041
1042    /// Active ownership leases preventing eviction.
1043    pub const fn pins(self) -> u64 {
1044        self.pins
1045    }
1046
1047    /// Exact transfer generation awaiting disposition.
1048    pub const fn in_flight(self) -> Option<u64> {
1049        self.in_flight
1050    }
1051}
1052
1053/// Point-in-time logical residency for one planned unit.
1054#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
1055pub struct UnitResidencyReport {
1056    id: OffloadUnitId,
1057    planned_tier: MemoryTier,
1058    policy: ResidencyPolicy,
1059    expected_bytes: u64,
1060    host_allocated_bytes: u64,
1061    device_allocated_bytes: u64,
1062    host_resident: bool,
1063    device_resident: bool,
1064    host_pins: u64,
1065    device_pins: u64,
1066    active_window: bool,
1067}
1068
1069impl UnitResidencyReport {
1070    /// Stable unit identifier.
1071    pub fn id(&self) -> &OffloadUnitId {
1072        &self.id
1073    }
1074    /// Initial tier selected by the plan.
1075    pub const fn planned_tier(&self) -> MemoryTier {
1076        self.planned_tier
1077    }
1078    /// Operational lifetime policy.
1079    pub const fn policy(&self) -> ResidencyPolicy {
1080        self.policy
1081    }
1082    /// Planned logical bytes.
1083    pub const fn expected_bytes(&self) -> u64 {
1084        self.expected_bytes
1085    }
1086    /// Charged host allocation capacity.
1087    pub const fn host_allocated_bytes(&self) -> u64 {
1088        self.host_allocated_bytes
1089    }
1090    /// Charged execution-memory bytes.
1091    pub const fn device_allocated_bytes(&self) -> u64 {
1092        self.device_allocated_bytes
1093    }
1094    /// Whether a host copy is logically resident.
1095    pub const fn host_resident(&self) -> bool {
1096        self.host_resident
1097    }
1098    /// Whether an execution-memory copy is logically resident.
1099    pub const fn device_resident(&self) -> bool {
1100        self.device_resident
1101    }
1102    /// Active host-copy leases.
1103    pub const fn host_pins(&self) -> u64 {
1104        self.host_pins
1105    }
1106    /// Active execution-copy leases.
1107    pub const fn device_pins(&self) -> u64 {
1108        self.device_pins
1109    }
1110    /// Whether any named execution window protects the unit.
1111    pub const fn active_window(&self) -> bool {
1112        self.active_window
1113    }
1114}
1115
1116/// A copy the backend must release after a ledger transition.
1117#[derive(Debug, Clone, Eq, PartialEq)]
1118pub struct EvictedResidencyCopy {
1119    /// Logical unit whose backend storage must be released.
1120    pub id: OffloadUnitId,
1121    /// Tier of the released backend storage.
1122    pub tier: MemoryTier,
1123    /// Bytes removed from ledger accounting.
1124    pub bytes: u64,
1125}
1126
1127/// One unit preventing an automatic capacity reservation.
1128#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
1129pub struct ResidencyBlocker {
1130    /// Stable logical identifier.
1131    pub id: OffloadUnitId,
1132    /// Whether lifetime policy forbids eviction.
1133    pub pinned: bool,
1134    /// Active ownership leases.
1135    pub in_use: u64,
1136    /// Whether an execution window protects the unit.
1137    pub active_window: bool,
1138    /// Whether the current atomic request protects the unit.
1139    pub request_protected: bool,
1140}
1141
1142#[derive(Debug, Clone, Copy, Eq, PartialEq)]
1143enum CopyLifecycle {
1144    Reserved,
1145    Resident,
1146}
1147
1148#[derive(Debug, Clone, Copy, Eq, PartialEq)]
1149struct LedgerCopy {
1150    lifecycle: CopyLifecycle,
1151    bytes: u64,
1152    pins: u64,
1153    last_used: u64,
1154    frequency: u64,
1155    in_flight: Option<u64>,
1156}
1157
1158impl LedgerCopy {
1159    fn status(self) -> Option<ResidentCopyStatus> {
1160        (self.lifecycle == CopyLifecycle::Resident).then_some(ResidentCopyStatus {
1161            bytes: self.bytes,
1162            pins: self.pins,
1163            in_flight: self.in_flight,
1164        })
1165    }
1166}
1167
1168#[derive(Debug, Clone)]
1169struct LedgerUnit {
1170    spec: OffloadUnitSpec,
1171    host: Option<LedgerCopy>,
1172    device: Option<LedgerCopy>,
1173}
1174
1175impl LedgerUnit {
1176    fn copy(&self, tier: MemoryTier) -> Option<&LedgerCopy> {
1177        match tier {
1178            MemoryTier::Host => self.host.as_ref(),
1179            MemoryTier::Device => self.device.as_ref(),
1180            MemoryTier::Disk => None,
1181        }
1182    }
1183
1184    fn copy_mut(&mut self, tier: MemoryTier) -> Option<&mut LedgerCopy> {
1185        match tier {
1186            MemoryTier::Host => self.host.as_mut(),
1187            MemoryTier::Device => self.device.as_mut(),
1188            MemoryTier::Disk => None,
1189        }
1190    }
1191
1192    fn slot_mut(&mut self, tier: MemoryTier) -> Option<&mut Option<LedgerCopy>> {
1193        match tier {
1194            MemoryTier::Host => Some(&mut self.host),
1195            MemoryTier::Device => Some(&mut self.device),
1196            MemoryTier::Disk => None,
1197        }
1198    }
1199}
1200
1201/// Backend-neutral ownership and capacity state for one residency plan.
1202///
1203/// Backends mirror each resident ledger copy with concrete storage. Ledger
1204/// transitions return every copy whose storage must be released, so tensor or
1205/// buffer destruction never needs to be hidden behind an untyped callback.
1206#[derive(Debug)]
1207pub struct ResidencyLedger {
1208    plan: OffloadPlan,
1209    units: BTreeMap<OffloadUnitId, LedgerUnit>,
1210    group_windows: BTreeMap<(String, MemoryTier), BTreeSet<OffloadUnitId>>,
1211    active_windows: BTreeMap<MemoryTier, BTreeSet<OffloadUnitId>>,
1212    telemetry: OffloadTelemetry,
1213    resident_bytes: TierByteTotals,
1214    tick: u64,
1215    next_transfer_generation: u64,
1216    initialized: bool,
1217}
1218
1219impl ResidencyLedger {
1220    /// Creates empty ownership state for every unit in a validated plan.
1221    pub fn new(plan: OffloadPlan) -> Self {
1222        let units = plan
1223            .units()
1224            .iter()
1225            .cloned()
1226            .map(|spec| {
1227                (
1228                    spec.id().clone(),
1229                    LedgerUnit {
1230                        spec,
1231                        host: None,
1232                        device: None,
1233                    },
1234                )
1235            })
1236            .collect();
1237        let telemetry = OffloadTelemetry::from_plan(&plan);
1238        Self {
1239            plan,
1240            units,
1241            group_windows: BTreeMap::new(),
1242            active_windows: BTreeMap::new(),
1243            telemetry,
1244            resident_bytes: TierByteTotals::default(),
1245            tick: 0,
1246            next_transfer_generation: 1,
1247            initialized: false,
1248        }
1249    }
1250
1251    /// Validated plan governing this ledger.
1252    pub const fn plan(&self) -> &OffloadPlan {
1253        &self.plan
1254    }
1255
1256    /// Whether initial planned materialization has completed.
1257    pub const fn initialized(&self) -> bool {
1258        self.initialized
1259    }
1260
1261    /// Marks initial planned materialization complete.
1262    pub fn mark_initialized(&mut self) {
1263        self.initialized = true;
1264    }
1265
1266    /// Fails unless initial planned materialization completed.
1267    pub fn require_initialized(&self) -> Result<(), ResidencyLedgerError> {
1268        self.initialized
1269            .then_some(())
1270            .ok_or(ResidencyLedgerError::NotInitialized)
1271    }
1272
1273    /// Returns whether the plan contains a unit.
1274    pub fn contains(&self, id: &OffloadUnitId) -> bool {
1275        self.units.contains_key(id)
1276    }
1277
1278    /// Returns one planned unit specification.
1279    pub fn spec(&self, id: &OffloadUnitId) -> Result<&OffloadUnitSpec, ResidencyLedgerError> {
1280        self.units
1281            .get(id)
1282            .map(|unit| &unit.spec)
1283            .ok_or_else(|| ResidencyLedgerError::UnknownUnit { id: id.clone() })
1284    }
1285
1286    /// Returns a resident copy's logical state.
1287    pub fn copy_status(
1288        &self,
1289        id: &OffloadUnitId,
1290        tier: MemoryTier,
1291    ) -> Result<Option<ResidentCopyStatus>, ResidencyLedgerError> {
1292        validate_ledger_tier(tier, "copy status")?;
1293        Ok(self
1294            .units
1295            .get(id)
1296            .ok_or_else(|| ResidencyLedgerError::UnknownUnit { id: id.clone() })?
1297            .copy(tier)
1298            .and_then(|copy| copy.status()))
1299    }
1300
1301    /// Returns whether a materialized copy exists, including in-flight copies.
1302    pub fn is_resident(
1303        &self,
1304        id: &OffloadUnitId,
1305        tier: MemoryTier,
1306    ) -> Result<bool, ResidencyLedgerError> {
1307        Ok(self.copy_status(id, tier)?.is_some())
1308    }
1309
1310    /// Validates one ordered, duplicate-free batch of known units.
1311    pub fn validate_batch(
1312        &self,
1313        ids: &[OffloadUnitId],
1314        tier: MemoryTier,
1315    ) -> Result<(), ResidencyLedgerError> {
1316        validate_ledger_tier(tier, "residency batch")?;
1317        let mut seen = BTreeSet::new();
1318        for id in ids {
1319            if !seen.insert(id) {
1320                return Err(ResidencyLedgerError::DuplicateBatchUnit);
1321            }
1322            self.spec(id)?;
1323        }
1324        Ok(())
1325    }
1326
1327    /// Reserves capacity for one missing copy and returns backend storage evictions.
1328    pub fn reserve_copy(
1329        &mut self,
1330        id: &OffloadUnitId,
1331        tier: MemoryTier,
1332        required_bytes: u64,
1333        protected: &BTreeSet<OffloadUnitId>,
1334    ) -> Result<Vec<EvictedResidencyCopy>, ResidencyLedgerError> {
1335        self.reserve_copies(&[(id.clone(), required_bytes)], tier, protected)
1336    }
1337
1338    /// Atomically reserves capacity for a batch of missing copies.
1339    ///
1340    /// Admission is planned before any ledger copy is removed or inserted. A
1341    /// failed request therefore leaves ownership and accounting unchanged.
1342    pub fn reserve_copies(
1343        &mut self,
1344        requests: &[(OffloadUnitId, u64)],
1345        tier: MemoryTier,
1346        protected: &BTreeSet<OffloadUnitId>,
1347    ) -> Result<Vec<EvictedResidencyCopy>, ResidencyLedgerError> {
1348        validate_ledger_tier(tier, "capacity reservation")?;
1349        let ids = requests
1350            .iter()
1351            .map(|(id, _)| id.clone())
1352            .collect::<Vec<_>>();
1353        self.validate_batch(&ids, tier)?;
1354        let mut total_required = 0u64;
1355        for (id, required_bytes) in requests {
1356            if *required_bytes == 0 {
1357                return Err(ResidencyLedgerError::ZeroReservation { id: id.clone() });
1358            }
1359            let unit = self
1360                .units
1361                .get(id)
1362                .ok_or_else(|| ResidencyLedgerError::UnknownUnit { id: id.clone() })?;
1363            if unit.copy(tier).is_some() {
1364                return Err(ResidencyLedgerError::CopyAlreadyExists {
1365                    id: id.clone(),
1366                    tier,
1367                });
1368            }
1369            total_required = total_required.checked_add(*required_bytes).ok_or(
1370                ResidencyLedgerError::ArithmeticOverflow {
1371                    context: "batch capacity reservation",
1372                },
1373            )?;
1374        }
1375
1376        if requests.is_empty() {
1377            return Ok(Vec::new());
1378        }
1379
1380        let mut victims = Vec::new();
1381        if let Some(budget) = self.budget(tier) {
1382            let needed = self.tier_bytes(tier).checked_add(total_required).ok_or(
1383                ResidencyLedgerError::ArithmeticOverflow {
1384                    context: "budget reservation",
1385                },
1386            )?;
1387            let required_release = needed.saturating_sub(budget);
1388            let mut releasable = 0u64;
1389            if required_release != 0 {
1390                for victim in self.eviction_candidates(tier, protected) {
1391                    let bytes = self
1392                        .units
1393                        .get(&victim)
1394                        .and_then(|unit| unit.copy(tier))
1395                        .ok_or_else(|| inconsistent(&victim, tier, "capacity eviction planning"))?
1396                        .bytes;
1397                    releasable = releasable.checked_add(bytes).ok_or(
1398                        ResidencyLedgerError::ArithmeticOverflow {
1399                            context: "eviction capacity planning",
1400                        },
1401                    )?;
1402                    victims.push(victim);
1403                    if releasable >= required_release {
1404                        break;
1405                    }
1406                }
1407                if releasable < required_release {
1408                    return Err(ResidencyLedgerError::BudgetExhausted {
1409                        requested: requests[0].0.clone(),
1410                        tier,
1411                        required_bytes: total_required,
1412                        budget_bytes: budget,
1413                        resident_bytes: self.tier_bytes(tier),
1414                        blocking_units: self.blockers(tier, protected),
1415                    });
1416                }
1417            }
1418        }
1419
1420        let mut evicted = Vec::with_capacity(victims.len());
1421        for victim in victims {
1422            evicted.push(self.remove_copy(&victim, tier, true)?);
1423        }
1424
1425        let charged = self.tier_bytes(tier).checked_add(total_required).ok_or(
1426            ResidencyLedgerError::ArithmeticOverflow {
1427                context: "resident byte reservation",
1428            },
1429        )?;
1430        self.set_tier_bytes(tier, charged);
1431        for (id, required_bytes) in requests {
1432            let tick = self.next_tick();
1433            *self
1434                .units
1435                .get_mut(id)
1436                .and_then(|unit| unit.slot_mut(tier))
1437                .ok_or_else(|| inconsistent(id, tier, "reservation insertion"))? =
1438                Some(LedgerCopy {
1439                    lifecycle: CopyLifecycle::Reserved,
1440                    bytes: *required_bytes,
1441                    pins: 0,
1442                    last_used: tick,
1443                    frequency: 0,
1444                    in_flight: None,
1445                });
1446        }
1447        self.update_resident_telemetry(tier);
1448        Ok(evicted)
1449    }
1450
1451    /// Publishes backend storage into an existing reservation.
1452    pub fn publish_reserved(
1453        &mut self,
1454        id: &OffloadUnitId,
1455        tier: MemoryTier,
1456        actual_bytes: u64,
1457        in_flight: Option<u64>,
1458    ) -> Result<(), ResidencyLedgerError> {
1459        validate_ledger_tier(tier, "copy publication")?;
1460        let reserved = *self
1461            .units
1462            .get(id)
1463            .and_then(|unit| unit.copy(tier))
1464            .ok_or_else(|| inconsistent(id, tier, "publication lookup"))?;
1465        if reserved.lifecycle != CopyLifecycle::Reserved {
1466            return Err(inconsistent(id, tier, "publication lifecycle"));
1467        }
1468        if actual_bytes == 0 {
1469            return Err(ResidencyLedgerError::ZeroPublication {
1470                id: id.clone(),
1471                tier,
1472            });
1473        }
1474        if actual_bytes > reserved.bytes {
1475            return Err(ResidencyLedgerError::PublicationExceedsReservation {
1476                id: id.clone(),
1477                tier,
1478                reserved_bytes: reserved.bytes,
1479                actual_bytes,
1480            });
1481        }
1482        let adjusted = self
1483            .tier_bytes(tier)
1484            .checked_sub(reserved.bytes)
1485            .and_then(|bytes| bytes.checked_add(actual_bytes))
1486            .ok_or_else(|| inconsistent(id, tier, "publication accounting"))?;
1487        self.set_tier_bytes(tier, adjusted);
1488        let tick = self.next_tick();
1489        let copy = self
1490            .units
1491            .get_mut(id)
1492            .and_then(|unit| unit.copy_mut(tier))
1493            .ok_or_else(|| inconsistent(id, tier, "publication mutation"))?;
1494        *copy = LedgerCopy {
1495            lifecycle: CopyLifecycle::Resident,
1496            bytes: actual_bytes,
1497            pins: 0,
1498            last_used: tick,
1499            frequency: 0,
1500            in_flight,
1501        };
1502        self.update_resident_telemetry(tier);
1503        Ok(())
1504    }
1505
1506    /// Rolls back an unpublished reservation without recording an eviction.
1507    pub fn rollback_reserved(
1508        &mut self,
1509        id: &OffloadUnitId,
1510        tier: MemoryTier,
1511    ) -> Result<(), ResidencyLedgerError> {
1512        validate_ledger_tier(tier, "reservation rollback")?;
1513        let copy = self.units.get(id).and_then(|unit| unit.copy(tier)).copied();
1514        let Some(copy) = copy else {
1515            return Ok(());
1516        };
1517        if copy.lifecycle != CopyLifecycle::Reserved {
1518            return Ok(());
1519        }
1520        self.remove_copy(id, tier, false)?;
1521        Ok(())
1522    }
1523
1524    /// Allocates a stable generation for one exact transfer submission.
1525    pub fn next_transfer_generation(&mut self) -> Result<u64, ResidencyLedgerError> {
1526        let generation = self.next_transfer_generation;
1527        self.next_transfer_generation = self.next_transfer_generation.checked_add(1).ok_or(
1528            ResidencyLedgerError::ArithmeticOverflow {
1529                context: "resident transfer generation",
1530            },
1531        )?;
1532        Ok(generation)
1533    }
1534
1535    /// Resolves exact transfer completion and returns failed backend copies to release.
1536    pub fn resolve_transfer(
1537        &mut self,
1538        ids: &[OffloadUnitId],
1539        tier: MemoryTier,
1540        generation: u64,
1541        succeeded: bool,
1542    ) -> Result<Vec<EvictedResidencyCopy>, ResidencyLedgerError> {
1543        validate_ledger_tier(tier, "transfer resolution")?;
1544        self.validate_batch(ids, tier)?;
1545        let mut removed = Vec::new();
1546        for id in ids {
1547            let matches = self
1548                .units
1549                .get(id)
1550                .and_then(|unit| unit.copy(tier))
1551                .is_some_and(|copy| {
1552                    copy.lifecycle == CopyLifecycle::Resident && copy.in_flight == Some(generation)
1553                });
1554            if !matches {
1555                continue;
1556            }
1557            if succeeded {
1558                self.units
1559                    .get_mut(id)
1560                    .and_then(|unit| unit.copy_mut(tier))
1561                    .ok_or_else(|| inconsistent(id, tier, "transfer success resolution"))?
1562                    .in_flight = None;
1563            } else {
1564                removed.push(self.remove_copy(id, tier, false)?);
1565            }
1566        }
1567        Ok(removed)
1568    }
1569
1570    /// Pins a resident copy and records weighted demand.
1571    pub fn pin(
1572        &mut self,
1573        id: &OffloadUnitId,
1574        tier: MemoryTier,
1575        demand: u64,
1576    ) -> Result<ResidentCopyStatus, ResidencyLedgerError> {
1577        validate_ledger_tier(tier, "pin")?;
1578        let tick = self.next_tick();
1579        let copy = self
1580            .units
1581            .get_mut(id)
1582            .ok_or_else(|| ResidencyLedgerError::UnknownUnit { id: id.clone() })?
1583            .copy_mut(tier)
1584            .filter(|copy| copy.lifecycle == CopyLifecycle::Resident)
1585            .ok_or_else(|| inconsistent(id, tier, "pin"))?;
1586        copy.pins = copy
1587            .pins
1588            .checked_add(1)
1589            .ok_or(ResidencyLedgerError::ArithmeticOverflow {
1590                context: "resident lease count",
1591            })?;
1592        copy.last_used = tick;
1593        copy.frequency = copy.frequency.saturating_add(demand);
1594        Ok(copy.status().expect("resident copy has status"))
1595    }
1596
1597    /// Releases one pin. Unknown or already released copies are ignored for drop safety.
1598    pub fn unpin(&mut self, id: &OffloadUnitId, tier: MemoryTier) {
1599        if let Some(copy) = self.units.get_mut(id).and_then(|unit| unit.copy_mut(tier)) {
1600            copy.pins = copy.pins.saturating_sub(1);
1601        }
1602    }
1603
1604    /// Updates recency for an already resident copy.
1605    pub fn touch(
1606        &mut self,
1607        id: &OffloadUnitId,
1608        tier: MemoryTier,
1609    ) -> Result<(), ResidencyLedgerError> {
1610        validate_ledger_tier(tier, "touch")?;
1611        let tick = self.next_tick();
1612        self.units
1613            .get_mut(id)
1614            .ok_or_else(|| ResidencyLedgerError::UnknownUnit { id: id.clone() })?
1615            .copy_mut(tier)
1616            .filter(|copy| copy.lifecycle == CopyLifecycle::Resident)
1617            .ok_or_else(|| inconsistent(id, tier, "touch"))?
1618            .last_used = tick;
1619        Ok(())
1620    }
1621
1622    /// Replaces one named protected window.
1623    pub fn set_group_window(
1624        &mut self,
1625        group: &str,
1626        active: &[OffloadUnitId],
1627        tier: MemoryTier,
1628    ) -> Result<(), ResidencyLedgerError> {
1629        validate_ledger_tier(tier, "protected window")?;
1630        if group.trim().is_empty() {
1631            return Err(ResidencyLedgerError::InvalidGroupId);
1632        }
1633        self.validate_batch(active, tier)?;
1634        let key = (group.to_string(), tier);
1635        if active.is_empty() {
1636            self.group_windows.remove(&key);
1637        } else {
1638            self.group_windows
1639                .insert(key, active.iter().cloned().collect());
1640        }
1641        let union = self
1642            .group_windows
1643            .iter()
1644            .filter(|((_, candidate_tier), _)| *candidate_tier == tier)
1645            .flat_map(|(_, window)| window.iter().cloned())
1646            .collect();
1647        self.active_windows.insert(tier, union);
1648        Ok(())
1649    }
1650
1651    /// Explicitly evicts an unpinned copy and returns backend storage to release.
1652    pub fn evict(
1653        &mut self,
1654        id: &OffloadUnitId,
1655        tier: MemoryTier,
1656    ) -> Result<Option<EvictedResidencyCopy>, ResidencyLedgerError> {
1657        validate_ledger_tier(tier, "evict")?;
1658        let unit = self
1659            .units
1660            .get(id)
1661            .ok_or_else(|| ResidencyLedgerError::UnknownUnit { id: id.clone() })?;
1662        let Some(copy) = unit.copy(tier).and_then(|copy| copy.status()) else {
1663            return Ok(None);
1664        };
1665        if unit.spec.policy() == ResidencyPolicy::Pinned {
1666            return Err(ResidencyLedgerError::PinnedEviction {
1667                id: id.clone(),
1668                tier,
1669            });
1670        }
1671        if copy.pins != 0 {
1672            return Err(ResidencyLedgerError::InUseEviction {
1673                id: id.clone(),
1674                tier,
1675                pin_count: copy.pins,
1676            });
1677        }
1678        self.remove_copy(id, tier, true).map(Some)
1679    }
1680
1681    /// Records one prefetch result in neutral telemetry.
1682    pub fn record_prefetch(&mut self, tier: MemoryTier, outcome: PrefetchOutcome) {
1683        self.telemetry.record_tier_prefetch(tier, outcome);
1684    }
1685
1686    /// Records backend transfer observations.
1687    pub fn record_transfer(
1688        &mut self,
1689        direction: TransferDirection,
1690        bytes: u64,
1691        duration: Duration,
1692    ) {
1693        self.telemetry.record_transfer(direction, bytes, duration);
1694    }
1695
1696    /// Records time spent waiting for demand materialization.
1697    pub fn record_prefetch_stall(&mut self, duration: Duration) {
1698        self.telemetry.record_prefetch_stall(duration);
1699    }
1700
1701    /// Records a backend allocator observation.
1702    pub fn record_allocator_memory(&mut self, metrics: AllocatorMemoryMetrics) {
1703        self.telemetry.record_allocator_memory(metrics);
1704    }
1705
1706    /// Samples optional process memory using the portable core sampler.
1707    pub fn sample_process_metrics(&mut self) {
1708        self.telemetry.sample_process_metrics();
1709    }
1710
1711    /// Returns immutable accounting telemetry.
1712    pub fn telemetry(&self) -> OffloadReport {
1713        self.telemetry.snapshot()
1714    }
1715
1716    /// Returns logical unit reports in stable identifier order.
1717    pub fn unit_reports(&self) -> Vec<UnitResidencyReport> {
1718        let active = self.active_window();
1719        self.units
1720            .values()
1721            .map(|unit| {
1722                let host = unit.host.and_then(LedgerCopy::status);
1723                let device = unit.device.and_then(LedgerCopy::status);
1724                UnitResidencyReport {
1725                    id: unit.spec.id().clone(),
1726                    planned_tier: unit.spec.tier(),
1727                    policy: unit.spec.policy(),
1728                    expected_bytes: unit.spec.bytes(),
1729                    host_allocated_bytes: host.map_or(0, ResidentCopyStatus::bytes),
1730                    device_allocated_bytes: device.map_or(0, ResidentCopyStatus::bytes),
1731                    host_resident: host.is_some(),
1732                    device_resident: device.is_some(),
1733                    host_pins: host.map_or(0, ResidentCopyStatus::pins),
1734                    device_pins: device.map_or(0, ResidentCopyStatus::pins),
1735                    active_window: active.contains(unit.spec.id()),
1736                }
1737            })
1738            .collect()
1739    }
1740
1741    /// Returns the union of every named protected window.
1742    pub fn active_window(&self) -> BTreeSet<OffloadUnitId> {
1743        self.active_windows
1744            .values()
1745            .flat_map(|window| window.iter().cloned())
1746            .collect()
1747    }
1748
1749    fn budget(&self, tier: MemoryTier) -> Option<u64> {
1750        match tier {
1751            MemoryTier::Host => self.plan.config().host_budget_bytes(),
1752            MemoryTier::Device => self.plan.config().device_budget_bytes(),
1753            MemoryTier::Disk => None,
1754        }
1755    }
1756
1757    fn tier_bytes(&self, tier: MemoryTier) -> u64 {
1758        self.resident_bytes.get(tier)
1759    }
1760
1761    fn set_tier_bytes(&mut self, tier: MemoryTier, bytes: u64) {
1762        self.resident_bytes.set(tier, bytes);
1763    }
1764
1765    fn eviction_candidates(
1766        &self,
1767        tier: MemoryTier,
1768        protected: &BTreeSet<OffloadUnitId>,
1769    ) -> Vec<OffloadUnitId> {
1770        let mut candidates = self
1771            .units
1772            .values()
1773            .filter_map(|unit| {
1774                let copy = unit.copy(tier)?;
1775                if copy.lifecycle != CopyLifecycle::Resident
1776                    || unit.spec.policy() == ResidencyPolicy::Pinned
1777                    || copy.pins != 0
1778                    || protected.contains(unit.spec.id())
1779                    || self
1780                        .active_windows
1781                        .get(&tier)
1782                        .is_some_and(|window| window.contains(unit.spec.id()))
1783                {
1784                    return None;
1785                }
1786                let priority = match unit.spec.policy() {
1787                    ResidencyPolicy::Windowed => 0u8,
1788                    ResidencyPolicy::Cacheable => 1u8,
1789                    ResidencyPolicy::Pinned => return None,
1790                };
1791                let frequency = match self.plan.config().eviction_policy() {
1792                    CacheEvictionPolicy::LeastRecentlyUsed => 0,
1793                    CacheEvictionPolicy::LeastFrequentlyUsed => copy.frequency,
1794                };
1795                Some((priority, frequency, copy.last_used, unit.spec.id().clone()))
1796            })
1797            .collect::<Vec<_>>();
1798        candidates.sort();
1799        candidates.into_iter().map(|(_, _, _, id)| id).collect()
1800    }
1801
1802    fn blockers(
1803        &self,
1804        tier: MemoryTier,
1805        protected: &BTreeSet<OffloadUnitId>,
1806    ) -> Vec<ResidencyBlocker> {
1807        self.units
1808            .values()
1809            .filter_map(|unit| {
1810                let copy = unit.copy(tier)?;
1811                if copy.lifecycle != CopyLifecycle::Resident {
1812                    return None;
1813                }
1814                let pinned = unit.spec.policy() == ResidencyPolicy::Pinned;
1815                let active_window = self
1816                    .active_windows
1817                    .get(&tier)
1818                    .is_some_and(|window| window.contains(unit.spec.id()));
1819                let request_protected = protected.contains(unit.spec.id());
1820                (pinned || copy.pins != 0 || active_window || request_protected).then(|| {
1821                    ResidencyBlocker {
1822                        id: unit.spec.id().clone(),
1823                        pinned,
1824                        in_use: copy.pins,
1825                        active_window,
1826                        request_protected,
1827                    }
1828                })
1829            })
1830            .collect()
1831    }
1832
1833    fn remove_copy(
1834        &mut self,
1835        id: &OffloadUnitId,
1836        tier: MemoryTier,
1837        record_eviction: bool,
1838    ) -> Result<EvictedResidencyCopy, ResidencyLedgerError> {
1839        let copy = self
1840            .units
1841            .get_mut(id)
1842            .and_then(|unit| unit.slot_mut(tier))
1843            .and_then(Option::take)
1844            .ok_or_else(|| inconsistent(id, tier, "copy removal"))?;
1845        let bytes = self
1846            .tier_bytes(tier)
1847            .checked_sub(copy.bytes)
1848            .ok_or_else(|| inconsistent(id, tier, "copy removal accounting"))?;
1849        self.set_tier_bytes(tier, bytes);
1850        self.update_resident_telemetry(tier);
1851        if record_eviction {
1852            self.telemetry.record_tier_eviction(tier, copy.bytes);
1853        }
1854        Ok(EvictedResidencyCopy {
1855            id: id.clone(),
1856            tier,
1857            bytes: copy.bytes,
1858        })
1859    }
1860
1861    fn update_resident_telemetry(&mut self, tier: MemoryTier) {
1862        let units = self
1863            .units
1864            .values()
1865            .filter(|unit| unit.copy(tier).and_then(|copy| copy.status()).is_some())
1866            .count();
1867        self.telemetry
1868            .set_resident_bytes(tier, self.tier_bytes(tier));
1869        self.telemetry.set_resident_units(tier, units);
1870    }
1871
1872    fn next_tick(&mut self) -> u64 {
1873        if self.tick == u64::MAX {
1874            for unit in self.units.values_mut() {
1875                if let Some(copy) = unit.host.as_mut() {
1876                    copy.last_used /= 2;
1877                }
1878                if let Some(copy) = unit.device.as_mut() {
1879                    copy.last_used /= 2;
1880                }
1881            }
1882            self.tick /= 2;
1883        }
1884        self.tick += 1;
1885        self.tick
1886    }
1887}
1888
1889fn validate_ledger_tier(
1890    tier: MemoryTier,
1891    operation: &'static str,
1892) -> Result<(), ResidencyLedgerError> {
1893    if tier == MemoryTier::Disk {
1894        Err(ResidencyLedgerError::InvalidTargetTier { operation })
1895    } else {
1896        Ok(())
1897    }
1898}
1899
1900fn inconsistent(
1901    id: &OffloadUnitId,
1902    tier: MemoryTier,
1903    operation: &'static str,
1904) -> ResidencyLedgerError {
1905    ResidencyLedgerError::StateInconsistent {
1906        id: id.clone(),
1907        tier,
1908        operation,
1909    }
1910}
1911
1912/// Backend-neutral residency ownership failure.
1913#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
1914pub enum ResidencyLedgerError {
1915    /// Initial planned materialization has not completed.
1916    #[error("residency manager has not been initialized")]
1917    NotInitialized,
1918    /// A named protected window had no stable identity.
1919    #[error("resident execution group id must not be empty")]
1920    InvalidGroupId,
1921    /// One batch named the same logical unit more than once.
1922    #[error("batched residency acquisition contains a duplicate unit")]
1923    DuplicateBatchUnit,
1924    /// A request referenced an unknown logical unit.
1925    #[error("unknown residency unit: {id}")]
1926    UnknownUnit {
1927        /// Missing unit.
1928        id: OffloadUnitId,
1929    },
1930    /// Disk is not a materialized target tier.
1931    #[error("{operation} requires a host or device tier")]
1932    InvalidTargetTier {
1933        /// Invalid operation.
1934        operation: &'static str,
1935    },
1936    /// A zero-byte physical reservation was requested.
1937    #[error("residency unit {id} cannot reserve zero bytes")]
1938    ZeroReservation {
1939        /// Invalid unit.
1940        id: OffloadUnitId,
1941    },
1942    /// A backend attempted to reserve an existing copy.
1943    #[error("residency unit {id} already has a {tier:?} copy")]
1944    CopyAlreadyExists {
1945        /// Existing unit.
1946        id: OffloadUnitId,
1947        /// Existing tier.
1948        tier: MemoryTier,
1949    },
1950    /// Backend storage exceeded the capacity reserved for its publication.
1951    #[error("residency unit {id} published {actual_bytes} bytes to {tier:?} after reserving only {reserved_bytes}")]
1952    PublicationExceedsReservation {
1953        /// Published unit.
1954        id: OffloadUnitId,
1955        /// Target tier.
1956        tier: MemoryTier,
1957        /// Capacity charged before materialization.
1958        reserved_bytes: u64,
1959        /// Actual materialized capacity.
1960        actual_bytes: u64,
1961    },
1962    /// Backend storage publication had no physical capacity.
1963    #[error("residency unit {id} cannot publish a zero-byte copy to {tier:?}")]
1964    ZeroPublication {
1965        /// Published unit.
1966        id: OffloadUnitId,
1967        /// Target tier.
1968        tier: MemoryTier,
1969    },
1970    /// Checked lifecycle arithmetic overflowed.
1971    #[error("residency arithmetic overflow during {context}")]
1972    ArithmeticOverflow {
1973        /// Stable operation description.
1974        context: &'static str,
1975    },
1976    /// No safe victim could satisfy a finite tier budget.
1977    #[error("cannot reserve {required_bytes} bytes for {requested} in {tier:?}: {resident_bytes}/{budget_bytes} bytes resident")]
1978    BudgetExhausted {
1979        /// Requested unit.
1980        requested: OffloadUnitId,
1981        /// Requested tier.
1982        tier: MemoryTier,
1983        /// Required physical bytes.
1984        required_bytes: u64,
1985        /// Finite tier budget.
1986        budget_bytes: u64,
1987        /// Currently charged bytes.
1988        resident_bytes: u64,
1989        /// Units preventing eviction.
1990        blocking_units: Vec<ResidencyBlocker>,
1991    },
1992    /// Explicit eviction contradicted pinned lifetime policy.
1993    #[error("cannot evict pinned residency unit {id} from {tier:?}")]
1994    PinnedEviction {
1995        /// Pinned unit.
1996        id: OffloadUnitId,
1997        /// Resident tier.
1998        tier: MemoryTier,
1999    },
2000    /// Explicit eviction targeted leased storage.
2001    #[error("cannot evict residency unit {id} from {tier:?} while {pin_count} leases are active")]
2002    InUseEviction {
2003        /// Leased unit.
2004        id: OffloadUnitId,
2005        /// Resident tier.
2006        tier: MemoryTier,
2007        /// Active leases.
2008        pin_count: u64,
2009    },
2010    /// Backend storage and ledger transitions were applied out of order.
2011    #[error("residency ledger state is inconsistent for {id} in {tier:?} during {operation}")]
2012    StateInconsistent {
2013        /// Unit whose ownership invariant was violated.
2014        id: OffloadUnitId,
2015        /// Tier whose copy state was inconsistent.
2016        tier: MemoryTier,
2017        /// Stable transition description.
2018        operation: &'static str,
2019    },
2020}
2021
2022#[cfg(test)]
2023mod tests {
2024    use super::*;
2025
2026    fn unit(id: &str, bytes: u64, policy: ResidencyPolicy, tier: MemoryTier) -> OffloadUnitSpec {
2027        OffloadUnitSpec::new(OffloadUnitId::new(id).unwrap(), bytes, policy, tier).unwrap()
2028    }
2029
2030    #[test]
2031    fn explicit_plan_is_validated_sorted_and_inspectable() {
2032        let config = OffloadConfig::new(Some(80), Some(40), 2).unwrap();
2033        let plan = OffloadPlan::new(
2034            config,
2035            [
2036                unit("layer.2", 20, ResidencyPolicy::Windowed, MemoryTier::Disk),
2037                unit("layer.0", 40, ResidencyPolicy::Pinned, MemoryTier::Device),
2038                unit("layer.1", 30, ResidencyPolicy::Cacheable, MemoryTier::Host),
2039            ],
2040        )
2041        .unwrap();
2042
2043        assert_eq!(
2044            plan.units()
2045                .iter()
2046                .map(|unit| unit.id().as_str())
2047                .collect::<Vec<_>>(),
2048            ["layer.0", "layer.1", "layer.2"]
2049        );
2050        assert_eq!(plan.config(), config);
2051        assert_eq!(plan.planned_bytes(), TierByteTotals::new(40, 30, 20));
2052        assert_eq!(
2053            plan.unit(&OffloadUnitId::new("layer.1").unwrap())
2054                .unwrap()
2055                .bytes(),
2056            30
2057        );
2058    }
2059
2060    #[test]
2061    fn plan_and_report_serialization_preserve_validated_state() {
2062        let plan = OffloadPlan::new(
2063            OffloadConfig::new(Some(80), Some(40), 2)
2064                .unwrap()
2065                .with_eviction_policy(CacheEvictionPolicy::LeastFrequentlyUsed),
2066            [
2067                unit("layer.1", 20, ResidencyPolicy::Windowed, MemoryTier::Disk),
2068                unit("layer.0", 40, ResidencyPolicy::Pinned, MemoryTier::Device),
2069            ],
2070        )
2071        .unwrap();
2072        let encoded = serde_json::to_string(&plan).unwrap();
2073        let decoded: OffloadPlan = serde_json::from_str(&encoded).unwrap();
2074        assert_eq!(decoded, plan);
2075
2076        let report = OffloadTelemetry::from_plan(&plan).snapshot();
2077        let encoded = serde_json::to_string(&report).unwrap();
2078        assert_eq!(
2079            serde_json::from_str::<OffloadReport>(&encoded).unwrap(),
2080            report
2081        );
2082    }
2083
2084    #[test]
2085    fn deserialization_rejects_invalid_plan_instead_of_bypassing_constructors() {
2086        let invalid = serde_json::json!({
2087            "schema_version": OFFLOAD_PLAN_SCHEMA_VERSION,
2088            "config": {
2089                "device_budget_bytes": 1,
2090                "host_budget_bytes": null,
2091                "prefetch_depth": 1,
2092                "eviction_policy": "least_recently_used"
2093            },
2094            "units": [{
2095                "id": "layer.0",
2096                "bytes": 2,
2097                "policy": "pinned",
2098                "tier": "device"
2099            }]
2100        });
2101        assert!(serde_json::from_value::<OffloadPlan>(invalid).is_err());
2102    }
2103
2104    #[test]
2105    fn duplicate_identifiers_are_rejected_deterministically() {
2106        let duplicate = OffloadPlan::new(
2107            OffloadConfig::default(),
2108            [
2109                unit("b", 1, ResidencyPolicy::Cacheable, MemoryTier::Host),
2110                unit("a", 1, ResidencyPolicy::Pinned, MemoryTier::Device),
2111                unit("a", 2, ResidencyPolicy::Cacheable, MemoryTier::Host),
2112            ],
2113        )
2114        .unwrap_err();
2115        assert_eq!(
2116            duplicate,
2117            OffloadError::DuplicateUnitId {
2118                id: OffloadUnitId::new("a").unwrap()
2119            }
2120        );
2121    }
2122
2123    #[test]
2124    fn finite_tier_budgets_are_enforced() {
2125        let error = OffloadPlan::new(
2126            OffloadConfig::new(Some(9), None, 1).unwrap(),
2127            [unit(
2128                "weights",
2129                10,
2130                ResidencyPolicy::Pinned,
2131                MemoryTier::Device,
2132            )],
2133        )
2134        .unwrap_err();
2135        assert_eq!(
2136            error,
2137            OffloadError::BudgetExceeded {
2138                tier: MemoryTier::Device,
2139                planned_bytes: 10,
2140                budget_bytes: 9,
2141            }
2142        );
2143    }
2144
2145    #[test]
2146    fn byte_total_overflow_is_reported() {
2147        let error = OffloadPlan::new(
2148            OffloadConfig::default(),
2149            [
2150                unit("a", u64::MAX, ResidencyPolicy::Cacheable, MemoryTier::Host),
2151                unit("b", 1, ResidencyPolicy::Cacheable, MemoryTier::Host),
2152            ],
2153        )
2154        .unwrap_err();
2155        assert_eq!(
2156            error,
2157            OffloadError::ByteTotalOverflow {
2158                tier: MemoryTier::Host
2159            }
2160        );
2161    }
2162
2163    #[test]
2164    fn meaningless_and_contradictory_inputs_are_rejected() {
2165        assert_eq!(
2166            OffloadConfig::new(None, None, 0),
2167            Err(OffloadError::ZeroPrefetchDepth)
2168        );
2169        let id = OffloadUnitId::new("empty").unwrap();
2170        assert_eq!(
2171            OffloadUnitSpec::new(id.clone(), 0, ResidencyPolicy::Cacheable, MemoryTier::Host),
2172            Err(OffloadError::ZeroSizedUnit { id })
2173        );
2174        assert!(matches!(
2175            OffloadUnitSpec::new(
2176                OffloadUnitId::new("pinned-disk").unwrap(),
2177                1,
2178                ResidencyPolicy::Pinned,
2179                MemoryTier::Disk
2180            ),
2181            Err(OffloadError::ContradictoryAssignment { .. })
2182        ));
2183    }
2184
2185    #[test]
2186    fn telemetry_accounts_for_residency_activity_and_runtime_samples() {
2187        let plan = OffloadPlan::new(
2188            OffloadConfig::default(),
2189            [unit("a", 10, ResidencyPolicy::Pinned, MemoryTier::Device)],
2190        )
2191        .unwrap();
2192        let mut telemetry = OffloadTelemetry::from_plan(&plan);
2193        telemetry.set_resident_bytes(MemoryTier::Device, 8);
2194        telemetry.set_resident_bytes(MemoryTier::Device, 12);
2195        telemetry.set_resident_bytes(MemoryTier::Device, 6);
2196        telemetry.set_resident_units(MemoryTier::Device, 1);
2197        telemetry.set_resident_units(MemoryTier::Device, 2);
2198        telemetry.set_resident_units(MemoryTier::Device, 1);
2199        telemetry.record_transfer(TransferDirection::HostToDevice, 5, Duration::from_millis(2));
2200        telemetry.record_transfer(TransferDirection::HostToDevice, 7, Duration::from_millis(3));
2201        telemetry.record_tier_prefetch(MemoryTier::Device, PrefetchOutcome::Hit);
2202        telemetry.record_tier_prefetch(MemoryTier::Host, PrefetchOutcome::Miss);
2203        telemetry.record_prefetch_stall(Duration::from_millis(4));
2204        telemetry.record_tier_eviction(MemoryTier::Device, 3);
2205        telemetry.record_tier_eviction(MemoryTier::Host, 4);
2206        telemetry.record_allocator_memory(AllocatorMemoryMetrics::new(11, 12, 13));
2207        telemetry.record_process_metrics(ProcessMetrics::new(Some(14), Some(15), Some(16)));
2208
2209        let report = telemetry.snapshot();
2210        assert_eq!(report.planned_bytes().get(MemoryTier::Device), 10);
2211        assert_eq!(report.resident_bytes().get(MemoryTier::Device), 6);
2212        assert_eq!(report.peak_resident_bytes().get(MemoryTier::Device), 12);
2213        assert_eq!(report.resident_units().get(MemoryTier::Device), 1);
2214        assert_eq!(report.peak_resident_units().get(MemoryTier::Device), 2);
2215        assert_eq!(report.transfer(TransferDirection::HostToDevice).count(), 2);
2216        assert_eq!(report.transfer(TransferDirection::HostToDevice).bytes(), 12);
2217        assert_eq!(
2218            report.transfer(TransferDirection::HostToDevice).duration(),
2219            Duration::from_millis(5)
2220        );
2221        assert_eq!(report.prefetch().requests(), 2);
2222        assert_eq!(report.prefetch().hits(), 1);
2223        assert_eq!(report.prefetch().misses(), 1);
2224        assert_eq!(report.prefetch().stalls(), 1);
2225        assert_eq!(report.prefetch().stall_duration(), Duration::from_millis(4));
2226        assert_eq!(report.evictions().count(), 2);
2227        assert_eq!(report.evictions().bytes(), 7);
2228        assert_eq!(report.tier_prefetch(MemoryTier::Device).hits(), 1);
2229        assert_eq!(report.tier_prefetch(MemoryTier::Host).misses(), 1);
2230        assert_eq!(report.tier_evictions(MemoryTier::Device).bytes(), 3);
2231        assert_eq!(report.tier_evictions(MemoryTier::Host).bytes(), 4);
2232        assert_eq!(report.allocator_memory().unwrap().peak_bytes(), 13);
2233        assert_eq!(report.process_metrics().rss_bytes(), Some(14));
2234        assert!(report.process_sampled());
2235    }
2236
2237    #[test]
2238    fn snapshot_is_immutable_and_reset_clears_everything() {
2239        let mut telemetry = OffloadTelemetry::default();
2240        telemetry.set_planned_bytes(TierByteTotals::new(1, 2, 3));
2241        telemetry.set_resident_bytes(MemoryTier::Host, 4);
2242        telemetry.set_resident_units(MemoryTier::Host, 1);
2243        let snapshot = telemetry.snapshot();
2244
2245        telemetry.set_resident_bytes(MemoryTier::Host, 9);
2246        telemetry.record_eviction(5);
2247        assert_eq!(snapshot.resident_bytes().get(MemoryTier::Host), 4);
2248        assert_eq!(snapshot.resident_units().get(MemoryTier::Host), 1);
2249        assert_eq!(snapshot.evictions(), EvictionMetrics::default());
2250        assert!(!snapshot.process_sampled());
2251
2252        telemetry.reset();
2253        assert_eq!(telemetry.snapshot(), OffloadTelemetry::default().snapshot());
2254    }
2255
2256    #[test]
2257    fn telemetry_counters_saturate() {
2258        let mut telemetry = OffloadTelemetry::default();
2259        telemetry.transfers[TransferDirection::DiskToHost.index()] = TransferMetrics {
2260            count: u64::MAX,
2261            bytes: u64::MAX,
2262            duration: Duration::MAX,
2263        };
2264        telemetry.evictions = EvictionMetrics {
2265            count: u64::MAX,
2266            bytes: u64::MAX,
2267        };
2268        telemetry.record_transfer(TransferDirection::DiskToHost, 1, Duration::from_nanos(1));
2269        telemetry.record_eviction(1);
2270        let report = telemetry.snapshot();
2271        assert_eq!(
2272            report.transfer(TransferDirection::DiskToHost).count(),
2273            u64::MAX
2274        );
2275        assert_eq!(
2276            report.transfer(TransferDirection::DiskToHost).bytes(),
2277            u64::MAX
2278        );
2279        assert_eq!(
2280            report.transfer(TransferDirection::DiskToHost).duration(),
2281            Duration::MAX
2282        );
2283        assert_eq!(report.evictions().count(), u64::MAX);
2284        assert_eq!(report.evictions().bytes(), u64::MAX);
2285    }
2286
2287    #[test]
2288    fn optional_process_sampler_never_requires_platform_support() {
2289        let metrics = sample_process_metrics();
2290        if let Some(rss_bytes) = metrics.rss_bytes() {
2291            assert!(rss_bytes > 0);
2292        }
2293        let mut telemetry = OffloadTelemetry::default();
2294        telemetry.sample_process_metrics();
2295        let report = telemetry.snapshot();
2296        let _ = report.process_metrics();
2297        assert!(report.process_sampled());
2298    }
2299
2300    fn disk_ledger(
2301        device_budget: Option<u64>,
2302        units: impl IntoIterator<Item = (&'static str, u64, ResidencyPolicy)>,
2303    ) -> ResidencyLedger {
2304        let plan = OffloadPlan::new(
2305            OffloadConfig::new(device_budget, None, 1).unwrap(),
2306            units
2307                .into_iter()
2308                .map(|(id, bytes, policy)| unit(id, bytes, policy, MemoryTier::Disk)),
2309        )
2310        .unwrap();
2311        ResidencyLedger::new(plan)
2312    }
2313
2314    fn id(value: &str) -> OffloadUnitId {
2315        OffloadUnitId::new(value).unwrap()
2316    }
2317
2318    fn publish_device(
2319        ledger: &mut ResidencyLedger,
2320        unit: &str,
2321        bytes: u64,
2322        generation: Option<u64>,
2323    ) {
2324        ledger
2325            .reserve_copy(&id(unit), MemoryTier::Device, bytes, &BTreeSet::new())
2326            .unwrap();
2327        ledger
2328            .publish_reserved(&id(unit), MemoryTier::Device, bytes, generation)
2329            .unwrap();
2330    }
2331
2332    #[test]
2333    fn ledger_owns_publication_pins_and_exact_completion() {
2334        let mut ledger = disk_ledger(Some(8), [("a", 8, ResidencyPolicy::Cacheable)]);
2335        assert_eq!(
2336            ledger.require_initialized(),
2337            Err(ResidencyLedgerError::NotInitialized)
2338        );
2339        ledger.mark_initialized();
2340
2341        let generation = ledger.next_transfer_generation().unwrap();
2342        ledger
2343            .reserve_copy(&id("a"), MemoryTier::Device, 8, &BTreeSet::new())
2344            .unwrap();
2345        assert_eq!(
2346            ledger.copy_status(&id("a"), MemoryTier::Device).unwrap(),
2347            None
2348        );
2349        ledger
2350            .publish_reserved(&id("a"), MemoryTier::Device, 8, Some(generation))
2351            .unwrap();
2352        let pinned = ledger.pin(&id("a"), MemoryTier::Device, 3).unwrap();
2353        assert_eq!(pinned.pins(), 1);
2354        assert_eq!(pinned.in_flight(), Some(generation));
2355        assert!(matches!(
2356            ledger.evict(&id("a"), MemoryTier::Device),
2357            Err(ResidencyLedgerError::InUseEviction { pin_count: 1, .. })
2358        ));
2359
2360        assert!(ledger
2361            .resolve_transfer(&[id("a")], MemoryTier::Device, generation + 1, true)
2362            .unwrap()
2363            .is_empty());
2364        assert_eq!(
2365            ledger
2366                .copy_status(&id("a"), MemoryTier::Device)
2367                .unwrap()
2368                .unwrap()
2369                .in_flight(),
2370            Some(generation)
2371        );
2372        ledger
2373            .resolve_transfer(&[id("a")], MemoryTier::Device, generation, true)
2374            .unwrap();
2375        ledger.unpin(&id("a"), MemoryTier::Device);
2376        assert!(ledger
2377            .evict(&id("a"), MemoryTier::Device)
2378            .unwrap()
2379            .is_some());
2380        assert_eq!(
2381            ledger.telemetry().resident_bytes().get(MemoryTier::Device),
2382            0
2383        );
2384        assert_eq!(ledger.telemetry().evictions().count(), 1);
2385    }
2386
2387    #[test]
2388    fn failed_exact_completion_discards_copy_and_accounting() {
2389        let mut ledger = disk_ledger(Some(8), [("a", 8, ResidencyPolicy::Cacheable)]);
2390        let generation = ledger.next_transfer_generation().unwrap();
2391        publish_device(&mut ledger, "a", 8, Some(generation));
2392
2393        let removed = ledger
2394            .resolve_transfer(&[id("a")], MemoryTier::Device, generation, false)
2395            .unwrap();
2396        assert_eq!(removed.len(), 1);
2397        assert_eq!(removed[0].id, id("a"));
2398        assert!(!ledger.is_resident(&id("a"), MemoryTier::Device).unwrap());
2399        assert_eq!(
2400            ledger.telemetry().resident_bytes().get(MemoryTier::Device),
2401            0
2402        );
2403        assert_eq!(ledger.telemetry().evictions().count(), 0);
2404    }
2405
2406    #[test]
2407    fn failed_batch_admission_is_atomic() {
2408        let mut ledger = disk_ledger(
2409            Some(16),
2410            [
2411                ("a", 8, ResidencyPolicy::Cacheable),
2412                ("b", 8, ResidencyPolicy::Cacheable),
2413                ("c", 8, ResidencyPolicy::Cacheable),
2414                ("d", 8, ResidencyPolicy::Cacheable),
2415            ],
2416        );
2417        publish_device(&mut ledger, "a", 8, None);
2418        publish_device(&mut ledger, "b", 8, None);
2419        ledger.pin(&id("b"), MemoryTier::Device, 1).unwrap();
2420
2421        let error = ledger
2422            .reserve_copies(
2423                &[(id("c"), 8), (id("d"), 8)],
2424                MemoryTier::Device,
2425                &BTreeSet::new(),
2426            )
2427            .unwrap_err();
2428        assert!(matches!(
2429            error,
2430            ResidencyLedgerError::BudgetExhausted {
2431                required_bytes: 16,
2432                resident_bytes: 16,
2433                ..
2434            }
2435        ));
2436        assert!(ledger.is_resident(&id("a"), MemoryTier::Device).unwrap());
2437        assert!(ledger.is_resident(&id("b"), MemoryTier::Device).unwrap());
2438        assert!(!ledger.is_resident(&id("c"), MemoryTier::Device).unwrap());
2439        assert!(!ledger.is_resident(&id("d"), MemoryTier::Device).unwrap());
2440        assert_eq!(ledger.telemetry().evictions().count(), 0);
2441    }
2442
2443    #[test]
2444    fn batch_reservation_evicts_only_after_complete_admission() {
2445        let mut ledger = disk_ledger(
2446            Some(16),
2447            [
2448                ("a", 8, ResidencyPolicy::Cacheable),
2449                ("b", 8, ResidencyPolicy::Cacheable),
2450                ("c", 8, ResidencyPolicy::Cacheable),
2451                ("d", 8, ResidencyPolicy::Cacheable),
2452            ],
2453        );
2454        publish_device(&mut ledger, "a", 8, None);
2455        publish_device(&mut ledger, "b", 8, None);
2456
2457        let evicted = ledger
2458            .reserve_copies(
2459                &[(id("c"), 8), (id("d"), 8)],
2460                MemoryTier::Device,
2461                &BTreeSet::new(),
2462            )
2463            .unwrap();
2464        assert_eq!(
2465            evicted
2466                .iter()
2467                .map(|copy| copy.id.clone())
2468                .collect::<Vec<_>>(),
2469            [id("a"), id("b")]
2470        );
2471        assert!(!ledger.is_resident(&id("a"), MemoryTier::Device).unwrap());
2472        assert!(!ledger.is_resident(&id("b"), MemoryTier::Device).unwrap());
2473        assert_eq!(
2474            ledger.copy_status(&id("c"), MemoryTier::Device).unwrap(),
2475            None
2476        );
2477        assert_eq!(
2478            ledger.copy_status(&id("d"), MemoryTier::Device).unwrap(),
2479            None
2480        );
2481        ledger
2482            .rollback_reserved(&id("c"), MemoryTier::Device)
2483            .unwrap();
2484        ledger
2485            .rollback_reserved(&id("d"), MemoryTier::Device)
2486            .unwrap();
2487        assert_eq!(
2488            ledger.telemetry().resident_bytes().get(MemoryTier::Device),
2489            0
2490        );
2491    }
2492
2493    #[test]
2494    fn equal_eviction_priority_uses_stable_unit_identity() {
2495        let mut ledger = disk_ledger(
2496            Some(16),
2497            [
2498                ("a", 8, ResidencyPolicy::Cacheable),
2499                ("b", 8, ResidencyPolicy::Cacheable),
2500                ("c", 8, ResidencyPolicy::Cacheable),
2501            ],
2502        );
2503        publish_device(&mut ledger, "a", 8, None);
2504        publish_device(&mut ledger, "b", 8, None);
2505        for unit in [id("a"), id("b")] {
2506            ledger
2507                .units
2508                .get_mut(&unit)
2509                .unwrap()
2510                .device
2511                .as_mut()
2512                .unwrap()
2513                .last_used = 10;
2514        }
2515
2516        let evicted = ledger
2517            .reserve_copy(&id("c"), MemoryTier::Device, 8, &BTreeSet::new())
2518            .unwrap();
2519        assert_eq!(evicted[0].id, id("a"));
2520    }
2521
2522    #[test]
2523    fn request_protection_is_reported_as_a_capacity_blocker() {
2524        let mut ledger = disk_ledger(
2525            Some(8),
2526            [
2527                ("active", 8, ResidencyPolicy::Cacheable),
2528                ("next", 8, ResidencyPolicy::Cacheable),
2529            ],
2530        );
2531        publish_device(&mut ledger, "active", 8, None);
2532        let protected = BTreeSet::from([id("active")]);
2533        let error = ledger
2534            .reserve_copy(&id("next"), MemoryTier::Device, 8, &protected)
2535            .unwrap_err();
2536        let ResidencyLedgerError::BudgetExhausted { blocking_units, .. } = error else {
2537            panic!("unexpected error: {error}");
2538        };
2539        assert_eq!(blocking_units.len(), 1);
2540        assert!(blocking_units[0].request_protected);
2541    }
2542
2543    #[test]
2544    fn windows_are_group_scoped_and_protect_capacity() {
2545        let mut ledger = disk_ledger(
2546            Some(16),
2547            [
2548                ("text", 8, ResidencyPolicy::Windowed),
2549                ("vision", 8, ResidencyPolicy::Windowed),
2550                ("next", 8, ResidencyPolicy::Cacheable),
2551            ],
2552        );
2553        publish_device(&mut ledger, "text", 8, None);
2554        publish_device(&mut ledger, "vision", 8, None);
2555        ledger
2556            .set_group_window("text", &[id("text")], MemoryTier::Device)
2557            .unwrap();
2558        ledger
2559            .set_group_window("vision", &[id("vision")], MemoryTier::Device)
2560            .unwrap();
2561        assert_eq!(
2562            ledger.active_window(),
2563            BTreeSet::from([id("text"), id("vision")])
2564        );
2565        assert!(matches!(
2566            ledger.reserve_copy(&id("next"), MemoryTier::Device, 8, &BTreeSet::new()),
2567            Err(ResidencyLedgerError::BudgetExhausted { .. })
2568        ));
2569
2570        ledger
2571            .set_group_window("text", &[], MemoryTier::Device)
2572            .unwrap();
2573        assert_eq!(ledger.active_window(), BTreeSet::from([id("vision")]));
2574        let evicted = ledger
2575            .reserve_copy(&id("next"), MemoryTier::Device, 8, &BTreeSet::new())
2576            .unwrap();
2577        assert_eq!(evicted[0].id, id("text"));
2578    }
2579
2580    #[test]
2581    fn unit_reports_round_trip_without_backend_state() {
2582        let mut ledger = disk_ledger(Some(8), [("a", 8, ResidencyPolicy::Cacheable)]);
2583        publish_device(&mut ledger, "a", 8, None);
2584        ledger.pin(&id("a"), MemoryTier::Device, 2).unwrap();
2585        let report = ledger.unit_reports().remove(0);
2586        let encoded = serde_json::to_string(&report).unwrap();
2587        assert_eq!(
2588            serde_json::from_str::<UnitResidencyReport>(&encoded).unwrap(),
2589            report
2590        );
2591        assert_eq!(report.device_allocated_bytes(), 8);
2592        assert_eq!(report.device_pins(), 1);
2593    }
2594
2595    #[test]
2596    fn publication_cannot_exceed_reserved_capacity() {
2597        let mut ledger = disk_ledger(Some(16), [("a", 8, ResidencyPolicy::Cacheable)]);
2598        ledger
2599            .reserve_copy(&id("a"), MemoryTier::Device, 8, &BTreeSet::new())
2600            .unwrap();
2601        assert!(matches!(
2602            ledger.publish_reserved(&id("a"), MemoryTier::Device, 0, None),
2603            Err(ResidencyLedgerError::ZeroPublication { .. })
2604        ));
2605        assert!(matches!(
2606            ledger.publish_reserved(&id("a"), MemoryTier::Device, 9, None),
2607            Err(ResidencyLedgerError::PublicationExceedsReservation {
2608                reserved_bytes: 8,
2609                actual_bytes: 9,
2610                ..
2611            })
2612        ));
2613        ledger
2614            .rollback_reserved(&id("a"), MemoryTier::Device)
2615            .unwrap();
2616        assert_eq!(
2617            ledger.telemetry().resident_bytes().get(MemoryTier::Device),
2618            0
2619        );
2620    }
2621}