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 let Some(mismatch) = checkpoint_schema_mismatch(stream, full_spec) else {
482 return Ok(());
483 };
484 Err(match mismatch {
485 CheckpointSchemaMismatch::Count {
486 stream_fields,
487 full_spec_fields,
488 } => StorageError::IncompleteCheckpointStream {
489 stream: stream.name.clone(),
490 reason: format!(
491 "stream declares {stream_fields} fields but the full state declares {full_spec_fields}",
492 ),
493 },
494 CheckpointSchemaMismatch::Field { position, stored, expected } => {
495 StorageError::IncompleteCheckpointStream {
496 stream: stream.name.clone(),
497 reason: format!("field {position} is `{stored}` but full state requires `{expected}`"),
498 }
499 }
500 })
501}
502
503pub(crate) fn is_complete_checkpoint_stream(
504 stream: &StateStreamMetadata,
505 full_spec: &SystemStateSchema,
506) -> bool {
507 checkpoint_schema_mismatch(stream, full_spec).is_none()
508}
509
510enum CheckpointSchemaMismatch {
511 Count {
512 stream_fields: usize,
513 full_spec_fields: usize,
514 },
515 Field {
516 position: usize,
517 stored: String,
518 expected: String,
519 },
520}
521
522fn checkpoint_schema_mismatch(
523 stream: &StateStreamMetadata,
524 full_spec: &SystemStateSchema,
525) -> Option<CheckpointSchemaMismatch> {
526 if stream.fields.len() != full_spec.len() {
527 return Some(CheckpointSchemaMismatch::Count {
528 stream_fields: stream.fields.len(),
529 full_spec_fields: full_spec.len(),
530 });
531 }
532 for (position, (stored, expected)) in stream
533 .fields
534 .iter()
535 .zip(full_spec.field_schemas())
536 .enumerate()
537 {
538 if stored.name != expected.name() || stored.description.as_deref() != expected.description()
539 {
540 return Some(CheckpointSchemaMismatch::Field {
541 position,
542 stored: stored.name.clone(),
543 expected: expected.name().to_string(),
544 });
545 }
546 }
547 None
548}
549
550fn decode_state_record_with_decoders(
552 record: &[u8],
553 path: &Path,
554 stream: &StateStreamMetadata,
555 full_spec: &SystemStateSchema,
556 decoders: &JsonPayloadDecoderRegistry,
557) -> Result<SystemState, StorageError> {
558 let record: BorrowedRecord<'_> = serde_json::from_slice(record)
559 .map_err(|source| invalid_record(path, 1, format!("invalid JSON record: {source}")))?;
560 let time = match record.physical_time {
561 Some(physical_time) => {
562 SimulationTime::from_iteration_and_physical_time(record.iteration, physical_time)
563 .ok_or_else(|| invalid_record(path, 1, "physical time must be finite"))?
564 }
565 None => SimulationTime::from_iteration(record.iteration),
566 };
567 let mut state = full_spec.create_empty_state(time);
568 decode_values(decoders, stream, path, 1, record.values, &mut state)?;
569 Ok(state)
570}
571
572#[derive(Deserialize)]
574#[serde(deny_unknown_fields)]
575struct BorrowedRecord<'a> {
576 iteration: u64,
577 #[serde(default)]
578 physical_time: Option<f64>,
579 #[serde(borrow)]
580 values: BorrowedValues<'a>,
581}
582
583struct BorrowedValues<'a> {
585 entries: Vec<&'a RawValue>,
586}
587
588impl<'de: 'a, 'a> Deserialize<'de> for BorrowedValues<'a> {
589 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
591 where
592 D: Deserializer<'de>,
593 {
594 deserializer.deserialize_seq(BorrowedValuesVisitor {
595 output: std::marker::PhantomData,
596 })
597 }
598}
599
600struct BorrowedValuesVisitor<'a> {
602 output: std::marker::PhantomData<&'a RawValue>,
604}
605
606impl<'de: 'a, 'a> Visitor<'de> for BorrowedValuesVisitor<'a> {
607 type Value = BorrowedValues<'a>;
608
609 fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
611 formatter.write_str("an array of raw JSON payload values")
612 }
613
614 fn visit_seq<A>(self, mut sequence: A) -> Result<Self::Value, A::Error>
616 where
617 A: SeqAccess<'de>,
618 {
619 let mut entries = Vec::with_capacity(sequence.size_hint().unwrap_or(0));
620 while let Some(value) = sequence.next_element::<&'de RawValue>()? {
621 let value: &'a RawValue = value;
622 entries.push(value);
623 }
624 Ok(BorrowedValues { entries })
625 }
626}
627
628#[derive(Serialize)]
630struct StreamTemplateRef<'a> {
631 fields: &'a [super::jsonl_format::StateFieldMetadata],
632}
633
634fn stream_spec(
636 metadata_path: &Path,
637 stream: &StateStreamMetadata,
638) -> Result<SystemStateSchema, StorageError> {
639 let bytes = serde_json::to_vec(&StreamTemplateRef {
640 fields: &stream.fields,
641 })
642 .map_err(|source| StorageError::Json {
643 operation: "serialize stream schema",
644 path: metadata_path.to_path_buf(),
645 source,
646 })?;
647 SystemStateSchema::parse(metadata_path.to_path_buf(), &bytes).map_err(|source| {
648 StorageError::InvalidMetadata {
649 path: metadata_path.to_path_buf(),
650 reason: format!(
651 "stream `{}` has an invalid state schema: {source}",
652 stream.name
653 ),
654 }
655 })
656}
657
658fn verify_file_size(path: &Path, expected: u64) -> Result<(), StorageError> {
660 let actual = match fs::metadata(path) {
661 Ok(metadata) => metadata.len(),
662 Err(source) if source.kind() == std::io::ErrorKind::NotFound => {
663 return Err(StorageError::MissingChunk {
664 path: path.to_path_buf(),
665 });
666 }
667 Err(source) => {
668 return Err(StorageError::Io {
669 operation: "inspect chunk",
670 path: path.to_path_buf(),
671 source,
672 });
673 }
674 };
675 if actual != expected {
676 return Err(StorageError::ChunkSizeMismatch {
677 path: path.to_path_buf(),
678 expected,
679 actual,
680 });
681 }
682 Ok(())
683}
684
685pub(crate) fn read_verified_chunk(
687 metadata_path: &Path,
688 path: &Path,
689 chunk: &ChunkMetadata,
690) -> Result<Vec<u8>, StorageError> {
691 verify_file_size(path, chunk.bytes)?;
692 let bytes = fs::read(path).map_err(|source| {
693 if source.kind() == std::io::ErrorKind::NotFound {
694 StorageError::MissingChunk {
695 path: path.to_path_buf(),
696 }
697 } else {
698 StorageError::Io {
699 operation: "read verified chunk",
700 path: path.to_path_buf(),
701 source,
702 }
703 }
704 })?;
705 verify_checksum(metadata_path, path, &chunk.checksum, Sha256::digest(&bytes))?;
706 Ok(bytes)
707}
708
709fn validate_iteration(
711 path: &Path,
712 line: u64,
713 iteration: u64,
714 previous: Option<u64>,
715) -> Result<(), StorageError> {
716 if let Some(previous) = previous
717 && iteration <= previous
718 {
719 return Err(invalid_record(
720 path,
721 line,
722 format!("iteration {iteration} is not greater than previous iteration {previous}"),
723 ));
724 }
725 Ok(())
726}
727
728fn decode_values(
730 decoders: &JsonPayloadDecoderRegistry,
731 stream: &StateStreamMetadata,
732 path: &Path,
733 line: u64,
734 values: BorrowedValues<'_>,
735 state: &mut crate::system_state::SystemState,
736) -> Result<(), StorageError> {
737 if values.entries.len() != stream.fields.len() {
738 return Err(invalid_record(
739 path,
740 line,
741 format!(
742 "record contains {} payload values but stream `{}` declares {} fields",
743 values.entries.len(),
744 stream.name,
745 stream.fields.len()
746 ),
747 ));
748 }
749 for (field, raw) in stream.fields.iter().zip(values.entries) {
750 decoders.decode_into(
751 &stream.name,
752 state.simulation_time().iteration(),
753 &field.name,
754 raw.get(),
755 state,
756 )?;
757 }
758 Ok(())
759}
760
761fn validate_chunk_facts(
763 metadata_path: &Path,
764 stream: &StateStreamMetadata,
765 chunk: &ChunkMetadata,
766 records: u64,
767 first_iteration: Option<u64>,
768 last_iteration: Option<u64>,
769) -> Result<(), StorageError> {
770 if records != chunk.records
771 || first_iteration != Some(chunk.first_iteration)
772 || last_iteration != Some(chunk.last_iteration)
773 {
774 return Err(StorageError::InvalidMetadata {
775 path: metadata_path.to_path_buf(),
776 reason: format!(
777 "stream `{}` chunk {} declares {} records at {}..={}, but contains {} records at {:?}..={:?}",
778 stream.name,
779 chunk.ordinal,
780 chunk.records,
781 chunk.first_iteration,
782 chunk.last_iteration,
783 records,
784 first_iteration,
785 last_iteration
786 ),
787 });
788 }
789 Ok(())
790}
791
792fn verify_checksum(
794 metadata_path: &Path,
795 path: &Path,
796 expected: &str,
797 digest: impl AsRef<[u8]>,
798) -> Result<(), StorageError> {
799 let Some(expected_digest) = expected.strip_prefix("sha256:") else {
800 return Err(StorageError::InvalidMetadata {
801 path: metadata_path.to_path_buf(),
802 reason: format!("unsupported chunk checksum algorithm in `{expected}`"),
803 });
804 };
805 let actual_digest = lowercase_hex(digest.as_ref());
806 if actual_digest != expected_digest {
807 return Err(StorageError::ChecksumMismatch {
808 path: path.to_path_buf(),
809 expected: expected.to_owned(),
810 actual: format!("sha256:{actual_digest}"),
811 });
812 }
813 Ok(())
814}
815
816fn lowercase_hex(bytes: &[u8]) -> String {
818 use std::fmt::Write as _;
819
820 let mut encoded = String::with_capacity(bytes.len() * 2);
821 for byte in bytes {
822 write!(&mut encoded, "{byte:02x}").expect("writing to String cannot fail");
823 }
824 encoded
825}
826
827fn invalid_record(path: &Path, line: u64, reason: impl Into<String>) -> StorageError {
829 StorageError::InvalidRecord {
830 path: path.to_path_buf(),
831 line,
832 reason: reason.into(),
833 }
834}