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