Skip to main content

eredu_runtime/capture/
checkpoint.rs

1//! In-process checkpoints of the shared record owner, independent of native state.
2
3use super::*;
4use eredu_core::intervention::{
5    AdmittedInterventionPlan, InterventionDiscovery, InterventionEstimator, InterventionPlan,
6};
7use std::sync::Arc;
8
9/// Reusable checkpoint at a drained portable record boundary. This is only the
10/// capture/intervention component of a generation snapshot; callers must also
11/// establish native completion and save model, sampler, and semantic state.
12///
13/// It contains no estimator or native handle. Admission proofs remain opaque and
14/// cannot be deserialized. Clones share only immutable host admission data.
15#[derive(Clone)]
16pub struct CaptureCheckpoint {
17    owner: Arc<()>,
18    artifact_identity: String,
19    plan: AdmittedCapturePlan,
20    intervention: Option<AdmittedInterventionPlan>,
21    prediction: u64,
22    phase: CapturePhase,
23    has_step: bool,
24    usage: CaptureUsage,
25}
26
27/// Explicit admission inputs for a child capture run. The backend supplies the
28/// actual child's retained discovery and current estimator facts.
29pub struct CaptureForkRequest<'a> {
30    /// Child's actual loaded source, observation catalog and support.
31    pub discovery: &'a CaptureDiscovery,
32    /// Absolute output limit, including the inherited prefix.
33    pub max_predictions: u64,
34    /// Explicit child limits. Cumulative limits include inherited consumption.
35    pub limits: CaptureLimits,
36    /// Required when inheriting interventions; also permits adding a new plan.
37    pub intervention: Option<InterventionForkRequest<'a>>,
38}
39
40/// Shared re-admission of inherited or prospectively replaced interventions.
41pub struct InterventionForkRequest<'a> {
42    /// Actual child's backend/source/target discovery.
43    pub discovery: &'a InterventionDiscovery,
44    /// Fresh child facade session identity, never the source admission identity.
45    pub session_id: &'a str,
46    /// None inherits the original operations. Some replaces future operations;
47    /// an empty plan removes them. Earlier native state is not recomputed.
48    pub replacement: Option<InterventionPlan>,
49    /// Child's current side-effect-free native facts. Not part of the checkpoint.
50    pub estimator: Arc<dyn InterventionEstimator>,
51}
52
53/// Validated, move-only same-run schedule restoration. Dropping this preparation
54/// changes nothing. Commit is infallible so native and portable state can be
55/// installed atomically after every fallible operation has succeeded.
56pub struct PreparedCaptureRestore<'a> {
57    run: &'a mut CaptureSession,
58    prediction: u64,
59    phase: CapturePhase,
60    has_step: bool,
61}
62
63impl PreparedCaptureRestore<'_> {
64    /// Installs only rewindable schedule/delivery state, leaving consumption intact.
65    pub fn commit(self) {
66        self.run.prediction = self.prediction;
67        self.run.phase = self.phase;
68        self.run.has_step = self.has_step;
69        self.run.capture_seconds = 0.0;
70        self.run.checkpoint_ready = true;
71    }
72}
73
74impl CaptureSession {
75    /// Known logical host storage for a copied portable checkpoint. Includes all
76    /// admitted plan payloads and declarations, but no native estimator or handle.
77    /// Callers reserve this before `checkpoint` clones the admission data.
78    pub fn checkpoint_storage_bytes(&self, discovery: &CaptureDiscovery) -> Option<u64> {
79        checkpoint_storage_bytes(
80            &self.plan,
81            self.interventions.as_ref().map(|run| &run.plan),
82            &discovery.artifact_identity,
83        )
84    }
85    /// Saves portable schedule position only after the current records have been
86    /// delivered and all intervention finish checks succeeded. The initial run is
87    /// also a valid boundary. The caller supplies retained loaded discovery.
88    pub fn checkpoint(
89        &self,
90        discovery: &CaptureDiscovery,
91    ) -> Result<CaptureCheckpoint, CaptureError> {
92        if !self.checkpoint_ready
93            || self.records.is_some()
94            || self
95                .interventions
96                .as_ref()
97                .is_some_and(|run| run.records.is_some() || run.routing_pending.is_some())
98        {
99            return Err(CaptureError::Invalid(
100                "capture checkpoint requires a successful, drained record boundary".into(),
101            ));
102        }
103        let checked = self.plan.plan().clone().admit(
104            &discovery.catalog,
105            &discovery.support,
106            &discovery.support.capture,
107            self.plan.request(),
108        )?;
109        if checked.identity() != self.plan.identity()
110            || self
111                .interventions
112                .as_ref()
113                .is_some_and(|run| run.plan.artifact_identity() != discovery.artifact_identity)
114        {
115            return Err(CaptureError::Invalid(
116                "checkpoint source/admission mismatch".into(),
117            ));
118        }
119        Ok(CaptureCheckpoint {
120            owner: Arc::clone(&self.owner),
121            artifact_identity: discovery.artifact_identity.clone(),
122            plan: self.plan.clone(),
123            intervention: self.interventions.as_ref().map(|run| run.plan.clone()),
124            prediction: self.prediction,
125            phase: self.phase,
126            has_step: self.has_step,
127            usage: self.ledger.total(),
128        })
129    }
130
131    /// Checks a same-run restore before any native or portable mutation. Undelivered
132    /// records must be consumed even when the previous operation failed.
133    pub fn validate_restore(&self, checkpoint: &CaptureCheckpoint) -> Result<(), CaptureError> {
134        if !Arc::ptr_eq(&self.owner, &checkpoint.owner)
135            || self.plan.identity() != checkpoint.plan.identity()
136            || self.interventions.as_ref().map(|run| run.plan.identity())
137                != checkpoint.intervention.as_ref().map(|plan| plan.identity())
138        {
139            return Err(CaptureError::Invalid(
140                "capture checkpoint belongs to another run".into(),
141            ));
142        }
143        if self.records.is_some()
144            || self
145                .interventions
146                .as_ref()
147                .is_some_and(|run| run.records.is_some() || run.routing_pending.is_some())
148        {
149            return Err(CaptureError::Invalid(
150                "restore requires drained records and resolved routing".into(),
151            ));
152        }
153        Ok(())
154    }
155
156    /// Rewinds schedule state without refunding any consumed resource or publishing
157    /// old records again. Native restoration must succeed before this is committed.
158    pub fn restore(&mut self, checkpoint: &CaptureCheckpoint) -> Result<(), CaptureError> {
159        self.prepare_restore(checkpoint)?.commit();
160        Ok(())
161    }
162
163    /// Validates restoration while exclusively borrowing the run until commit.
164    pub fn prepare_restore(
165        &mut self,
166        checkpoint: &CaptureCheckpoint,
167    ) -> Result<PreparedCaptureRestore<'_>, CaptureError> {
168        self.validate_restore(checkpoint)?;
169        Ok(PreparedCaptureRestore {
170            run: self,
171            prediction: checkpoint.prediction,
172            phase: checkpoint.phase,
173            has_step: checkpoint.has_step,
174        })
175    }
176
177    /// Cumulative usage of this run, including the explicitly inherited child base.
178    pub fn cumulative_usage(&self) -> CaptureUsage {
179        self.ledger.total()
180    }
181}
182
183impl CaptureCheckpoint {
184    /// Logical storage of this immutable admission/schedule checkpoint.
185    pub fn logical_storage_bytes(&self) -> Option<u64> {
186        checkpoint_storage_bytes(
187            &self.plan,
188            self.intervention.as_ref(),
189            &self.artifact_identity,
190        )
191    }
192
193    /// Conservative logical host storage for shared child re-admission, measured
194    /// against actual child declarations before cloning plans or payloads.
195    pub fn fork_storage_bytes(&self, request: &CaptureForkRequest<'_>) -> Option<u64> {
196        use crate::execution_control::storage::heap_bytes;
197        let mut bytes = self
198            .logical_storage_bytes()?
199            .checked_add(u64::try_from(std::mem::size_of::<CaptureSession>()).ok()?)?;
200        for selection in &self.plan.plan().selections {
201            let point = request
202                .discovery
203                .catalog
204                .points
205                .iter()
206                .find(|point| point.path == selection.path)?;
207            bytes = bytes
208                .checked_add(u64::try_from(std::mem::size_of_val(point)).ok()?)?
209                .checked_add(heap_bytes(point)?)?;
210        }
211        if let Some(child) = &request.intervention {
212            let plan = child
213                .replacement
214                .as_ref()
215                .or_else(|| self.intervention.as_ref().map(|plan| plan.plan()))?;
216            bytes = bytes
217                .checked_add(u64::try_from(std::mem::size_of::<AdmittedInterventionPlan>()).ok()?)?
218                .checked_add(heap_bytes(plan)?)?
219                .checked_add(u64::try_from(child.discovery.artifact_identity.len()).ok()?)?
220                .checked_add(u64::try_from(child.session_id.len()).ok()?)?
221                .checked_add(64)?;
222            for operation in &plan.operations {
223                let point = child
224                    .discovery
225                    .points
226                    .iter()
227                    .find(|point| point.path == operation.target)?;
228                bytes = bytes
229                    .checked_add(u64::try_from(std::mem::size_of_val(point)).ok()?)?
230                    .checked_add(heap_bytes(point)?)?;
231            }
232        }
233        Some(bytes)
234    }
235    /// Absolute next prediction; zero means prompt prefill has not executed.
236    pub fn next_prediction(&self) -> u64 {
237        // begin_step requires prediction < max_predictions, hence this cannot overflow.
238        if self.has_step {
239            self.prediction + 1
240        } else {
241            0
242        }
243    }
244
245    /// Usage inherited by a child. Same-run restore never writes this to the ledger.
246    pub fn inherited_usage(&self) -> CaptureUsage {
247        self.usage
248    }
249
250    /// Original source provenance, retained when admitting child plans.
251    pub fn artifact_identity(&self) -> &str {
252        &self.artifact_identity
253    }
254
255    /// Original intervention admission for lineage/override provenance.
256    pub fn intervention_plan(&self) -> Option<&AdmittedInterventionPlan> {
257        self.intervention.as_ref()
258    }
259
260    /// Re-admits child plans with current discovery and estimator facts before
261    /// installing a new shared run. Absolute schedules keep their original origin.
262    /// Child budgets include consumption at this checkpoint; subsequent work in
263    /// either parent or child is charged independently. No native state is copied.
264    pub fn fork(
265        &self,
266        request: CaptureForkRequest<'_>,
267        estimate: impl FnMut(
268            &[u64],
269            &CaptureSelection,
270            &ResolvedCaptureSlice,
271        ) -> Result<CaptureUsage, CaptureError>,
272    ) -> Result<CaptureSession, CaptureError> {
273        if request.discovery.artifact_identity != self.artifact_identity {
274            return Err(CaptureError::Invalid(
275                "child prepared source differs from checkpoint".into(),
276            ));
277        }
278        let mut geometry = self.plan.request();
279        geometry.max_predictions = request.max_predictions;
280        let mut plan = self.plan.plan().clone();
281        plan.limits = request.limits;
282        let plan = plan.admit(
283            &request.discovery.catalog,
284            &request.discovery.support,
285            &request.discovery.support.capture,
286            geometry,
287        )?;
288        super::validate_continuation(
289            &plan,
290            request.discovery,
291            self.next_prediction(),
292            self.usage,
293            estimate,
294        )?;
295
296        let intervention = match request.intervention {
297            Some(child) => {
298                if child.session_id.is_empty()
299                    || self
300                        .intervention
301                        .as_ref()
302                        .is_some_and(|parent| parent.session_id() == child.session_id)
303                    || child.discovery.artifact_identity != self.artifact_identity
304                {
305                    return Err(CaptureError::Invalid(
306                        "child intervention identity/source mismatch".into(),
307                    ));
308                }
309                let operations = child
310                    .replacement
311                    .or_else(|| self.intervention.as_ref().map(|p| p.plan().clone()))
312                    .ok_or_else(|| {
313                        CaptureError::Invalid("child intervention plan is absent".into())
314                    })?;
315                let admitted = operations.admit(child.discovery, geometry, child.session_id)?;
316                crate::intervention::validate_continuation(
317                    &plan,
318                    &admitted,
319                    child.discovery,
320                    child.estimator.as_ref(),
321                    self.next_prediction(),
322                    self.usage,
323                )?;
324                Some((admitted, child.estimator))
325            }
326            None if self.intervention.is_some() => {
327                return Err(CaptureError::Invalid(
328                    "inherited interventions require child discovery and re-admission".into(),
329                ))
330            }
331            None => None,
332        };
333        // Installation stays with the same shared owner. A fully empty child still
334        // carries its inherited ledger, so removing controls cannot erase usage.
335        let mut child = CaptureSession::new(plan);
336        if let Some((plan, estimator)) = intervention {
337            child.enable_interventions(plan, estimator)?;
338        }
339        child.ledger = CaptureLedger::with_inherited_usage(&child.plan, self.usage)?;
340        child.prediction = self.prediction;
341        child.phase = self.phase;
342        child.has_step = self.has_step;
343        Ok(child)
344    }
345}
346
347fn checkpoint_storage_bytes(
348    plan: &AdmittedCapturePlan,
349    intervention: Option<&AdmittedInterventionPlan>,
350    artifact: &str,
351) -> Option<u64> {
352    use crate::execution_control::storage::heap_bytes;
353    let mut total = u64::try_from(std::mem::size_of::<CaptureCheckpoint>())
354        .ok()?
355        .checked_add(u64::try_from(artifact.len()).ok()?)?
356        .checked_add(heap_bytes(plan.plan())?)?
357        .checked_add(heap_bytes(plan.points())?)?
358        .checked_add(u64::try_from(plan.identity().len()).ok()?)?;
359    if let Some(plan) = intervention {
360        total = total
361            .checked_add(heap_bytes(plan.plan())?)?
362            .checked_add(heap_bytes(plan.points())?)?
363            .checked_add(u64::try_from(plan.identity().len()).ok()?)?
364            .checked_add(u64::try_from(plan.artifact_identity().len()).ok()?)?
365            .checked_add(u64::try_from(plan.session_id().len()).ok()?)?;
366    }
367    Some(total)
368}