1use std::fmt;
30use std::fs::{self, File};
31use std::io::{BufRead, BufReader};
32use std::path::{Path, PathBuf};
33
34use serde::de::{SeqAccess, Visitor};
35use serde::{Deserialize, Deserializer, Serialize};
36use serde_json::value::RawValue;
37use sha2::{Digest, Sha256};
38
39use crate::system_state::{SimulationTime, SystemState, SystemStateSchema};
40use crate::time_series::StateSeries;
41
42use super::RecordingTiming;
43use super::error::StorageError;
44use super::json_payload_decoder::JsonPayloadDecoderRegistry;
45use super::jsonl_format::{ChunkMetadata, RecordingMetadata, RecordingStatus, StateStreamMetadata};
46
47const METADATA_FILE: &str = "metadata.json";
49
50pub struct StoredStateSeriesReader {
57 root: PathBuf,
58 metadata_path: PathBuf,
59 metadata: RecordingMetadata,
60 timing: RecordingTiming,
61 decoders: JsonPayloadDecoderRegistry,
62}
63
64impl StoredStateSeriesReader {
65 pub fn open_completed_recording(
82 root: impl AsRef<Path>,
83 decoders: JsonPayloadDecoderRegistry,
84 ) -> Result<Self, StorageError> {
85 let root = root.as_ref().to_path_buf();
86 let metadata_path = root.join(METADATA_FILE);
87 let bytes = fs::read(&metadata_path).map_err(|source| StorageError::Io {
88 operation: "read metadata",
89 path: metadata_path.clone(),
90 source,
91 })?;
92 let metadata: RecordingMetadata =
93 serde_json::from_slice(&bytes).map_err(|source| StorageError::Json {
94 operation: "parse metadata",
95 path: metadata_path.clone(),
96 source,
97 })?;
98 metadata.validate(&metadata_path)?;
99 if !matches!(metadata.status, RecordingStatus::Complete) {
100 return Err(StorageError::RecordingNotComplete {
101 path: metadata_path,
102 });
103 }
104 let timing = RecordingTiming::from_stored(&metadata.timing, &metadata_path)?;
105 Ok(Self {
106 root,
107 metadata_path,
108 metadata,
109 timing,
110 decoders,
111 })
112 }
113
114 pub fn recording_directory(&self) -> &Path {
116 &self.root
117 }
118
119 pub fn stream_names(&self) -> impl ExactSizeIterator<Item = &str> {
121 self.metadata
122 .streams
123 .iter()
124 .map(|stream| stream.name.as_str())
125 }
126
127 pub fn format_version(&self) -> u32 {
129 self.metadata.version
130 }
131
132 pub fn user_metadata(&self) -> &serde_json::Map<String, serde_json::Value> {
134 &self.metadata.user_metadata
135 }
136
137 pub fn terminal_metadata(&self) -> &serde_json::Map<String, serde_json::Value> {
139 &self.metadata.terminal_metadata
140 }
141
142 pub fn recording_timing(&self) -> &RecordingTiming {
144 &self.timing
145 }
146
147 pub fn stream_record_count(&self, stream: &str) -> Result<u64, StorageError> {
149 let declaration =
150 self.metadata
151 .stream(stream)
152 .ok_or_else(|| StorageError::UnknownStateStream {
153 stream: stream.to_owned(),
154 })?;
155 declaration
156 .chunks
157 .iter()
158 .try_fold(0_u64, |total, chunk| total.checked_add(chunk.records))
159 .ok_or_else(|| StorageError::ByteCountOverflow {
160 stream: declaration.name.clone(),
161 })
162 }
163
164 pub fn stream_encoded_bytes(&self, stream: &str) -> Result<u64, StorageError> {
166 let declaration =
167 self.metadata
168 .stream(stream)
169 .ok_or_else(|| StorageError::UnknownStateStream {
170 stream: stream.to_owned(),
171 })?;
172 declaration
173 .chunks
174 .iter()
175 .try_fold(0_u64, |total, chunk| total.checked_add(chunk.bytes))
176 .ok_or_else(|| StorageError::ByteCountOverflow {
177 stream: declaration.name.clone(),
178 })
179 }
180
181 pub fn read_stream_as_state_series(&self, stream: &str) -> Result<StateSeries, StorageError> {
194 let declaration =
195 self.metadata
196 .stream(stream)
197 .ok_or_else(|| StorageError::UnknownStateStream {
198 stream: stream.to_owned(),
199 })?;
200 self.decoders
201 .require(declaration.fields.iter().map(|field| field.name.as_str()))?;
202
203 let spec = stream_spec(&self.metadata_path, declaration)?;
204 let total_records = self.stream_record_count(stream)?;
205 let capacity =
206 usize::try_from(total_records).map_err(|_| StorageError::ByteCountOverflow {
207 stream: declaration.name.clone(),
208 })?;
209 let mut series = StateSeries::with_capacity(spec, capacity);
210 let mut previous_iteration = None;
211
212 for chunk in &declaration.chunks {
213 self.read_chunk(declaration, chunk, &mut previous_iteration, &mut series)?;
214 }
215 Ok(series)
216 }
217
218 pub fn read_all_streams_as_state_series(
224 &self,
225 ) -> Result<Vec<(String, StateSeries)>, StorageError> {
226 self.metadata
227 .streams
228 .iter()
229 .map(|stream| {
230 self.read_stream_as_state_series(&stream.name)
231 .map(|series| (stream.name.clone(), series))
232 })
233 .collect()
234 }
235
236 pub fn read_latest_state_from_stream(&self, stream: &str) -> Result<SystemState, StorageError> {
244 let declaration =
245 self.metadata
246 .stream(stream)
247 .ok_or_else(|| StorageError::UnknownStateStream {
248 stream: stream.to_owned(),
249 })?;
250 self.decoders
251 .require(declaration.fields.iter().map(|field| field.name.as_str()))?;
252 let chunk = declaration
253 .chunks
254 .last()
255 .ok_or_else(|| StorageError::NoRecordedState {
256 stream: declaration.name.clone(),
257 })?;
258 let path = self.root.join(&declaration.directory).join(&chunk.file);
259 let bytes = read_verified_chunk(&self.metadata_path, &path, chunk)?;
260 let record = final_jsonl_record(&path, &bytes)?;
261 let spec = stream_spec(&self.metadata_path, declaration)?;
262 let state =
263 decode_state_record_with_decoders(record, &path, declaration, &spec, &self.decoders)?;
264 if state.simulation_time().iteration() != chunk.last_iteration {
265 return Err(invalid_record(
266 &path,
267 chunk.records,
268 format!(
269 "latest record iteration {} differs from chunk descriptor {}",
270 state.simulation_time().iteration(),
271 chunk.last_iteration
272 ),
273 ));
274 }
275 Ok(state)
276 }
277
278 fn read_chunk(
280 &self,
281 stream: &StateStreamMetadata,
282 chunk: &ChunkMetadata,
283 previous_iteration: &mut Option<u64>,
284 series: &mut StateSeries,
285 ) -> Result<(), StorageError> {
286 let path = self.root.join(&stream.directory).join(&chunk.file);
287 verify_file_size(&path, chunk.bytes)?;
288 let file = File::open(&path).map_err(|source| {
289 if source.kind() == std::io::ErrorKind::NotFound {
290 StorageError::MissingChunk { path: path.clone() }
291 } else {
292 StorageError::Io {
293 operation: "open chunk",
294 path: path.clone(),
295 source,
296 }
297 }
298 })?;
299 let mut input = BufReader::new(file);
300 let mut line = Vec::new();
301 let mut line_number = 0_u64;
302 let mut records = 0_u64;
303 let mut first_iteration = None;
304 let mut last_iteration = None;
305 let mut hasher = Sha256::new();
306
307 loop {
308 line.clear();
309 let bytes_read =
310 input
311 .read_until(b'\n', &mut line)
312 .map_err(|source| StorageError::Io {
313 operation: "read chunk",
314 path: path.clone(),
315 source,
316 })?;
317 if bytes_read == 0 {
318 break;
319 }
320 line_number =
321 line_number
322 .checked_add(1)
323 .ok_or_else(|| StorageError::ByteCountOverflow {
324 stream: stream.name.clone(),
325 })?;
326 hasher.update(&line);
327 if line.last() != Some(&b'\n') {
328 return Err(invalid_record(
329 &path,
330 line_number,
331 "record is not terminated by a newline",
332 ));
333 }
334 line.pop();
335 if line.is_empty() {
336 return Err(invalid_record(
337 &path,
338 line_number,
339 "record line must not be empty",
340 ));
341 }
342
343 let record: BorrowedRecord<'_> = serde_json::from_slice(&line).map_err(|source| {
344 invalid_record(&path, line_number, format!("invalid JSON record: {source}"))
345 })?;
346 validate_iteration(&path, line_number, record.iteration, *previous_iteration)?;
347 let iteration = record.iteration;
348 first_iteration.get_or_insert(iteration);
349 last_iteration = Some(iteration);
350 *previous_iteration = Some(iteration);
351
352 let time = match record.physical_time {
353 Some(physical_time) => {
354 SimulationTime::from_iteration_and_physical_time(iteration, physical_time)
355 .ok_or_else(|| {
356 invalid_record(&path, line_number, "physical time must be finite")
357 })?
358 }
359 None => SimulationTime::from_iteration(iteration),
360 };
361 let mut state = series.schema().create_empty_state(time);
362 decode_values(
363 &self.decoders,
364 stream,
365 &path,
366 line_number,
367 record.values,
368 &mut state,
369 )?;
370 series.push_state(state).map_err(|rejection| {
371 let (source, state) = rejection.into_parts();
372 let index = state.simulation_time().iteration();
373 drop(state);
374 StorageError::StateSeriesInvariant {
375 stream: stream.name.clone(),
376 iteration: index,
377 source,
378 }
379 })?;
380 records = records
381 .checked_add(1)
382 .ok_or_else(|| StorageError::ByteCountOverflow {
383 stream: stream.name.clone(),
384 })?;
385 }
386
387 validate_chunk_facts(
388 &self.metadata_path,
389 stream,
390 chunk,
391 records,
392 first_iteration,
393 last_iteration,
394 )?;
395 verify_checksum(
396 &self.metadata_path,
397 &path,
398 &chunk.checksum,
399 hasher.finalize(),
400 )
401 }
402}
403
404fn final_jsonl_record<'a>(path: &Path, bytes: &'a [u8]) -> Result<&'a [u8], StorageError> {
406 if !bytes.ends_with(b"\n") {
407 return Err(invalid_record(
408 path,
409 1,
410 "latest chunk is not terminated by a newline",
411 ));
412 }
413 let without_final_newline = &bytes[..bytes.len() - 1];
414 let start = without_final_newline
415 .iter()
416 .rposition(|byte| *byte == b'\n')
417 .map_or(0, |position| position + 1);
418 let record = &without_final_newline[start..];
419 if record.is_empty() {
420 return Err(invalid_record(path, 1, "latest record must not be empty"));
421 }
422 Ok(record)
423}
424
425impl fmt::Debug for StoredStateSeriesReader {
426 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
428 formatter
429 .debug_struct("StoredStateSeriesReader")
430 .field("root", &self.root)
431 .field("streams", &self.metadata.streams.len())
432 .field("decoders", &self.decoders)
433 .finish_non_exhaustive()
434 }
435}
436
437pub(crate) fn decode_resume_state(
443 root: &Path,
444 metadata_path: &Path,
445 stream: &StateStreamMetadata,
446 full_spec: &SystemStateSchema,
447 decoders: &JsonPayloadDecoderRegistry,
448) -> Result<SystemState, StorageError> {
449 validate_complete_resume_schema(stream, full_spec)?;
450 decoders.require(full_spec.field_schemas().iter().map(|field| field.name()))?;
451
452 let chunk = stream
453 .chunks
454 .last()
455 .ok_or_else(|| StorageError::NoCheckpointState {
456 stream: stream.name.clone(),
457 })?;
458 let path = root.join(&stream.directory).join(&chunk.file);
459 let bytes = read_verified_chunk(metadata_path, &path, chunk)?;
460 let record = final_jsonl_record(&path, &bytes)?;
461 let state = decode_state_record_with_decoders(record, &path, stream, full_spec, decoders)?;
462 if state.simulation_time().iteration() != chunk.last_iteration {
463 return Err(invalid_record(
464 &path,
465 chunk.records,
466 format!(
467 "latest record iteration {} differs from chunk descriptor {}",
468 state.simulation_time().iteration(),
469 chunk.last_iteration
470 ),
471 ));
472 }
473 Ok(state)
474}
475
476fn validate_complete_resume_schema(
478 stream: &StateStreamMetadata,
479 full_spec: &SystemStateSchema,
480) -> Result<(), StorageError> {
481 if stream.fields.len() != full_spec.len() {
482 return Err(StorageError::IncompleteCheckpointStream {
483 stream: stream.name.clone(),
484 reason: format!(
485 "stream declares {} fields but the full state declares {}",
486 stream.fields.len(),
487 full_spec.len()
488 ),
489 });
490 }
491 for (position, (stored, expected)) in stream
492 .fields
493 .iter()
494 .zip(full_spec.field_schemas())
495 .enumerate()
496 {
497 if stored.name != expected.name() || stored.description.as_deref() != expected.description()
498 {
499 return Err(StorageError::IncompleteCheckpointStream {
500 stream: stream.name.clone(),
501 reason: format!(
502 "field {position} is `{}` but full state requires `{}`",
503 stored.name,
504 expected.name()
505 ),
506 });
507 }
508 }
509 Ok(())
510}
511
512fn decode_state_record_with_decoders(
514 record: &[u8],
515 path: &Path,
516 stream: &StateStreamMetadata,
517 full_spec: &SystemStateSchema,
518 decoders: &JsonPayloadDecoderRegistry,
519) -> Result<SystemState, StorageError> {
520 let record: BorrowedRecord<'_> = serde_json::from_slice(record)
521 .map_err(|source| invalid_record(path, 1, format!("invalid JSON record: {source}")))?;
522 let time = match record.physical_time {
523 Some(physical_time) => {
524 SimulationTime::from_iteration_and_physical_time(record.iteration, physical_time)
525 .ok_or_else(|| invalid_record(path, 1, "physical time must be finite"))?
526 }
527 None => SimulationTime::from_iteration(record.iteration),
528 };
529 let mut state = full_spec.create_empty_state(time);
530 decode_values(decoders, stream, path, 1, record.values, &mut state)?;
531 Ok(state)
532}
533
534#[derive(Deserialize)]
536#[serde(deny_unknown_fields)]
537struct BorrowedRecord<'a> {
538 iteration: u64,
539 #[serde(default)]
540 physical_time: Option<f64>,
541 #[serde(borrow)]
542 values: BorrowedValues<'a>,
543}
544
545struct BorrowedValues<'a> {
547 entries: Vec<&'a RawValue>,
548}
549
550impl<'de: 'a, 'a> Deserialize<'de> for BorrowedValues<'a> {
551 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
553 where
554 D: Deserializer<'de>,
555 {
556 deserializer.deserialize_seq(BorrowedValuesVisitor {
557 output: std::marker::PhantomData,
558 })
559 }
560}
561
562struct BorrowedValuesVisitor<'a> {
564 output: std::marker::PhantomData<&'a RawValue>,
566}
567
568impl<'de: 'a, 'a> Visitor<'de> for BorrowedValuesVisitor<'a> {
569 type Value = BorrowedValues<'a>;
570
571 fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
573 formatter.write_str("an array of raw JSON payload values")
574 }
575
576 fn visit_seq<A>(self, mut sequence: A) -> Result<Self::Value, A::Error>
578 where
579 A: SeqAccess<'de>,
580 {
581 let mut entries = Vec::with_capacity(sequence.size_hint().unwrap_or(0));
582 while let Some(value) = sequence.next_element::<&'de RawValue>()? {
583 let value: &'a RawValue = value;
584 entries.push(value);
585 }
586 Ok(BorrowedValues { entries })
587 }
588}
589
590#[derive(Serialize)]
592struct StreamTemplateRef<'a> {
593 fields: &'a [super::jsonl_format::StateFieldMetadata],
594}
595
596fn stream_spec(
598 metadata_path: &Path,
599 stream: &StateStreamMetadata,
600) -> Result<SystemStateSchema, StorageError> {
601 let bytes = serde_json::to_vec(&StreamTemplateRef {
602 fields: &stream.fields,
603 })
604 .map_err(|source| StorageError::Json {
605 operation: "serialize stream schema",
606 path: metadata_path.to_path_buf(),
607 source,
608 })?;
609 SystemStateSchema::parse(metadata_path.to_path_buf(), &bytes).map_err(|source| {
610 StorageError::InvalidMetadata {
611 path: metadata_path.to_path_buf(),
612 reason: format!(
613 "stream `{}` has an invalid state schema: {source}",
614 stream.name
615 ),
616 }
617 })
618}
619
620fn verify_file_size(path: &Path, expected: u64) -> Result<(), StorageError> {
622 let actual = match fs::metadata(path) {
623 Ok(metadata) => metadata.len(),
624 Err(source) if source.kind() == std::io::ErrorKind::NotFound => {
625 return Err(StorageError::MissingChunk {
626 path: path.to_path_buf(),
627 });
628 }
629 Err(source) => {
630 return Err(StorageError::Io {
631 operation: "inspect chunk",
632 path: path.to_path_buf(),
633 source,
634 });
635 }
636 };
637 if actual != expected {
638 return Err(StorageError::ChunkSizeMismatch {
639 path: path.to_path_buf(),
640 expected,
641 actual,
642 });
643 }
644 Ok(())
645}
646
647fn read_verified_chunk(
649 metadata_path: &Path,
650 path: &Path,
651 chunk: &ChunkMetadata,
652) -> Result<Vec<u8>, StorageError> {
653 verify_file_size(path, chunk.bytes)?;
654 let bytes = fs::read(path).map_err(|source| {
655 if source.kind() == std::io::ErrorKind::NotFound {
656 StorageError::MissingChunk {
657 path: path.to_path_buf(),
658 }
659 } else {
660 StorageError::Io {
661 operation: "read verified chunk",
662 path: path.to_path_buf(),
663 source,
664 }
665 }
666 })?;
667 verify_checksum(metadata_path, path, &chunk.checksum, Sha256::digest(&bytes))?;
668 Ok(bytes)
669}
670
671fn validate_iteration(
673 path: &Path,
674 line: u64,
675 iteration: u64,
676 previous: Option<u64>,
677) -> Result<(), StorageError> {
678 if let Some(previous) = previous
679 && iteration <= previous
680 {
681 return Err(invalid_record(
682 path,
683 line,
684 format!("iteration {iteration} is not greater than previous iteration {previous}"),
685 ));
686 }
687 Ok(())
688}
689
690fn decode_values(
692 decoders: &JsonPayloadDecoderRegistry,
693 stream: &StateStreamMetadata,
694 path: &Path,
695 line: u64,
696 values: BorrowedValues<'_>,
697 state: &mut crate::system_state::SystemState,
698) -> Result<(), StorageError> {
699 if values.entries.len() != stream.fields.len() {
700 return Err(invalid_record(
701 path,
702 line,
703 format!(
704 "record contains {} payload values but stream `{}` declares {} fields",
705 values.entries.len(),
706 stream.name,
707 stream.fields.len()
708 ),
709 ));
710 }
711 for (field, raw) in stream.fields.iter().zip(values.entries) {
712 decoders.decode_into(
713 &stream.name,
714 state.simulation_time().iteration(),
715 &field.name,
716 raw.get(),
717 state,
718 )?;
719 }
720 Ok(())
721}
722
723fn validate_chunk_facts(
725 metadata_path: &Path,
726 stream: &StateStreamMetadata,
727 chunk: &ChunkMetadata,
728 records: u64,
729 first_iteration: Option<u64>,
730 last_iteration: Option<u64>,
731) -> Result<(), StorageError> {
732 if records != chunk.records
733 || first_iteration != Some(chunk.first_iteration)
734 || last_iteration != Some(chunk.last_iteration)
735 {
736 return Err(StorageError::InvalidMetadata {
737 path: metadata_path.to_path_buf(),
738 reason: format!(
739 "stream `{}` chunk {} declares {} records at {}..={}, but contains {} records at {:?}..={:?}",
740 stream.name,
741 chunk.ordinal,
742 chunk.records,
743 chunk.first_iteration,
744 chunk.last_iteration,
745 records,
746 first_iteration,
747 last_iteration
748 ),
749 });
750 }
751 Ok(())
752}
753
754fn verify_checksum(
756 metadata_path: &Path,
757 path: &Path,
758 expected: &str,
759 digest: impl AsRef<[u8]>,
760) -> Result<(), StorageError> {
761 let Some(expected_digest) = expected.strip_prefix("sha256:") else {
762 return Err(StorageError::InvalidMetadata {
763 path: metadata_path.to_path_buf(),
764 reason: format!("unsupported chunk checksum algorithm in `{expected}`"),
765 });
766 };
767 let actual_digest = lowercase_hex(digest.as_ref());
768 if actual_digest != expected_digest {
769 return Err(StorageError::ChecksumMismatch {
770 path: path.to_path_buf(),
771 expected: expected.to_owned(),
772 actual: format!("sha256:{actual_digest}"),
773 });
774 }
775 Ok(())
776}
777
778fn lowercase_hex(bytes: &[u8]) -> String {
780 use std::fmt::Write as _;
781
782 let mut encoded = String::with_capacity(bytes.len() * 2);
783 for byte in bytes {
784 write!(&mut encoded, "{byte:02x}").expect("writing to String cannot fail");
785 }
786 encoded
787}
788
789fn invalid_record(path: &Path, line: u64, reason: impl Into<String>) -> StorageError {
791 StorageError::InvalidRecord {
792 path: path.to_path_buf(),
793 line,
794 reason: reason.into(),
795 }
796}