1use std::collections::{HashMap, HashSet};
52use std::fs::{self, File, OpenOptions};
53use std::io::Write;
54use std::num::NonZeroU64;
55use std::path::{Path, PathBuf};
56use std::sync::{Arc, Mutex, MutexGuard};
57use std::time::{Duration, Instant};
58
59use fs2::FileExt;
60use serde::{Deserialize, Serialize};
61use serde_json::{Map, Value};
62
63use crate::clock::{duration_nanoseconds, utc_now_rfc3339};
64use crate::configuration::TaskParameters;
65use crate::system_state::{SystemState, SystemStateSchema};
66
67mod error;
68mod json_payload_decoder;
69mod json_state_record_encoder;
70mod jsonl_format;
71mod queued_state_writer;
72mod stored_state_series_reader;
73
74pub use error::StorageError;
75pub use json_payload_decoder::{
76 JsonPayloadDecoder, JsonPayloadDecoderRegistry, JsonStringDecoder, JsonVecF64Decoder,
77};
78pub use stored_state_series_reader::StoredStateSeriesReader;
79
80use json_state_record_encoder::JsonStateRecordEncoder;
81use jsonl_format::{
82 RecordingMetadata, RecordingStatus, StateFieldMetadata, StateStreamMetadata,
83 TimeAxisMetadata as StoredTimeAxis,
84};
85use queued_state_writer::{RecoveredStateStream, StateStreamStorageConfig, StateWriterWorker};
86
87const METADATA_FILE: &str = "metadata.json";
89
90const METADATA_TEMP_FILE: &str = ".metadata.json.tmp";
92
93#[derive(Clone, Debug, Eq, PartialEq)]
100pub struct TimeAxisMetadata {
101 iteration_name: String,
102 iteration_unit: Option<String>,
103 physical_time_name: Option<String>,
104 physical_time_unit: Option<String>,
105}
106
107impl TimeAxisMetadata {
108 pub fn new(iteration_name: impl Into<String>) -> Self {
114 Self {
115 iteration_name: iteration_name.into(),
116 iteration_unit: None,
117 physical_time_name: None,
118 physical_time_unit: None,
119 }
120 }
121
122 #[must_use]
124 pub fn with_iteration_unit(mut self, unit: impl Into<String>) -> Self {
125 self.iteration_unit = Some(unit.into());
126 self
127 }
128
129 #[must_use]
131 pub fn with_physical_time_name(mut self, name: impl Into<String>) -> Self {
132 self.physical_time_name = Some(name.into());
133 self
134 }
135
136 #[must_use]
141 pub fn with_physical_time_unit(mut self, unit: impl Into<String>) -> Self {
142 self.physical_time_unit = Some(unit.into());
143 self
144 }
145
146 #[must_use]
148 pub fn with_physical_axis(mut self, name: impl Into<String>, unit: impl Into<String>) -> Self {
149 self.physical_time_name = Some(name.into());
150 self.physical_time_unit = Some(unit.into());
151 self
152 }
153
154 fn into_stored(self) -> StoredTimeAxis {
156 StoredTimeAxis {
157 iteration_name: self.iteration_name,
158 iteration_unit: self.iteration_unit,
159 physical_time_name: self.physical_time_name,
160 physical_time_unit: self.physical_time_unit,
161 }
162 }
163}
164
165impl Default for TimeAxisMetadata {
166 fn default() -> Self {
169 Self::new("iteration")
170 }
171}
172
173#[derive(Clone, Debug, Eq, PartialEq)]
179pub struct RecordingTiming {
180 created_at_utc: String,
181 finalized_at_utc: String,
182 active_duration_ns: u64,
183 continuation_count: u64,
184}
185
186impl RecordingTiming {
187 fn from_stored(
189 timing: &jsonl_format::RecordingTiming,
190 metadata_path: &Path,
191 ) -> Result<Self, StorageError> {
192 let finalized_at_utc =
193 timing
194 .finalized_at_utc
195 .clone()
196 .ok_or_else(|| StorageError::InvalidMetadata {
197 path: metadata_path.to_path_buf(),
198 reason: "completed recording lacks finalized timestamp".to_owned(),
199 })?;
200 Ok(Self {
201 created_at_utc: timing.created_at_utc.clone(),
202 finalized_at_utc,
203 active_duration_ns: timing.active_duration_ns,
204 continuation_count: timing.continuation_count,
205 })
206 }
207
208 pub fn created_at_utc(&self) -> &str {
210 &self.created_at_utc
211 }
212
213 pub fn finalized_at_utc(&self) -> &str {
215 &self.finalized_at_utc
216 }
217
218 pub fn active_duration_ns(&self) -> u64 {
220 self.active_duration_ns
221 }
222
223 pub fn active_duration(&self) -> Duration {
225 Duration::from_nanos(self.active_duration_ns)
226 }
227
228 pub fn continuation_count(&self) -> u64 {
230 self.continuation_count
231 }
232}
233
234#[derive(Clone, Debug, Eq, PartialEq)]
236pub struct CompletedStreamSummary {
237 name: String,
238 chunk_count: u64,
239 record_count: u64,
240 encoded_bytes: u64,
241 first_iteration: Option<u64>,
242 last_iteration: Option<u64>,
243}
244
245impl CompletedStreamSummary {
246 pub fn name(&self) -> &str {
248 &self.name
249 }
250
251 pub fn chunk_count(&self) -> u64 {
253 self.chunk_count
254 }
255
256 pub fn record_count(&self) -> u64 {
258 self.record_count
259 }
260
261 pub fn encoded_bytes(&self) -> u64 {
263 self.encoded_bytes
264 }
265
266 pub fn first_iteration(&self) -> Option<u64> {
268 self.first_iteration
269 }
270
271 pub fn last_iteration(&self) -> Option<u64> {
273 self.last_iteration
274 }
275}
276
277#[derive(Clone, Debug, Eq, PartialEq)]
282pub struct CompletedRecording {
283 directory: PathBuf,
284 timing: RecordingTiming,
285 terminal_metadata: Map<String, Value>,
286 streams: Vec<CompletedStreamSummary>,
287}
288
289impl CompletedRecording {
290 pub fn directory(&self) -> &Path {
292 &self.directory
293 }
294
295 pub fn timing(&self) -> &RecordingTiming {
297 &self.timing
298 }
299
300 pub fn terminal_metadata(&self) -> &Map<String, Value> {
302 &self.terminal_metadata
303 }
304
305 pub fn stream_summaries(&self) -> &[CompletedStreamSummary] {
307 &self.streams
308 }
309
310 pub fn stream_summary(&self, name: &str) -> Option<&CompletedStreamSummary> {
312 self.streams.iter().find(|stream| stream.name == name)
313 }
314}
315
316#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
327#[serde(rename_all = "snake_case")]
328pub enum SamplingInterval {
329 Iterations(NonZeroU64),
331}
332
333#[derive(Deserialize)]
334#[serde(untagged)]
335enum SamplingIntervalInput {
336 Iterations(NonZeroU64),
338 Tagged { iterations: NonZeroU64 },
340}
341
342impl<'de> Deserialize<'de> for SamplingInterval {
343 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
344 where
345 D: serde::Deserializer<'de>,
346 {
347 match SamplingIntervalInput::deserialize(deserializer)? {
348 SamplingIntervalInput::Iterations(interval)
349 | SamplingIntervalInput::Tagged {
350 iterations: interval,
351 } => Ok(Self::Iterations(interval)),
352 }
353 }
354}
355
356impl SamplingInterval {
357 pub const fn iterations(interval: u64) -> Option<Self> {
359 match NonZeroU64::new(interval) {
360 Some(interval) => Some(Self::Iterations(interval)),
361 None => None,
362 }
363 }
364
365 const fn includes(self, iteration: u64) -> bool {
367 match self {
368 Self::Iterations(interval) => iteration.is_multiple_of(interval.get()),
369 }
370 }
371}
372
373#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
382#[serde(deny_unknown_fields)]
383pub struct StateStreamConfig {
384 name: String,
385 #[serde(default, skip_serializing_if = "Option::is_none")]
386 directory: Option<String>,
387 sampling_interval: SamplingInterval,
388 fields: Vec<String>,
389 #[serde(default, skip_serializing_if = "Option::is_none")]
390 storage_limits: Option<(NonZeroU64, NonZeroU64)>,
391}
392
393impl StateStreamConfig {
394 pub fn new<I, K>(
402 name: impl Into<String>,
403 fields: I,
404 sampling_interval: SamplingInterval,
405 storage_limits: Option<(NonZeroU64, NonZeroU64)>,
406 ) -> Self
407 where
408 I: IntoIterator<Item = K>,
409 K: Into<String>,
410 {
411 let name = name.into();
412 Self {
413 directory: None,
414 name,
415 sampling_interval,
416 fields: fields.into_iter().map(Into::into).collect(),
417 storage_limits,
418 }
419 }
420
421 #[must_use]
426 pub fn with_relative_directory(mut self, directory: impl Into<String>) -> Self {
427 self.directory = Some(directory.into());
428 self
429 }
430
431 pub fn name(&self) -> &str {
433 &self.name
434 }
435
436 pub fn relative_directory(&self) -> &str {
438 self.directory.as_deref().unwrap_or(&self.name)
439 }
440
441 pub const fn sampling_interval(&self) -> SamplingInterval {
443 self.sampling_interval
444 }
445
446 pub fn fields(&self) -> &[String] {
448 &self.fields
449 }
450
451 pub const fn storage_limits(&self) -> Option<(NonZeroU64, NonZeroU64)> {
453 self.storage_limits
454 }
455}
456
457#[derive(Debug)]
464pub struct SystemStateWriterBuilder {
465 root: PathBuf,
466 spec: SystemStateSchema,
467 time: TimeAxisMetadata,
468 user_metadata: Map<String, Value>,
469 shared_stream_limits: Option<(NonZeroU64, NonZeroU64)>,
470 streams: Vec<StateStreamConfig>,
471}
472
473impl SystemStateWriterBuilder {
474 pub fn new(root: impl Into<PathBuf>, spec: &SystemStateSchema) -> Self {
479 Self {
480 root: root.into(),
481 spec: spec.clone(),
482 time: TimeAxisMetadata::default(),
483 user_metadata: Map::new(),
484 shared_stream_limits: None,
485 streams: Vec::new(),
486 }
487 }
488
489 #[must_use]
491 pub fn with_time_axis_metadata(mut self, time: TimeAxisMetadata) -> Self {
492 self.time = time;
493 self
494 }
495
496 #[must_use]
502 pub fn with_user_metadata(mut self, metadata: Map<String, Value>) -> Self {
503 self.user_metadata.extend(metadata);
504 self
505 }
506
507 #[must_use]
513 pub fn with_shared_stream_limits(
514 mut self,
515 max_chunk_bytes: NonZeroU64,
516 queue_bytes: NonZeroU64,
517 ) -> Self {
518 self.shared_stream_limits = Some((max_chunk_bytes, queue_bytes));
519 self
520 }
521
522 #[must_use]
530 pub fn with_task_parameters(mut self, parameters: &TaskParameters) -> Self {
531 self.user_metadata.extend(
532 parameters
533 .iter()
534 .map(|(key, value)| (key.to_owned(), value.clone())),
535 );
536 self.user_metadata.insert(
537 "task_ordinal".to_owned(),
538 Value::from(parameters.task_ordinal()),
539 );
540 self
541 }
542
543 #[must_use]
548 pub fn add_state_stream(mut self, stream: StateStreamConfig) -> Self {
549 self.streams.push(stream);
550 self
551 }
552
553 pub fn create_new_recording(self) -> Result<SystemStateWriter, StorageError> {
565 SystemStateWriter::create_new_recording(self)
566 }
567
568 pub fn continue_existing_recording(self) -> Result<SystemStateWriter, StorageError> {
576 SystemStateWriter::continue_recording(self, None).map(|(writer, _)| writer)
577 }
578
579 pub fn continue_recording_from_latest_checkpoint(
587 self,
588 stream: &str,
589 decoders: JsonPayloadDecoderRegistry,
590 ) -> Result<(SystemStateWriter, SystemState), StorageError> {
591 let (writer, state) =
592 SystemStateWriter::continue_recording(self, Some((stream, decoders)))?;
593 Ok((
594 writer,
595 state.expect("checkpoint-aware resume always reconstructs one state"),
596 ))
597 }
598}
599
600pub struct SystemStateWriter {
607 root: PathBuf,
608 stream_order: Vec<String>,
609 manifest: Arc<RecordingManifest>,
610 streams: HashMap<String, ScheduledStateStream>,
611 writer: Option<StateWriterWorker>,
612 session_started: Instant,
613 _lease: RecordingLease,
616}
617
618impl SystemStateWriter {
619 pub fn builder(root: impl Into<PathBuf>, spec: &SystemStateSchema) -> SystemStateWriterBuilder {
621 SystemStateWriterBuilder::new(root, spec)
622 }
623
624 pub fn recording_directory(&self) -> &Path {
626 &self.root
627 }
628
629 pub fn stream_names(&self) -> impl ExactSizeIterator<Item = &str> {
631 self.stream_order.iter().map(String::as_str)
632 }
633
634 pub fn observe_state(&mut self, state: &SystemState) -> Result<(), StorageError> {
648 let iteration = state.simulation_time().iteration();
649 let writer = self
650 .writer
651 .as_ref()
652 .expect("an active recording owns its writer worker");
653 for name in &self.stream_order {
654 let stream = self
655 .streams
656 .get_mut(name)
657 .expect("stream order contains every configured stream");
658 if !stream.sampling_interval.includes(iteration)
659 || stream.last_recorded_iteration == Some(iteration)
660 {
661 continue;
662 }
663 let record = stream.encoder.encode(state)?;
664 writer.submit_record(name, record)?;
665 stream.last_recorded_iteration = Some(iteration);
666 }
667 Ok(())
668 }
669
670 pub fn flush_stream_to_storage(&self, stream: &str) -> Result<(), StorageError> {
677 if !self.streams.contains_key(stream) {
678 return Err(StorageError::UnknownStateStream {
679 stream: stream.to_owned(),
680 });
681 }
682 self.writer
683 .as_ref()
684 .expect("an active recording owns its writer worker")
685 .flush_state_stream(stream)
686 }
687
688 pub fn complete_recording(self) -> Result<CompletedRecording, StorageError> {
696 self.complete_recording_with_terminal_metadata(Map::new())
697 }
698
699 pub fn complete_recording_with_terminal_metadata(
705 mut self,
706 terminal_metadata: Map<String, Value>,
707 ) -> Result<CompletedRecording, StorageError> {
708 if let Err(error) = self.finish_writer() {
709 let _ = self.transition_terminal(
710 RecordingStatus::Failed {
711 message: error.to_string(),
712 },
713 Map::new(),
714 );
715 return Err(error);
716 }
717 self.transition_terminal(RecordingStatus::Complete, terminal_metadata)?;
718 self.completed_recording()
719 }
720
721 pub fn complete_recording_with_final_state(
728 mut self,
729 state: &SystemState,
730 ) -> Result<CompletedRecording, StorageError> {
731 self.record_final_state(state)?;
732 self.complete_recording()
733 }
734
735 pub fn complete_recording_with_final_state_and_terminal_metadata(
738 mut self,
739 state: &SystemState,
740 terminal_metadata: Map<String, Value>,
741 ) -> Result<CompletedRecording, StorageError> {
742 self.record_final_state(state)?;
743 self.complete_recording_with_terminal_metadata(terminal_metadata)
744 }
745
746 fn record_final_state(&mut self, state: &SystemState) -> Result<(), StorageError> {
748 let iteration = state.simulation_time().iteration();
749 let writer = self
750 .writer
751 .as_ref()
752 .expect("an active recording owns its writer worker");
753 for name in &self.stream_order {
754 let stream = self
755 .streams
756 .get_mut(name)
757 .expect("stream order contains every configured stream");
758 if stream.last_recorded_iteration == Some(iteration) {
759 continue;
760 }
761 let record = stream.encoder.encode(state)?;
762 writer.submit_record(name, record)?;
763 stream.last_recorded_iteration = Some(iteration);
764 }
765 Ok(())
766 }
767
768 pub fn mark_recording_failed(self, message: impl Into<String>) -> Result<(), StorageError> {
780 self.mark_recording_failed_with_terminal_metadata(message, Map::new())
781 }
782
783 pub fn mark_recording_failed_with_terminal_metadata(
785 mut self,
786 message: impl Into<String>,
787 terminal_metadata: Map<String, Value>,
788 ) -> Result<(), StorageError> {
789 let message = message.into();
790 if message.trim().is_empty() {
791 return Err(StorageError::InvalidConfiguration {
792 setting: "failure_message",
793 reason: "failed run message must not be empty".to_owned(),
794 });
795 }
796
797 if let Err(error) = self.finish_writer() {
798 let _ = self.transition_terminal(
799 RecordingStatus::Failed {
800 message: error.to_string(),
801 },
802 Map::new(),
803 );
804 return Err(error);
805 }
806 self.transition_terminal(RecordingStatus::Failed { message }, terminal_metadata)
807 }
808
809 fn create_new_recording(builder: SystemStateWriterBuilder) -> Result<Self, StorageError> {
811 ensure_absent(&builder.root)?;
812 let prepared = PreparedRecording::from_builder(builder)?;
813 create_root(&prepared.root)?;
814 let lease = RecordingLease::acquire(&prepared.root)?;
815 for stream in &prepared.streams {
816 stream.writer.create_directory()?;
817 }
818 commit_metadata(&prepared.root, &prepared.metadata_path, &prepared.metadata)?;
819 let manifest = Arc::new(RecordingManifest::new(
820 prepared.root.clone(),
821 prepared.metadata_path.clone(),
822 prepared.metadata,
823 ));
824 Self::start_new_prepared(prepared.root, prepared.streams, manifest, lease)
825 }
826
827 fn continue_recording(
829 builder: SystemStateWriterBuilder,
830 checkpoint: Option<(&str, JsonPayloadDecoderRegistry)>,
831 ) -> Result<(Self, Option<SystemState>), StorageError> {
832 let prepared = PreparedRecording::from_builder(builder)?;
833 let lease = RecordingLease::acquire(&prepared.root)?;
834 remove_stale_metadata_temp(&prepared.root)?;
835 let mut existing = load_metadata(&prepared.metadata_path)?;
836 if !matches!(existing.status, RecordingStatus::Running) {
837 return Err(StorageError::RecordingNotContinuable {
838 path: prepared.metadata_path,
839 });
840 }
841 ensure_resume_match(&prepared.metadata_path, &prepared.metadata, &existing)?;
842
843 let mut recovered = Vec::with_capacity(prepared.streams.len());
844 for stream in prepared.streams {
845 let declaration = existing
846 .stream(&stream.name)
847 .expect("matched metadata contains every prepared stream");
848 let seed = StateWriterWorker::recover_state_stream(&stream.writer, declaration)?;
849 recovered.push((stream, seed));
850 }
851
852 let state = if let Some((checkpoint_stream, decoders)) = checkpoint {
853 let declaration = existing.stream(checkpoint_stream).ok_or_else(|| {
854 StorageError::UnknownStateStream {
855 stream: checkpoint_stream.to_owned(),
856 }
857 })?;
858 let seed = recovered
859 .iter()
860 .find(|(stream, _)| stream.name == checkpoint_stream)
861 .map(|(_, seed)| seed)
862 .expect("matched stream has one recovered seed");
863 Some(stored_state_series_reader::decode_resume_state(
864 &prepared.root,
865 &prepared.metadata_path,
866 declaration,
867 &prepared.spec,
868 &decoders,
869 seed.latest_open_record(),
870 )?)
871 } else {
872 None
873 };
874
875 existing.timing.continuation_count = existing
876 .timing
877 .continuation_count
878 .checked_add(1)
879 .ok_or_else(|| StorageError::InvalidMetadata {
880 path: prepared.metadata_path.clone(),
881 reason: "timing.continuation_count overflowed".to_owned(),
882 })?;
883 commit_metadata(&prepared.root, &prepared.metadata_path, &existing)?;
884 let manifest = Arc::new(RecordingManifest::new(
885 prepared.root.clone(),
886 prepared.metadata_path.clone(),
887 existing,
888 ));
889
890 let output = Self::start_resumed_prepared(prepared.root, recovered, manifest, lease)?;
891 Ok((output, state))
892 }
893
894 fn start_new_prepared(
896 root: PathBuf,
897 streams: Vec<PreparedStateStream>,
898 manifest: Arc<RecordingManifest>,
899 lease: RecordingLease,
900 ) -> Result<Self, StorageError> {
901 let mut scheduled = HashMap::with_capacity(streams.len());
902 let mut configs = Vec::with_capacity(streams.len());
903 let mut stream_order = Vec::with_capacity(streams.len());
904 for prepared in streams {
905 let name = prepared.name;
906 stream_order.push(name.clone());
907 scheduled.insert(
908 name,
909 ScheduledStateStream {
910 encoder: prepared.encoder,
911 sampling_interval: prepared.sampling_interval,
912 last_recorded_iteration: None,
913 },
914 );
915 configs.push(prepared.writer);
916 }
917 let writer = StateWriterWorker::start_new_recording(configs, Arc::clone(&manifest))?;
918 Ok(Self {
919 root,
920 stream_order,
921 manifest,
922 streams: scheduled,
923 writer: Some(writer),
924 session_started: Instant::now(),
925 _lease: lease,
926 })
927 }
928
929 fn start_resumed_prepared(
931 root: PathBuf,
932 streams: Vec<(PreparedStateStream, RecoveredStateStream)>,
933 manifest: Arc<RecordingManifest>,
934 lease: RecordingLease,
935 ) -> Result<Self, StorageError> {
936 let mut scheduled = HashMap::with_capacity(streams.len());
937 let mut recovered_streams = Vec::with_capacity(streams.len());
938 let mut stream_order = Vec::with_capacity(streams.len());
939 for (prepared, seed) in streams {
940 let name = prepared.name;
941 stream_order.push(name.clone());
942 scheduled.insert(
943 name,
944 ScheduledStateStream {
945 encoder: prepared.encoder,
946 sampling_interval: prepared.sampling_interval,
947 last_recorded_iteration: seed.last_iteration(),
948 },
949 );
950 recovered_streams.push((prepared.writer, seed));
951 }
952 let writer = StateWriterWorker::continue_recovered_recording(
953 recovered_streams,
954 Arc::clone(&manifest),
955 )?;
956 Ok(Self {
957 root,
958 stream_order,
959 manifest,
960 streams: scheduled,
961 writer: Some(writer),
962 session_started: Instant::now(),
963 _lease: lease,
964 })
965 }
966
967 fn finish_writer(&mut self) -> Result<(), StorageError> {
969 let Some(writer) = self.writer.take() else {
970 return Ok(());
971 };
972 writer.finish_recording()
973 }
974
975 fn transition_terminal(
977 &self,
978 status: RecordingStatus,
979 terminal_metadata: Map<String, Value>,
980 ) -> Result<(), StorageError> {
981 let finalized_at_utc =
982 utc_now_rfc3339().map_err(|source| StorageError::OperationalTimestamp {
983 operation: "finalize recording",
984 source,
985 })?;
986 let active_duration_ns = duration_nanoseconds(self.session_started.elapsed())
987 .ok_or(StorageError::OperationalDurationOverflow)?;
988 self.manifest.transition_terminal(
989 status,
990 finalized_at_utc,
991 active_duration_ns,
992 terminal_metadata,
993 )
994 }
995
996 fn completed_recording(&self) -> Result<CompletedRecording, StorageError> {
998 let metadata = self.manifest.snapshot();
999 let timing = RecordingTiming::from_stored(&metadata.timing, &self.manifest.path)?;
1000 let streams = metadata
1001 .streams
1002 .iter()
1003 .map(completed_stream_summary)
1004 .collect::<Result<Vec<_>, _>>()?;
1005 Ok(CompletedRecording {
1006 directory: self.root.clone(),
1007 timing,
1008 terminal_metadata: metadata.terminal_metadata,
1009 streams,
1010 })
1011 }
1012}
1013
1014fn completed_stream_summary(
1016 stream: &StateStreamMetadata,
1017) -> Result<CompletedStreamSummary, StorageError> {
1018 let overflow = || StorageError::ByteCountOverflow {
1019 stream: stream.name.clone(),
1020 };
1021 let chunk_count = u64::try_from(stream.chunks.len()).map_err(|_| overflow())?;
1022 let record_count = stream
1023 .chunks
1024 .iter()
1025 .try_fold(0_u64, |total, chunk| total.checked_add(chunk.records))
1026 .ok_or_else(&overflow)?;
1027 let encoded_bytes = stream
1028 .chunks
1029 .iter()
1030 .try_fold(0_u64, |total, chunk| total.checked_add(chunk.bytes))
1031 .ok_or_else(overflow)?;
1032 Ok(CompletedStreamSummary {
1033 name: stream.name.clone(),
1034 chunk_count,
1035 record_count,
1036 encoded_bytes,
1037 first_iteration: stream.chunks.first().map(|chunk| chunk.first_iteration),
1038 last_iteration: stream.chunks.last().map(|chunk| chunk.last_iteration),
1039 })
1040}
1041
1042struct PreparedRecording {
1044 root: PathBuf,
1045 metadata_path: PathBuf,
1046 spec: SystemStateSchema,
1047 metadata: RecordingMetadata,
1048 streams: Vec<PreparedStateStream>,
1049}
1050
1051impl PreparedRecording {
1052 fn from_builder(builder: SystemStateWriterBuilder) -> Result<Self, StorageError> {
1054 let metadata_path = builder.root.join(METADATA_FILE);
1055 let stored_time = builder.time.into_stored();
1056 let mut names = HashSet::with_capacity(builder.streams.len());
1057 let mut directories = HashSet::with_capacity(builder.streams.len());
1058 let mut streams = Vec::with_capacity(builder.streams.len());
1059 let mut declarations = Vec::with_capacity(builder.streams.len());
1060
1061 for config in builder.streams {
1062 if !names.insert(config.name.clone()) {
1063 return Err(StorageError::DuplicateStateStream {
1064 stream: config.name,
1065 });
1066 }
1067 let directory = config.relative_directory().to_owned();
1068 if !directories.insert(directory.clone()) {
1069 return Err(StorageError::InvalidConfiguration {
1070 setting: "stream.directory",
1071 reason: format!("multiple streams use relative directory `{}`", directory),
1072 });
1073 }
1074
1075 let (max_chunk_bytes, queue_bytes) = config
1076 .storage_limits
1077 .or(builder.shared_stream_limits)
1078 .ok_or_else(|| StorageError::InvalidConfiguration {
1079 setting: "stream.storage_limits",
1080 reason: format!(
1081 "stream `{}` has no explicit limits and the writer has no shared limits",
1082 config.name
1083 ),
1084 })?;
1085 let encoder = JsonStateRecordEncoder::new(&config.name, &builder.spec, &config.fields)?;
1086 let fields = encoder
1087 .fields()
1088 .map(|name| {
1089 let field = builder
1090 .spec
1091 .field_schema(name)
1092 .expect("encoder fields were validated against this specification");
1093 StateFieldMetadata {
1094 name: name.to_owned(),
1095 description: field.description().map(str::to_owned),
1096 }
1097 })
1098 .collect::<Vec<_>>();
1099 declarations.push(StateStreamMetadata {
1100 name: config.name.clone(),
1101 directory: directory.clone(),
1102 sampling_interval: config.sampling_interval,
1103 fields,
1104 max_chunk_bytes: max_chunk_bytes.get(),
1105 queue_bytes: queue_bytes.get(),
1106 chunks: Vec::new(),
1107 });
1108 streams.push(PreparedStateStream {
1109 name: config.name.clone(),
1110 encoder,
1111 sampling_interval: config.sampling_interval,
1112 writer: StateStreamStorageConfig::new(
1113 &config.name,
1114 builder.root.join(&directory),
1115 max_chunk_bytes,
1116 queue_bytes,
1117 )?,
1118 });
1119 }
1120
1121 let created_at_utc =
1122 utc_now_rfc3339().map_err(|source| StorageError::OperationalTimestamp {
1123 operation: "create recording",
1124 source,
1125 })?;
1126 let metadata = RecordingMetadata::running(
1127 stored_time,
1128 builder.user_metadata,
1129 declarations,
1130 created_at_utc,
1131 );
1132 metadata.validate(&metadata_path)?;
1133 Ok(Self {
1134 root: builder.root,
1135 metadata_path,
1136 spec: builder.spec,
1137 metadata,
1138 streams,
1139 })
1140 }
1141}
1142
1143struct PreparedStateStream {
1145 name: String,
1146 encoder: JsonStateRecordEncoder,
1147 sampling_interval: SamplingInterval,
1148 writer: StateStreamStorageConfig,
1149}
1150
1151struct ScheduledStateStream {
1153 encoder: JsonStateRecordEncoder,
1154 sampling_interval: SamplingInterval,
1155 last_recorded_iteration: Option<u64>,
1156}
1157
1158pub(crate) struct RecordingManifest {
1164 root: PathBuf,
1165 path: PathBuf,
1166 metadata: Mutex<RecordingMetadata>,
1167}
1168
1169impl RecordingManifest {
1170 fn new(root: PathBuf, path: PathBuf, metadata: RecordingMetadata) -> Self {
1172 Self {
1173 root,
1174 path,
1175 metadata: Mutex::new(metadata),
1176 }
1177 }
1178
1179 pub(crate) fn prepare_chunk(
1181 &self,
1182 stream: &str,
1183 descriptor: jsonl_format::ChunkMetadata,
1184 ) -> Result<(), StorageError> {
1185 let mut current = lock_metadata(&self.metadata);
1186 if !matches!(current.status, RecordingStatus::Running) {
1187 return Err(StorageError::RecordingFinished);
1188 }
1189 let mut candidate = current.clone();
1190 let declaration =
1191 candidate
1192 .stream_mut(stream)
1193 .ok_or_else(|| StorageError::UnknownStateStream {
1194 stream: stream.to_owned(),
1195 })?;
1196 let expected = u64::try_from(declaration.chunks.len()).map_err(|_| {
1197 StorageError::ByteCountOverflow {
1198 stream: stream.to_owned(),
1199 }
1200 })?;
1201 if descriptor.ordinal != expected {
1202 return Err(StorageError::InvalidMetadata {
1203 path: self.path.clone(),
1204 reason: format!(
1205 "stream `{stream}` prepared chunk ordinal {}, expected {expected}",
1206 descriptor.ordinal
1207 ),
1208 });
1209 }
1210 declaration.chunks.push(descriptor);
1211 commit_metadata(&self.root, &self.path, &candidate)?;
1212 *current = candidate;
1213 Ok(())
1214 }
1215
1216 fn transition_terminal(
1218 &self,
1219 status: RecordingStatus,
1220 finalized_at_utc: String,
1221 active_duration_ns: u64,
1222 terminal_metadata: Map<String, Value>,
1223 ) -> Result<(), StorageError> {
1224 let mut current = lock_metadata(&self.metadata);
1225 let mut candidate = current.clone();
1226 candidate.status = status;
1227 candidate.timing.finalized_at_utc = Some(finalized_at_utc);
1228 candidate.timing.active_duration_ns = candidate
1229 .timing
1230 .active_duration_ns
1231 .checked_add(active_duration_ns)
1232 .ok_or(StorageError::OperationalDurationOverflow)?;
1233 candidate.terminal_metadata = terminal_metadata;
1234 commit_metadata(&self.root, &self.path, &candidate)?;
1235 *current = candidate;
1236 Ok(())
1237 }
1238
1239 fn snapshot(&self) -> RecordingMetadata {
1241 lock_metadata(&self.metadata).clone()
1242 }
1243}
1244
1245struct RecordingLease {
1250 _directory: File,
1251}
1252
1253impl RecordingLease {
1254 fn acquire(root: &Path) -> Result<Self, StorageError> {
1256 let directory = File::open(root).map_err(|source| StorageError::Io {
1257 operation: "open output root for exclusive ownership",
1258 path: root.to_path_buf(),
1259 source,
1260 })?;
1261 match FileExt::try_lock_exclusive(&directory) {
1262 Ok(()) => Ok(Self {
1263 _directory: directory,
1264 }),
1265 Err(source) if source.kind() == std::io::ErrorKind::WouldBlock => {
1266 Err(StorageError::RecordingDirectoryInUse {
1267 path: root.to_path_buf(),
1268 })
1269 }
1270 Err(source) => Err(StorageError::Io {
1271 operation: "acquire exclusive output ownership",
1272 path: root.to_path_buf(),
1273 source,
1274 }),
1275 }
1276 }
1277}
1278
1279fn load_metadata(path: &Path) -> Result<RecordingMetadata, StorageError> {
1281 let bytes = fs::read(path).map_err(|source| StorageError::Io {
1282 operation: "read metadata for resume",
1283 path: path.to_path_buf(),
1284 source,
1285 })?;
1286 let metadata: RecordingMetadata =
1287 serde_json::from_slice(&bytes).map_err(|source| StorageError::Json {
1288 operation: "parse metadata for resume",
1289 path: path.to_path_buf(),
1290 source,
1291 })?;
1292 metadata.validate(path)?;
1293 Ok(metadata)
1294}
1295
1296fn ensure_resume_match(
1298 path: &Path,
1299 expected: &RecordingMetadata,
1300 existing: &RecordingMetadata,
1301) -> Result<(), StorageError> {
1302 let mut configuration = existing.clone();
1303 for stream in &mut configuration.streams {
1304 stream.chunks.clear();
1305 }
1306 configuration.status = RecordingStatus::Running;
1307 configuration.timing = expected.timing.clone();
1308 configuration.terminal_metadata.clear();
1309 if &configuration != expected {
1310 return Err(StorageError::RecordingConfigurationMismatch {
1311 path: path.to_path_buf(),
1312 reason: "builder time axis, user metadata, or stream declarations differ".to_owned(),
1313 });
1314 }
1315 Ok(())
1316}
1317
1318fn remove_stale_metadata_temp(root: &Path) -> Result<(), StorageError> {
1320 let path = root.join(METADATA_TEMP_FILE);
1321 match fs::remove_file(&path) {
1322 Ok(()) => File::open(root)
1323 .and_then(|directory| directory.sync_all())
1324 .map_err(|source| StorageError::Io {
1325 operation: "synchronize stale metadata cleanup",
1326 path: root.to_path_buf(),
1327 source,
1328 }),
1329 Err(source) if source.kind() == std::io::ErrorKind::NotFound => Ok(()),
1330 Err(source) => Err(StorageError::Io {
1331 operation: "remove stale temporary metadata",
1332 path,
1333 source,
1334 }),
1335 }
1336}
1337
1338fn lock_metadata(metadata: &Mutex<RecordingMetadata>) -> MutexGuard<'_, RecordingMetadata> {
1340 metadata
1341 .lock()
1342 .unwrap_or_else(|poisoned| poisoned.into_inner())
1343}
1344
1345fn ensure_absent(root: &Path) -> Result<(), StorageError> {
1347 match root.try_exists() {
1348 Ok(false) => Ok(()),
1349 Ok(true) => Err(StorageError::RecordingDirectoryExists {
1350 path: root.to_path_buf(),
1351 }),
1352 Err(source) => Err(StorageError::Io {
1353 operation: "inspect output root",
1354 path: root.to_path_buf(),
1355 source,
1356 }),
1357 }
1358}
1359
1360fn create_root(root: &Path) -> Result<(), StorageError> {
1362 match fs::create_dir(root) {
1363 Ok(()) => Ok(()),
1364 Err(source) if source.kind() == std::io::ErrorKind::AlreadyExists => {
1365 Err(StorageError::RecordingDirectoryExists {
1366 path: root.to_path_buf(),
1367 })
1368 }
1369 Err(source) => Err(StorageError::Io {
1370 operation: "create output root",
1371 path: root.to_path_buf(),
1372 source,
1373 }),
1374 }
1375}
1376
1377fn commit_metadata(
1385 root: &Path,
1386 metadata_path: &Path,
1387 metadata: &RecordingMetadata,
1388) -> Result<(), StorageError> {
1389 metadata.validate(metadata_path)?;
1390 let mut bytes = serde_json::to_vec_pretty(metadata).map_err(|source| StorageError::Json {
1391 operation: "serialize metadata",
1392 path: metadata_path.to_path_buf(),
1393 source,
1394 })?;
1395 bytes.push(b'\n');
1396
1397 let temporary_path = root.join(METADATA_TEMP_FILE);
1398 let result = write_and_replace_metadata(root, metadata_path, &temporary_path, &bytes);
1399 if result.is_err() {
1400 let _ = fs::remove_file(&temporary_path);
1401 }
1402 result
1403}
1404
1405fn write_and_replace_metadata(
1407 root: &Path,
1408 metadata_path: &Path,
1409 temporary_path: &Path,
1410 bytes: &[u8],
1411) -> Result<(), StorageError> {
1412 let mut temporary = OpenOptions::new()
1413 .write(true)
1414 .create_new(true)
1415 .open(temporary_path)
1416 .map_err(|source| StorageError::Io {
1417 operation: "create temporary metadata",
1418 path: temporary_path.to_path_buf(),
1419 source,
1420 })?;
1421 temporary
1422 .write_all(bytes)
1423 .map_err(|source| StorageError::Io {
1424 operation: "write temporary metadata",
1425 path: temporary_path.to_path_buf(),
1426 source,
1427 })?;
1428 temporary.sync_all().map_err(|source| StorageError::Io {
1429 operation: "sync temporary metadata",
1430 path: temporary_path.to_path_buf(),
1431 source,
1432 })?;
1433 drop(temporary);
1434
1435 fs::rename(temporary_path, metadata_path).map_err(|source| StorageError::Io {
1436 operation: "publish metadata",
1437 path: metadata_path.to_path_buf(),
1438 source,
1439 })?;
1440
1441 File::open(root)
1442 .and_then(|directory| directory.sync_all())
1443 .map_err(|source| StorageError::Io {
1444 operation: "sync output root",
1445 path: root.to_path_buf(),
1446 source,
1447 })
1448}