1use super::*;
4use eredu_core::intervention::{
5 AdmittedInterventionPlan, InterventionDiscovery, InterventionEstimator, InterventionPlan,
6};
7use std::sync::Arc;
8
9#[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
27pub struct CaptureForkRequest<'a> {
30 pub discovery: &'a CaptureDiscovery,
32 pub max_predictions: u64,
34 pub limits: CaptureLimits,
36 pub intervention: Option<InterventionForkRequest<'a>>,
38}
39
40pub struct InterventionForkRequest<'a> {
42 pub discovery: &'a InterventionDiscovery,
44 pub session_id: &'a str,
46 pub replacement: Option<InterventionPlan>,
49 pub estimator: Arc<dyn InterventionEstimator>,
51}
52
53pub struct PreparedCaptureRestore<'a> {
57 run: &'a mut CaptureSession,
58 prediction: u64,
59 phase: CapturePhase,
60 has_step: bool,
61}
62
63impl PreparedCaptureRestore<'_> {
64 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 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 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 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 pub fn restore(&mut self, checkpoint: &CaptureCheckpoint) -> Result<(), CaptureError> {
159 self.prepare_restore(checkpoint)?.commit();
160 Ok(())
161 }
162
163 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 pub fn cumulative_usage(&self) -> CaptureUsage {
179 self.ledger.total()
180 }
181}
182
183impl CaptureCheckpoint {
184 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 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 pub fn next_prediction(&self) -> u64 {
237 if self.has_step {
239 self.prediction + 1
240 } else {
241 0
242 }
243 }
244
245 pub fn inherited_usage(&self) -> CaptureUsage {
247 self.usage
248 }
249
250 pub fn artifact_identity(&self) -> &str {
252 &self.artifact_identity
253 }
254
255 pub fn intervention_plan(&self) -> Option<&AdmittedInterventionPlan> {
257 self.intervention.as_ref()
258 }
259
260 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 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}