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::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 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 Ok(Self {
105 root,
106 metadata_path,
107 metadata,
108 decoders,
109 })
110 }
111
112 pub fn recording_directory(&self) -> &Path {
114 &self.root
115 }
116
117 pub fn stream_names(&self) -> impl ExactSizeIterator<Item = &str> {
119 self.metadata
120 .streams
121 .iter()
122 .map(|stream| stream.name.as_str())
123 }
124
125 pub fn read_stream_as_state_series(&self, stream: &str) -> Result<StateSeries, StorageError> {
138 let declaration =
139 self.metadata
140 .stream(stream)
141 .ok_or_else(|| StorageError::UnknownStateStream {
142 stream: stream.to_owned(),
143 })?;
144 self.decoders
145 .require(declaration.fields.iter().map(|field| field.name.as_str()))?;
146
147 let spec = stream_spec(&self.metadata_path, declaration)?;
148 let capacity = declaration
149 .chunks
150 .iter()
151 .try_fold(0_u64, |total, chunk| total.checked_add(chunk.records))
152 .and_then(|total| usize::try_from(total).ok())
153 .ok_or_else(|| StorageError::ByteCountOverflow {
154 stream: declaration.name.clone(),
155 })?;
156 let mut series = StateSeries::with_capacity(spec, capacity);
157 let mut previous_iteration = None;
158
159 for chunk in &declaration.chunks {
160 self.read_chunk(declaration, chunk, &mut previous_iteration, &mut series)?;
161 }
162 Ok(series)
163 }
164
165 pub fn read_all_streams_as_state_series(
171 &self,
172 ) -> Result<Vec<(String, StateSeries)>, StorageError> {
173 self.metadata
174 .streams
175 .iter()
176 .map(|stream| {
177 self.read_stream_as_state_series(&stream.name)
178 .map(|series| (stream.name.clone(), series))
179 })
180 .collect()
181 }
182
183 fn read_chunk(
185 &self,
186 stream: &StateStreamMetadata,
187 chunk: &ChunkMetadata,
188 previous_iteration: &mut Option<u64>,
189 series: &mut StateSeries,
190 ) -> Result<(), StorageError> {
191 let path = self.root.join(&stream.directory).join(&chunk.file);
192 verify_file_size(&path, chunk.bytes)?;
193 let file = File::open(&path).map_err(|source| {
194 if source.kind() == std::io::ErrorKind::NotFound {
195 StorageError::MissingChunk { path: path.clone() }
196 } else {
197 StorageError::Io {
198 operation: "open chunk",
199 path: path.clone(),
200 source,
201 }
202 }
203 })?;
204 let mut input = BufReader::new(file);
205 let mut line = Vec::new();
206 let mut line_number = 0_u64;
207 let mut records = 0_u64;
208 let mut first_iteration = None;
209 let mut last_iteration = None;
210 let mut hasher = Sha256::new();
211
212 loop {
213 line.clear();
214 let bytes_read =
215 input
216 .read_until(b'\n', &mut line)
217 .map_err(|source| StorageError::Io {
218 operation: "read chunk",
219 path: path.clone(),
220 source,
221 })?;
222 if bytes_read == 0 {
223 break;
224 }
225 line_number =
226 line_number
227 .checked_add(1)
228 .ok_or_else(|| StorageError::ByteCountOverflow {
229 stream: stream.name.clone(),
230 })?;
231 hasher.update(&line);
232 if line.last() != Some(&b'\n') {
233 return Err(invalid_record(
234 &path,
235 line_number,
236 "record is not terminated by a newline",
237 ));
238 }
239 line.pop();
240 if line.is_empty() {
241 return Err(invalid_record(
242 &path,
243 line_number,
244 "record line must not be empty",
245 ));
246 }
247
248 let record: BorrowedRecord<'_> = serde_json::from_slice(&line).map_err(|source| {
249 invalid_record(&path, line_number, format!("invalid JSON record: {source}"))
250 })?;
251 validate_iteration(&path, line_number, record.iteration, *previous_iteration)?;
252 let iteration = record.iteration;
253 first_iteration.get_or_insert(iteration);
254 last_iteration = Some(iteration);
255 *previous_iteration = Some(iteration);
256
257 let time = match record.physical_time {
258 Some(physical_time) => {
259 SimulationTime::from_iteration_and_physical_time(iteration, physical_time)
260 .ok_or_else(|| {
261 invalid_record(&path, line_number, "physical time must be finite")
262 })?
263 }
264 None => SimulationTime::from_iteration(iteration),
265 };
266 let mut state = series.schema().create_empty_state(time);
267 decode_values(
268 &self.decoders,
269 stream,
270 &path,
271 line_number,
272 record.values,
273 &mut state,
274 )?;
275 series.push_state(state).map_err(|rejection| {
276 let (source, state) = rejection.into_parts();
277 let index = state.simulation_time().iteration();
278 drop(state);
279 StorageError::StateSeriesInvariant {
280 stream: stream.name.clone(),
281 iteration: index,
282 source,
283 }
284 })?;
285 records = records
286 .checked_add(1)
287 .ok_or_else(|| StorageError::ByteCountOverflow {
288 stream: stream.name.clone(),
289 })?;
290 }
291
292 validate_chunk_facts(
293 &self.metadata_path,
294 stream,
295 chunk,
296 records,
297 first_iteration,
298 last_iteration,
299 )?;
300 verify_checksum(
301 &self.metadata_path,
302 &path,
303 &chunk.checksum,
304 hasher.finalize(),
305 )
306 }
307}
308
309impl fmt::Debug for StoredStateSeriesReader {
310 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
312 formatter
313 .debug_struct("StoredStateSeriesReader")
314 .field("root", &self.root)
315 .field("streams", &self.metadata.streams.len())
316 .field("decoders", &self.decoders)
317 .finish_non_exhaustive()
318 }
319}
320
321pub(crate) fn decode_resume_state(
328 root: &Path,
329 _metadata_path: &Path,
330 stream: &StateStreamMetadata,
331 full_spec: &SystemStateSchema,
332 decoders: &JsonPayloadDecoderRegistry,
333 open_record: Option<&RecoveredUnsealedRecord>,
334) -> Result<SystemState, StorageError> {
335 validate_complete_resume_schema(stream, full_spec)?;
336 decoders.require(full_spec.field_schemas().iter().map(|field| field.name()))?;
337
338 if let Some(record) = open_record {
339 let bytes = read_open_record(record, &stream.name)?;
340 return decode_state_record_with_decoders(
341 &bytes,
342 record.path(),
343 stream,
344 full_spec,
345 decoders,
346 );
347 }
348
349 let chunk = stream
350 .chunks
351 .last()
352 .ok_or_else(|| StorageError::NoCheckpointState {
353 stream: stream.name.clone(),
354 })?;
355 let path = root.join(&stream.directory).join(&chunk.file);
356 let record = read_last_sealed_record(&path)?;
357 if record.is_empty() {
358 return Err(StorageError::NoCheckpointState {
359 stream: stream.name.clone(),
360 });
361 }
362 decode_state_record_with_decoders(&record, &path, stream, full_spec, decoders)
363}
364
365fn read_open_record(
367 record: &RecoveredUnsealedRecord,
368 stream: &str,
369) -> Result<Vec<u8>, StorageError> {
370 let length = usize::try_from(record.bytes()).map_err(|_| StorageError::ByteCountOverflow {
371 stream: stream.to_owned(),
372 })?;
373 let mut file = File::open(record.path()).map_err(|source| StorageError::Io {
374 operation: "open latest recoverable record",
375 path: record.path().to_path_buf(),
376 source,
377 })?;
378 file.seek(SeekFrom::Start(record.offset()))
379 .map_err(|source| StorageError::Io {
380 operation: "seek latest recoverable record",
381 path: record.path().to_path_buf(),
382 source,
383 })?;
384 let mut bytes = vec![0_u8; length];
385 file.read_exact(&mut bytes)
386 .map_err(|source| StorageError::Io {
387 operation: "read latest recoverable record",
388 path: record.path().to_path_buf(),
389 source,
390 })?;
391 Ok(bytes)
392}
393
394fn validate_complete_resume_schema(
396 stream: &StateStreamMetadata,
397 full_spec: &SystemStateSchema,
398) -> Result<(), StorageError> {
399 if stream.fields.len() != full_spec.len() {
400 return Err(StorageError::IncompleteCheckpointStream {
401 stream: stream.name.clone(),
402 reason: format!(
403 "stream declares {} fields but the full state declares {}",
404 stream.fields.len(),
405 full_spec.len()
406 ),
407 });
408 }
409 for (position, (stored, expected)) in stream
410 .fields
411 .iter()
412 .zip(full_spec.field_schemas())
413 .enumerate()
414 {
415 if stored.name != expected.name() || stored.description.as_deref() != expected.description()
416 {
417 return Err(StorageError::IncompleteCheckpointStream {
418 stream: stream.name.clone(),
419 reason: format!(
420 "field {position} is `{}` but full state requires `{}`",
421 stored.name,
422 expected.name()
423 ),
424 });
425 }
426 }
427 Ok(())
428}
429
430fn decode_state_record_with_decoders(
432 record: &[u8],
433 path: &Path,
434 stream: &StateStreamMetadata,
435 full_spec: &SystemStateSchema,
436 decoders: &JsonPayloadDecoderRegistry,
437) -> Result<SystemState, StorageError> {
438 let record: BorrowedRecord<'_> = serde_json::from_slice(record)
439 .map_err(|source| invalid_record(path, 1, format!("invalid JSON record: {source}")))?;
440 let time = match record.physical_time {
441 Some(physical_time) => {
442 SimulationTime::from_iteration_and_physical_time(record.iteration, physical_time)
443 .ok_or_else(|| invalid_record(path, 1, "physical time must be finite"))?
444 }
445 None => SimulationTime::from_iteration(record.iteration),
446 };
447 let mut state = full_spec.create_empty_state(time);
448 decode_values(decoders, stream, path, 1, record.values, &mut state)?;
449 Ok(state)
450}
451
452fn read_last_sealed_record(path: &Path) -> Result<Vec<u8>, StorageError> {
454 const BLOCK: u64 = 8 * 1024;
455 let mut file = File::open(path).map_err(|source| {
456 if source.kind() == std::io::ErrorKind::NotFound {
457 StorageError::MissingChunk {
458 path: path.to_path_buf(),
459 }
460 } else {
461 StorageError::Io {
462 operation: "open latest sealed checkpoint",
463 path: path.to_path_buf(),
464 source,
465 }
466 }
467 })?;
468 let length = file
469 .metadata()
470 .map_err(|source| StorageError::Io {
471 operation: "inspect latest sealed checkpoint",
472 path: path.to_path_buf(),
473 source,
474 })?
475 .len();
476 if length == 0 {
477 return Ok(Vec::new());
478 }
479
480 file.seek(SeekFrom::Start(length - 1))
481 .map_err(|source| StorageError::Io {
482 operation: "seek latest sealed checkpoint",
483 path: path.to_path_buf(),
484 source,
485 })?;
486 let mut final_byte = [0_u8; 1];
487 file.read_exact(&mut final_byte)
488 .map_err(|source| StorageError::Io {
489 operation: "read latest sealed checkpoint",
490 path: path.to_path_buf(),
491 source,
492 })?;
493 if final_byte[0] != b'\n' {
494 return Err(invalid_record(
495 path,
496 1,
497 "sealed chunk does not end with a record newline",
498 ));
499 }
500
501 let record_end = length - 1;
502 let mut cursor = record_end;
503 let mut record_start = 0_u64;
504 while cursor > 0 {
505 let block_start = cursor.saturating_sub(BLOCK);
506 let block_len =
507 usize::try_from(cursor - block_start).map_err(|_| StorageError::ByteCountOverflow {
508 stream: path.display().to_string(),
509 })?;
510 let mut block = vec![0_u8; block_len];
511 file.seek(SeekFrom::Start(block_start))
512 .and_then(|_| file.read_exact(&mut block))
513 .map_err(|source| StorageError::Io {
514 operation: "scan latest sealed record boundary",
515 path: path.to_path_buf(),
516 source,
517 })?;
518 if let Some(position) = block.iter().rposition(|byte| *byte == b'\n') {
519 record_start = block_start + position as u64 + 1;
520 break;
521 }
522 cursor = block_start;
523 }
524
525 let record_len = usize::try_from(record_end - record_start).map_err(|_| {
526 StorageError::ByteCountOverflow {
527 stream: path.display().to_string(),
528 }
529 })?;
530 let mut record = vec![0_u8; record_len];
531 file.seek(SeekFrom::Start(record_start))
532 .and_then(|_| file.read_exact(&mut record))
533 .map_err(|source| StorageError::Io {
534 operation: "read latest sealed record",
535 path: path.to_path_buf(),
536 source,
537 })?;
538 Ok(record)
539}
540
541#[derive(Deserialize)]
543#[serde(deny_unknown_fields)]
544struct BorrowedRecord<'a> {
545 iteration: u64,
546 #[serde(default)]
547 physical_time: Option<f64>,
548 #[serde(borrow)]
549 values: BorrowedValues<'a>,
550}
551
552struct BorrowedValues<'a> {
554 entries: Vec<(String, &'a RawValue)>,
555}
556
557impl<'de: 'a, 'a> Deserialize<'de> for BorrowedValues<'a> {
558 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
560 where
561 D: Deserializer<'de>,
562 {
563 deserializer.deserialize_map(BorrowedValuesVisitor {
564 output: std::marker::PhantomData,
565 })
566 }
567}
568
569struct BorrowedValuesVisitor<'a> {
571 output: std::marker::PhantomData<&'a RawValue>,
573}
574
575impl<'de: 'a, 'a> Visitor<'de> for BorrowedValuesVisitor<'a> {
576 type Value = BorrowedValues<'a>;
577
578 fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
580 formatter.write_str("an object of unique field keys and raw JSON values")
581 }
582
583 fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
586 where
587 A: MapAccess<'de>,
588 {
589 let mut entries = Vec::with_capacity(map.size_hint().unwrap_or(0));
590 while let Some((key, value)) = map.next_entry::<Cow<'de, str>, &'de RawValue>()? {
591 if entries.iter().any(|(existing, _)| existing == key.as_ref()) {
592 return Err(serde::de::Error::custom(format!(
593 "duplicate payload field `{key}`"
594 )));
595 }
596 let value: &'a RawValue = value;
597 entries.push((key.into_owned(), value));
598 }
599 Ok(BorrowedValues { entries })
600 }
601}
602
603#[derive(Serialize)]
605struct StreamTemplateRef<'a> {
606 fields: &'a [super::jsonl_format::StateFieldMetadata],
607}
608
609fn stream_spec(
611 metadata_path: &Path,
612 stream: &StateStreamMetadata,
613) -> Result<SystemStateSchema, StorageError> {
614 let bytes = serde_json::to_vec(&StreamTemplateRef {
615 fields: &stream.fields,
616 })
617 .map_err(|source| StorageError::Json {
618 operation: "serialize stream schema",
619 path: metadata_path.to_path_buf(),
620 source,
621 })?;
622 SystemStateSchema::parse(metadata_path.to_path_buf(), &bytes).map_err(|source| {
623 StorageError::InvalidMetadata {
624 path: metadata_path.to_path_buf(),
625 reason: format!(
626 "stream `{}` has an invalid state schema: {source}",
627 stream.name
628 ),
629 }
630 })
631}
632
633fn verify_file_size(path: &Path, expected: u64) -> Result<(), StorageError> {
635 let actual = match fs::metadata(path) {
636 Ok(metadata) => metadata.len(),
637 Err(source) if source.kind() == std::io::ErrorKind::NotFound => {
638 return Err(StorageError::MissingChunk {
639 path: path.to_path_buf(),
640 });
641 }
642 Err(source) => {
643 return Err(StorageError::Io {
644 operation: "inspect chunk",
645 path: path.to_path_buf(),
646 source,
647 });
648 }
649 };
650 if actual != expected {
651 return Err(StorageError::ChunkSizeMismatch {
652 path: path.to_path_buf(),
653 expected,
654 actual,
655 });
656 }
657 Ok(())
658}
659
660fn validate_iteration(
662 path: &Path,
663 line: u64,
664 iteration: u64,
665 previous: Option<u64>,
666) -> Result<(), StorageError> {
667 if let Some(previous) = previous
668 && iteration <= previous
669 {
670 return Err(invalid_record(
671 path,
672 line,
673 format!("iteration {iteration} is not greater than previous iteration {previous}"),
674 ));
675 }
676 Ok(())
677}
678
679fn decode_values(
681 decoders: &JsonPayloadDecoderRegistry,
682 stream: &StateStreamMetadata,
683 path: &Path,
684 line: u64,
685 mut values: BorrowedValues<'_>,
686 state: &mut crate::system_state::SystemState,
687) -> Result<(), StorageError> {
688 for field in &stream.fields {
689 let Some(position) = values
690 .entries
691 .iter()
692 .position(|(name, _)| name == &field.name)
693 else {
694 return Err(invalid_record(
695 path,
696 line,
697 format!("missing payload field `{}`", field.name),
698 ));
699 };
700 let (_, raw) = values.entries.swap_remove(position);
701 decoders.decode_into(
702 &stream.name,
703 state.simulation_time().iteration(),
704 &field.name,
705 raw.get(),
706 state,
707 )?;
708 }
709 if !values.entries.is_empty() {
710 let mut extra = values
711 .entries
712 .into_iter()
713 .map(|(name, _)| name)
714 .collect::<Vec<_>>();
715 extra.sort_unstable();
716 return Err(invalid_record(
717 path,
718 line,
719 format!("undeclared payload fields: {}", extra.join(", ")),
720 ));
721 }
722 Ok(())
723}
724
725fn validate_chunk_facts(
727 metadata_path: &Path,
728 stream: &StateStreamMetadata,
729 chunk: &ChunkMetadata,
730 records: u64,
731 first_iteration: Option<u64>,
732 last_iteration: Option<u64>,
733) -> Result<(), StorageError> {
734 if records != chunk.records
735 || first_iteration != Some(chunk.first_iteration)
736 || last_iteration != Some(chunk.last_iteration)
737 {
738 return Err(StorageError::InvalidMetadata {
739 path: metadata_path.to_path_buf(),
740 reason: format!(
741 "stream `{}` chunk {} declares {} records at {}..={}, but contains {} records at {:?}..={:?}",
742 stream.name,
743 chunk.ordinal,
744 chunk.records,
745 chunk.first_iteration,
746 chunk.last_iteration,
747 records,
748 first_iteration,
749 last_iteration
750 ),
751 });
752 }
753 Ok(())
754}
755
756fn verify_checksum(
758 metadata_path: &Path,
759 path: &Path,
760 expected: &str,
761 digest: impl AsRef<[u8]>,
762) -> Result<(), StorageError> {
763 let Some(expected_digest) = expected.strip_prefix("sha256:") else {
764 return Err(StorageError::InvalidMetadata {
765 path: metadata_path.to_path_buf(),
766 reason: format!("unsupported chunk checksum algorithm in `{expected}`"),
767 });
768 };
769 let actual_digest = lowercase_hex(digest.as_ref());
770 if actual_digest != expected_digest {
771 return Err(StorageError::ChecksumMismatch {
772 path: path.to_path_buf(),
773 expected: expected.to_owned(),
774 actual: format!("sha256:{actual_digest}"),
775 });
776 }
777 Ok(())
778}
779
780fn lowercase_hex(bytes: &[u8]) -> String {
782 use std::fmt::Write as _;
783
784 let mut encoded = String::with_capacity(bytes.len() * 2);
785 for byte in bytes {
786 write!(&mut encoded, "{byte:02x}").expect("writing to String cannot fail");
787 }
788 encoded
789}
790
791fn invalid_record(path: &Path, line: u64, reason: impl Into<String>) -> StorageError {
793 StorageError::InvalidRecord {
794 path: path.to_path_buf(),
795 line,
796 reason: reason.into(),
797 }
798}