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