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