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