1use std::collections::{HashMap, HashSet};
65use std::fs::{self, File, OpenOptions};
66use std::io::Write;
67use std::num::NonZeroU64;
68use std::path::{Path, PathBuf};
69use std::sync::{Arc, Mutex, MutexGuard};
70use std::time::{Duration, Instant};
71
72use fs2::FileExt;
73use serde::{Deserialize, Serialize};
74use serde_json::{Map, Value};
75
76use crate::clock::{duration_nanoseconds, utc_now_rfc3339};
77use crate::configuration::ResolvedConfiguration;
78use crate::system_state::{StateSchemaSource, SystemState, SystemStateSchema};
79
80mod error;
81mod json_payload_decoder;
82mod json_state_record_encoder;
83mod jsonl_format;
84mod queued_state_writer;
85mod resume;
86mod stored_state_series_reader;
87
88pub use error::StorageError;
89pub use json_payload_decoder::{
90 JsonPayloadDecoder, JsonPayloadDecoderRegistry, JsonStringDecoder, JsonVecF64Decoder,
91};
92pub use stored_state_series_reader::StoredStateSeriesReader;
93
94use json_state_record_encoder::JsonStateRecordEncoder;
95use jsonl_format::{
96 RecordingMetadata, RecordingStatus, StateFieldMetadata, StateStreamMetadata,
97 TimeAxisMetadata as StoredTimeAxis,
98};
99use queued_state_writer::{RecoveredStateStream, StateStreamStorageConfig, StateWriterWorker};
100
101const METADATA_FILE: &str = "metadata.json";
103
104const METADATA_TEMP_FILE: &str = ".metadata.json.tmp";
106
107#[derive(Clone, Debug, Eq, PartialEq)]
114pub struct TimeAxisMetadata {
115 iteration_name: String,
116 iteration_unit: Option<String>,
117 physical_time_name: Option<String>,
118 physical_time_unit: Option<String>,
119}
120
121impl TimeAxisMetadata {
122 pub fn new(iteration_name: impl Into<String>) -> Self {
128 Self {
129 iteration_name: iteration_name.into(),
130 iteration_unit: None,
131 physical_time_name: None,
132 physical_time_unit: None,
133 }
134 }
135
136 #[must_use]
138 pub fn with_iteration_unit(mut self, unit: impl Into<String>) -> Self {
139 self.iteration_unit = Some(unit.into());
140 self
141 }
142
143 #[must_use]
145 pub fn with_physical_time_name(mut self, name: impl Into<String>) -> Self {
146 self.physical_time_name = Some(name.into());
147 self
148 }
149
150 #[must_use]
155 pub fn with_physical_time_unit(mut self, unit: impl Into<String>) -> Self {
156 self.physical_time_unit = Some(unit.into());
157 self
158 }
159
160 #[must_use]
162 pub fn with_physical_axis(mut self, name: impl Into<String>, unit: impl Into<String>) -> Self {
163 self.physical_time_name = Some(name.into());
164 self.physical_time_unit = Some(unit.into());
165 self
166 }
167
168 fn into_stored(self) -> StoredTimeAxis {
170 StoredTimeAxis {
171 iteration_name: self.iteration_name,
172 iteration_unit: self.iteration_unit,
173 physical_time_name: self.physical_time_name,
174 physical_time_unit: self.physical_time_unit,
175 }
176 }
177}
178
179impl Default for TimeAxisMetadata {
180 fn default() -> Self {
183 Self::new("iteration")
184 }
185}
186
187#[derive(Clone, Debug, Eq, PartialEq)]
193pub struct RecordingTiming {
194 created_at_utc: String,
195 finalized_at_utc: String,
196 active_duration_ns: u64,
197 continuation_count: u64,
198}
199
200impl RecordingTiming {
201 fn from_stored(
203 timing: &jsonl_format::RecordingTiming,
204 metadata_path: &Path,
205 ) -> Result<Self, StorageError> {
206 let finalized_at_utc =
207 timing
208 .finalized_at_utc
209 .clone()
210 .ok_or_else(|| StorageError::InvalidMetadata {
211 path: metadata_path.to_path_buf(),
212 reason: "completed recording lacks finalized timestamp".to_owned(),
213 })?;
214 Ok(Self {
215 created_at_utc: timing.created_at_utc.clone(),
216 finalized_at_utc,
217 active_duration_ns: timing.active_duration_ns,
218 continuation_count: timing.continuation_count,
219 })
220 }
221
222 pub fn created_at_utc(&self) -> &str {
224 &self.created_at_utc
225 }
226
227 pub fn finalized_at_utc(&self) -> &str {
229 &self.finalized_at_utc
230 }
231
232 pub fn active_duration_ns(&self) -> u64 {
234 self.active_duration_ns
235 }
236
237 pub fn active_duration(&self) -> Duration {
239 Duration::from_nanos(self.active_duration_ns)
240 }
241
242 pub fn continuation_count(&self) -> u64 {
244 self.continuation_count
245 }
246}
247
248#[derive(Clone, Debug, Eq, PartialEq)]
250pub struct CompletedStreamSummary {
251 name: String,
252 chunk_count: u64,
253 record_count: u64,
254 encoded_bytes: u64,
255 first_iteration: Option<u64>,
256 last_iteration: Option<u64>,
257}
258
259impl CompletedStreamSummary {
260 pub fn name(&self) -> &str {
262 &self.name
263 }
264
265 pub fn chunk_count(&self) -> u64 {
267 self.chunk_count
268 }
269
270 pub fn record_count(&self) -> u64 {
272 self.record_count
273 }
274
275 pub fn encoded_bytes(&self) -> u64 {
277 self.encoded_bytes
278 }
279
280 pub fn first_iteration(&self) -> Option<u64> {
282 self.first_iteration
283 }
284
285 pub fn last_iteration(&self) -> Option<u64> {
287 self.last_iteration
288 }
289}
290
291#[derive(Clone, Debug, Eq, PartialEq)]
296pub struct CompletedRecording {
297 directory: PathBuf,
298 timing: RecordingTiming,
299 terminal_metadata: Map<String, Value>,
300 streams: Vec<CompletedStreamSummary>,
301}
302
303impl CompletedRecording {
304 pub fn directory(&self) -> &Path {
306 &self.directory
307 }
308
309 pub fn timing(&self) -> &RecordingTiming {
311 &self.timing
312 }
313
314 pub fn terminal_metadata(&self) -> &Map<String, Value> {
316 &self.terminal_metadata
317 }
318
319 pub fn stream_summaries(&self) -> &[CompletedStreamSummary] {
321 &self.streams
322 }
323
324 pub fn stream_summary(&self, name: &str) -> Option<&CompletedStreamSummary> {
326 self.streams.iter().find(|stream| stream.name == name)
327 }
328}
329
330#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
341#[serde(rename_all = "snake_case")]
342pub enum SamplingInterval {
343 Iterations(NonZeroU64),
345}
346
347#[derive(Deserialize)]
348#[serde(untagged)]
349enum SamplingIntervalInput {
350 Iterations(NonZeroU64),
352 Tagged { iterations: NonZeroU64 },
354}
355
356impl<'de> Deserialize<'de> for SamplingInterval {
357 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
358 where
359 D: serde::Deserializer<'de>,
360 {
361 match SamplingIntervalInput::deserialize(deserializer)? {
362 SamplingIntervalInput::Iterations(interval)
363 | SamplingIntervalInput::Tagged {
364 iterations: interval,
365 } => Ok(Self::Iterations(interval)),
366 }
367 }
368}
369
370impl SamplingInterval {
371 pub const fn iterations(interval: u64) -> Option<Self> {
373 match NonZeroU64::new(interval) {
374 Some(interval) => Some(Self::Iterations(interval)),
375 None => None,
376 }
377 }
378
379 const fn includes(self, iteration: u64) -> bool {
381 match self {
382 Self::Iterations(interval) => iteration.is_multiple_of(interval.get()),
383 }
384 }
385}
386
387#[derive(Clone, Copy, Debug, Eq, PartialEq, Deserialize, Serialize)]
389#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
390pub enum StateStreamLayout {
391 Chunked { target_bytes: NonZeroU64 },
393 IndividualFiles,
395}
396
397#[derive(Clone, Copy, Debug, Eq, PartialEq, Deserialize, Serialize)]
399#[serde(deny_unknown_fields)]
400pub struct StateStreamStorage {
401 layout: StateStreamLayout,
402 storage_queue_bytes: NonZeroU64,
403}
404
405impl StateStreamStorage {
406 pub const fn chunked(target_bytes: NonZeroU64, storage_queue_bytes: NonZeroU64) -> Self {
408 Self {
409 layout: StateStreamLayout::Chunked { target_bytes },
410 storage_queue_bytes,
411 }
412 }
413
414 pub const fn individual_files(storage_queue_bytes: NonZeroU64) -> Self {
416 Self {
417 layout: StateStreamLayout::IndividualFiles,
418 storage_queue_bytes,
419 }
420 }
421
422 pub const fn layout(self) -> StateStreamLayout {
423 self.layout
424 }
425
426 pub const fn storage_queue_bytes(self) -> NonZeroU64 {
428 self.storage_queue_bytes
429 }
430}
431
432#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
439#[serde(deny_unknown_fields)]
440pub struct StateStreamConfig {
441 name: String,
442 #[serde(default, skip_serializing_if = "Option::is_none")]
443 directory: Option<String>,
444 sampling_interval: SamplingInterval,
445 fields: Vec<String>,
446 #[serde(default, skip_serializing_if = "Option::is_none")]
447 storage: Option<StateStreamStorage>,
448}
449
450impl StateStreamConfig {
451 pub fn new<I, K>(
459 name: impl Into<String>,
460 fields: I,
461 sampling_interval: SamplingInterval,
462 storage: Option<StateStreamStorage>,
463 ) -> Self
464 where
465 I: IntoIterator<Item = K>,
466 K: Into<String>,
467 {
468 let name = name.into();
469 Self {
470 directory: None,
471 name,
472 sampling_interval,
473 fields: fields.into_iter().map(Into::into).collect(),
474 storage,
475 }
476 }
477
478 #[must_use]
483 pub fn with_relative_directory(mut self, directory: impl Into<String>) -> Self {
484 self.directory = Some(directory.into());
485 self
486 }
487
488 pub fn name(&self) -> &str {
490 &self.name
491 }
492
493 pub fn relative_directory(&self) -> &str {
495 self.directory.as_deref().unwrap_or(&self.name)
496 }
497
498 pub const fn sampling_interval(&self) -> SamplingInterval {
500 self.sampling_interval
501 }
502
503 pub fn fields(&self) -> &[String] {
505 &self.fields
506 }
507
508 pub const fn storage(&self) -> Option<StateStreamStorage> {
510 self.storage
511 }
512}
513
514#[derive(Debug)]
521pub struct SystemStateWriterBuilder {
522 root: PathBuf,
523 spec: SystemStateSchema,
524 time: TimeAxisMetadata,
525 user_metadata: Map<String, Value>,
526 shared_stream_storage: Option<StateStreamStorage>,
527 streams: Vec<StateStreamConfig>,
528}
529
530impl SystemStateWriterBuilder {
531 pub fn new<S>(root: impl Into<PathBuf>, source: &S) -> Self
537 where
538 S: StateSchemaSource + ?Sized,
539 {
540 Self {
541 root: root.into(),
542 spec: source.state_schema().clone(),
543 time: TimeAxisMetadata::default(),
544 user_metadata: Map::new(),
545 shared_stream_storage: None,
546 streams: Vec::new(),
547 }
548 }
549
550 #[must_use]
552 pub fn with_time_axis_metadata(mut self, time: TimeAxisMetadata) -> Self {
553 self.time = time;
554 self
555 }
556
557 #[must_use]
563 pub fn with_user_metadata(mut self, metadata: Map<String, Value>) -> Self {
564 self.user_metadata.extend(metadata);
565 self
566 }
567
568 #[must_use]
574 pub fn with_shared_stream_storage(mut self, storage: StateStreamStorage) -> Self {
575 self.shared_stream_storage = Some(storage);
576 self
577 }
578
579 #[must_use]
587 pub fn with_configuration(mut self, configuration: &ResolvedConfiguration) -> Self {
588 self.user_metadata
589 .extend(configuration.resolved_object().clone());
590 self.user_metadata
591 .insert("ordinal".to_owned(), Value::from(configuration.ordinal()));
592 self
593 }
594
595 #[must_use]
600 pub fn add_state_stream(mut self, stream: StateStreamConfig) -> Self {
601 self.streams.push(stream);
602 self
603 }
604
605 pub fn create_new_recording(self) -> Result<SystemStateWriter, StorageError> {
617 SystemStateWriter::create_new_recording(self)
618 }
619
620 pub fn open_or_resume_from_latest_checkpoint(
628 self,
629 decoders: JsonPayloadDecoderRegistry,
630 ) -> Result<(SystemStateWriter, Option<SystemState>), StorageError> {
631 match self.root.try_exists() {
632 Ok(false) => Self::create_new_recording(self).map(|writer| (writer, None)),
633 Ok(true) => SystemStateWriter::continue_recording(
634 self,
635 Some(CheckpointRequest::LatestComplete(decoders)),
636 ),
637 Err(source) => Err(StorageError::Io {
638 operation: "inspect recording root for automatic resume",
639 path: self.root.clone(),
640 source,
641 }),
642 }
643 }
644
645 pub fn continue_existing_recording(self) -> Result<SystemStateWriter, StorageError> {
653 SystemStateWriter::continue_recording(self, None).map(|(writer, _)| writer)
654 }
655
656 pub fn continue_recording_from_latest_checkpoint(
665 self,
666 stream: &str,
667 decoders: JsonPayloadDecoderRegistry,
668 ) -> Result<(SystemStateWriter, SystemState), StorageError> {
669 let (writer, state) = SystemStateWriter::continue_recording(
670 self,
671 Some(CheckpointRequest::Named(stream.to_owned(), decoders)),
672 )?;
673 Ok((
674 writer,
675 state.expect("checkpoint-aware resume always reconstructs one state"),
676 ))
677 }
678}
679
680enum CheckpointRequest {
681 Named(String, JsonPayloadDecoderRegistry),
682 LatestComplete(JsonPayloadDecoderRegistry),
683}
684
685pub struct SystemStateWriter {
692 root: PathBuf,
693 stream_order: Vec<String>,
694 manifest: Arc<RecordingManifest>,
695 streams: HashMap<String, ScheduledStateStream>,
696 writer: Option<StateWriterWorker>,
697 session_started: Instant,
698 _lease: RecordingLease,
701}
702
703impl SystemStateWriter {
704 pub fn builder<S>(root: impl Into<PathBuf>, source: &S) -> SystemStateWriterBuilder
706 where
707 S: StateSchemaSource + ?Sized,
708 {
709 SystemStateWriterBuilder::new(root, source)
710 }
711
712 pub fn recording_directory(&self) -> &Path {
714 &self.root
715 }
716
717 pub fn stream_names(&self) -> impl ExactSizeIterator<Item = &str> {
719 self.stream_order.iter().map(String::as_str)
720 }
721
722 pub fn observe_state(&mut self, state: &SystemState) -> Result<(), StorageError> {
736 let iteration = state.simulation_time().iteration();
737 let writer = self
738 .writer
739 .as_ref()
740 .expect("an active recording owns its writer worker");
741 for name in &self.stream_order {
742 let stream = self
743 .streams
744 .get_mut(name)
745 .expect("stream order contains every configured stream");
746 if !stream.sampling_interval.includes(iteration)
747 || stream.last_recorded_iteration == Some(iteration)
748 {
749 continue;
750 }
751 let record = stream.encoder.encode(state)?;
752 writer.submit_record(name, record)?;
753 stream.last_recorded_iteration = Some(iteration);
754 }
755 Ok(())
756 }
757
758 pub fn flush_stream_to_storage(&self, stream: &str) -> Result<(), StorageError> {
765 if !self.streams.contains_key(stream) {
766 return Err(StorageError::UnknownStateStream {
767 stream: stream.to_owned(),
768 });
769 }
770 self.writer
771 .as_ref()
772 .expect("an active recording owns its writer worker")
773 .flush_state_stream(stream)
774 }
775
776 pub fn complete_recording(self) -> Result<CompletedRecording, StorageError> {
784 self.complete_recording_with_terminal_metadata(Map::new())
785 }
786
787 pub fn complete_recording_with_terminal_metadata(
793 mut self,
794 terminal_metadata: Map<String, Value>,
795 ) -> Result<CompletedRecording, StorageError> {
796 if let Err(error) = self.finish_writer() {
797 let _ = self.transition_terminal(
798 RecordingStatus::Failed {
799 message: error.to_string(),
800 },
801 Map::new(),
802 );
803 return Err(error);
804 }
805 self.transition_terminal(RecordingStatus::Complete, terminal_metadata)?;
806 self.completed_recording()
807 }
808
809 pub fn complete_recording_with_final_state(
816 mut self,
817 state: &SystemState,
818 ) -> Result<CompletedRecording, StorageError> {
819 self.record_final_state(state)?;
820 self.complete_recording()
821 }
822
823 pub fn complete_recording_with_final_state_and_terminal_metadata(
826 mut self,
827 state: &SystemState,
828 terminal_metadata: Map<String, Value>,
829 ) -> Result<CompletedRecording, StorageError> {
830 self.record_final_state(state)?;
831 self.complete_recording_with_terminal_metadata(terminal_metadata)
832 }
833
834 fn record_final_state(&mut self, state: &SystemState) -> Result<(), StorageError> {
836 let iteration = state.simulation_time().iteration();
837 let writer = self
838 .writer
839 .as_ref()
840 .expect("an active recording owns its writer worker");
841 for name in &self.stream_order {
842 let stream = self
843 .streams
844 .get_mut(name)
845 .expect("stream order contains every configured stream");
846 if stream.last_recorded_iteration == Some(iteration) {
847 continue;
848 }
849 let record = stream.encoder.encode(state)?;
850 writer.submit_record(name, record)?;
851 stream.last_recorded_iteration = Some(iteration);
852 }
853 Ok(())
854 }
855
856 pub fn mark_recording_failed(self, message: impl Into<String>) -> Result<(), StorageError> {
868 self.mark_recording_failed_with_terminal_metadata(message, Map::new())
869 }
870
871 pub fn mark_recording_failed_with_terminal_metadata(
873 mut self,
874 message: impl Into<String>,
875 terminal_metadata: Map<String, Value>,
876 ) -> Result<(), StorageError> {
877 let message = message.into();
878 if message.trim().is_empty() {
879 return Err(StorageError::InvalidConfiguration {
880 setting: "failure_message",
881 reason: "failed run message must not be empty".to_owned(),
882 });
883 }
884
885 if let Err(error) = self.finish_writer() {
886 let _ = self.transition_terminal(
887 RecordingStatus::Failed {
888 message: error.to_string(),
889 },
890 Map::new(),
891 );
892 return Err(error);
893 }
894 self.transition_terminal(RecordingStatus::Failed { message }, terminal_metadata)
895 }
896
897 fn create_new_recording(builder: SystemStateWriterBuilder) -> Result<Self, StorageError> {
899 ensure_absent(&builder.root)?;
900 let prepared = PreparedRecording::from_builder(builder)?;
901 create_root(&prepared.root)?;
902 let lease = RecordingLease::acquire(&prepared.root)?;
903 for stream in &prepared.streams {
904 stream.writer.create_directory()?;
905 }
906 commit_metadata(&prepared.root, &prepared.metadata_path, &prepared.metadata)?;
907 let manifest = Arc::new(RecordingManifest::new(
908 prepared.root.clone(),
909 prepared.metadata_path.clone(),
910 prepared.metadata,
911 ));
912 Self::start_new_prepared(prepared.root, prepared.streams, manifest, lease)
913 }
914
915 fn continue_recording(
917 builder: SystemStateWriterBuilder,
918 checkpoint: Option<CheckpointRequest>,
919 ) -> Result<(Self, Option<SystemState>), StorageError> {
920 let prepared = PreparedRecording::from_builder(builder)?;
921 let lease = RecordingLease::acquire(&prepared.root)?;
922 remove_stale_metadata_temp(&prepared.root)?;
923 let mut existing = load_metadata(&prepared.metadata_path)?;
924 if !matches!(existing.status, RecordingStatus::Running) {
925 return Err(StorageError::RecordingNotContinuable {
926 path: prepared.metadata_path,
927 });
928 }
929 ensure_resume_match(&prepared.metadata_path, &prepared.metadata, &existing)?;
930
931 if checkpoint.is_some() {
932 for stream in &prepared.streams {
933 let declaration = existing
934 .stream(&stream.name)
935 .expect("matched metadata contains every prepared stream");
936 StateWriterWorker::recover_state_stream(&stream.writer, declaration)?;
937 }
938 }
939
940 let state = if let Some(checkpoint) = checkpoint {
941 let (checkpoint_stream, decoders) = match checkpoint {
942 CheckpointRequest::Named(stream, decoders) => (stream, decoders),
943 CheckpointRequest::LatestComplete(decoders) => {
944 let stream = existing
945 .streams
946 .iter()
947 .filter(|stream| {
948 stored_state_series_reader::is_complete_checkpoint_stream(
949 stream,
950 &prepared.spec,
951 ) && !stream.chunks.is_empty()
952 })
953 .max_by_key(|stream| stream.chunks.last().map(|chunk| chunk.last_iteration))
954 .map(|stream| stream.name.clone())
955 .ok_or(StorageError::NoCompleteCheckpoint)?;
956 (stream, decoders)
957 }
958 };
959 let declaration = existing.stream(&checkpoint_stream).ok_or_else(|| {
960 StorageError::UnknownStateStream {
961 stream: checkpoint_stream.clone(),
962 }
963 })?;
964 let state = stored_state_series_reader::decode_resume_state(
965 &prepared.root,
966 &prepared.metadata_path,
967 declaration,
968 &prepared.spec,
969 &decoders,
970 )?;
971 resume::prepare_rewind_after_checkpoint(
972 &prepared.root,
973 &prepared.metadata_path,
974 &mut existing,
975 state.simulation_time().iteration(),
976 )?;
977 commit_metadata(&prepared.root, &prepared.metadata_path, &existing)?;
978 Some(state)
979 } else {
980 None
981 };
982
983 let mut recovered = Vec::with_capacity(prepared.streams.len());
984 for stream in prepared.streams {
985 let declaration = existing
986 .stream(&stream.name)
987 .expect("matched metadata contains every prepared stream");
988 let seed = StateWriterWorker::recover_state_stream(&stream.writer, declaration)?;
989 recovered.push((stream, seed));
990 }
991
992 existing.timing.continuation_count = existing
993 .timing
994 .continuation_count
995 .checked_add(1)
996 .ok_or_else(|| StorageError::InvalidMetadata {
997 path: prepared.metadata_path.clone(),
998 reason: "timing.continuation_count overflowed".to_owned(),
999 })?;
1000 commit_metadata(&prepared.root, &prepared.metadata_path, &existing)?;
1001 let manifest = Arc::new(RecordingManifest::new(
1002 prepared.root.clone(),
1003 prepared.metadata_path.clone(),
1004 existing,
1005 ));
1006
1007 let output = Self::start_resumed_prepared(prepared.root, recovered, manifest, lease)?;
1008 Ok((output, state))
1009 }
1010
1011 fn start_new_prepared(
1013 root: PathBuf,
1014 streams: Vec<PreparedStateStream>,
1015 manifest: Arc<RecordingManifest>,
1016 lease: RecordingLease,
1017 ) -> Result<Self, StorageError> {
1018 let mut scheduled = HashMap::with_capacity(streams.len());
1019 let mut configs = Vec::with_capacity(streams.len());
1020 let mut stream_order = Vec::with_capacity(streams.len());
1021 for prepared in streams {
1022 let name = prepared.name;
1023 stream_order.push(name.clone());
1024 scheduled.insert(
1025 name,
1026 ScheduledStateStream {
1027 encoder: prepared.encoder,
1028 sampling_interval: prepared.sampling_interval,
1029 last_recorded_iteration: None,
1030 },
1031 );
1032 configs.push(prepared.writer);
1033 }
1034 let writer = StateWriterWorker::start_new_recording(configs, Arc::clone(&manifest))?;
1035 Ok(Self {
1036 root,
1037 stream_order,
1038 manifest,
1039 streams: scheduled,
1040 writer: Some(writer),
1041 session_started: Instant::now(),
1042 _lease: lease,
1043 })
1044 }
1045
1046 fn start_resumed_prepared(
1048 root: PathBuf,
1049 streams: Vec<(PreparedStateStream, RecoveredStateStream)>,
1050 manifest: Arc<RecordingManifest>,
1051 lease: RecordingLease,
1052 ) -> Result<Self, StorageError> {
1053 let mut scheduled = HashMap::with_capacity(streams.len());
1054 let mut recovered_streams = Vec::with_capacity(streams.len());
1055 let mut stream_order = Vec::with_capacity(streams.len());
1056 for (prepared, seed) in streams {
1057 let name = prepared.name;
1058 stream_order.push(name.clone());
1059 scheduled.insert(
1060 name,
1061 ScheduledStateStream {
1062 encoder: prepared.encoder,
1063 sampling_interval: prepared.sampling_interval,
1064 last_recorded_iteration: seed.last_iteration(),
1065 },
1066 );
1067 recovered_streams.push((prepared.writer, seed));
1068 }
1069 let writer = StateWriterWorker::continue_recovered_recording(
1070 recovered_streams,
1071 Arc::clone(&manifest),
1072 )?;
1073 Ok(Self {
1074 root,
1075 stream_order,
1076 manifest,
1077 streams: scheduled,
1078 writer: Some(writer),
1079 session_started: Instant::now(),
1080 _lease: lease,
1081 })
1082 }
1083
1084 fn finish_writer(&mut self) -> Result<(), StorageError> {
1086 let Some(writer) = self.writer.take() else {
1087 return Ok(());
1088 };
1089 writer.finish_recording()
1090 }
1091
1092 fn transition_terminal(
1094 &self,
1095 status: RecordingStatus,
1096 terminal_metadata: Map<String, Value>,
1097 ) -> Result<(), StorageError> {
1098 let finalized_at_utc =
1099 utc_now_rfc3339().map_err(|source| StorageError::OperationalTimestamp {
1100 operation: "finalize recording",
1101 source,
1102 })?;
1103 let active_duration_ns = duration_nanoseconds(self.session_started.elapsed())
1104 .ok_or(StorageError::OperationalDurationOverflow)?;
1105 self.manifest.transition_terminal(
1106 status,
1107 finalized_at_utc,
1108 active_duration_ns,
1109 terminal_metadata,
1110 )
1111 }
1112
1113 fn completed_recording(&self) -> Result<CompletedRecording, StorageError> {
1115 let metadata = self.manifest.snapshot();
1116 let timing = RecordingTiming::from_stored(&metadata.timing, &self.manifest.path)?;
1117 let streams = metadata
1118 .streams
1119 .iter()
1120 .map(completed_stream_summary)
1121 .collect::<Result<Vec<_>, _>>()?;
1122 Ok(CompletedRecording {
1123 directory: self.root.clone(),
1124 timing,
1125 terminal_metadata: metadata.terminal_metadata,
1126 streams,
1127 })
1128 }
1129}
1130
1131fn completed_stream_summary(
1133 stream: &StateStreamMetadata,
1134) -> Result<CompletedStreamSummary, StorageError> {
1135 let overflow = || StorageError::ByteCountOverflow {
1136 stream: stream.name.clone(),
1137 };
1138 let chunk_count = u64::try_from(stream.chunks.len()).map_err(|_| overflow())?;
1139 let record_count = stream
1140 .chunks
1141 .iter()
1142 .try_fold(0_u64, |total, chunk| total.checked_add(chunk.records))
1143 .ok_or_else(&overflow)?;
1144 let encoded_bytes = stream
1145 .chunks
1146 .iter()
1147 .try_fold(0_u64, |total, chunk| total.checked_add(chunk.bytes))
1148 .ok_or_else(overflow)?;
1149 Ok(CompletedStreamSummary {
1150 name: stream.name.clone(),
1151 chunk_count,
1152 record_count,
1153 encoded_bytes,
1154 first_iteration: stream.chunks.first().map(|chunk| chunk.first_iteration),
1155 last_iteration: stream.chunks.last().map(|chunk| chunk.last_iteration),
1156 })
1157}
1158
1159struct PreparedRecording {
1161 root: PathBuf,
1162 metadata_path: PathBuf,
1163 spec: SystemStateSchema,
1164 metadata: RecordingMetadata,
1165 streams: Vec<PreparedStateStream>,
1166}
1167
1168impl PreparedRecording {
1169 fn from_builder(builder: SystemStateWriterBuilder) -> Result<Self, StorageError> {
1171 let metadata_path = builder.root.join(METADATA_FILE);
1172 let stored_time = builder.time.into_stored();
1173 let mut names = HashSet::with_capacity(builder.streams.len());
1174 let mut directories = HashSet::with_capacity(builder.streams.len());
1175 let mut streams = Vec::with_capacity(builder.streams.len());
1176 let mut declarations = Vec::with_capacity(builder.streams.len());
1177
1178 for config in builder.streams {
1179 if !names.insert(config.name.clone()) {
1180 return Err(StorageError::DuplicateStateStream {
1181 stream: config.name,
1182 });
1183 }
1184 let directory = config.relative_directory().to_owned();
1185 if !directories.insert(directory.clone()) {
1186 return Err(StorageError::InvalidConfiguration {
1187 setting: "stream.directory",
1188 reason: format!("multiple streams use relative directory `{}`", directory),
1189 });
1190 }
1191
1192 let storage = config
1193 .storage
1194 .or(builder.shared_stream_storage)
1195 .ok_or_else(|| StorageError::InvalidConfiguration {
1196 setting: "stream.storage",
1197 reason: format!(
1198 "stream `{}` has no explicit storage and the writer has no shared storage",
1199 config.name
1200 ),
1201 })?;
1202 let encoder = JsonStateRecordEncoder::new(&config.name, &builder.spec, &config.fields)?;
1203 let fields = encoder
1204 .fields()
1205 .map(|name| {
1206 let field = builder
1207 .spec
1208 .field_schema(name)
1209 .expect("encoder fields were validated against this specification");
1210 StateFieldMetadata {
1211 name: name.to_owned(),
1212 description: field.description().map(str::to_owned),
1213 }
1214 })
1215 .collect::<Vec<_>>();
1216 declarations.push(StateStreamMetadata {
1217 name: config.name.clone(),
1218 directory: directory.clone(),
1219 sampling_interval: config.sampling_interval,
1220 fields,
1221 storage,
1222 chunks: Vec::new(),
1223 });
1224 streams.push(PreparedStateStream {
1225 name: config.name.clone(),
1226 encoder,
1227 sampling_interval: config.sampling_interval,
1228 writer: StateStreamStorageConfig::new(
1229 &config.name,
1230 builder.root.join(&directory),
1231 storage,
1232 )?,
1233 });
1234 }
1235
1236 let created_at_utc =
1237 utc_now_rfc3339().map_err(|source| StorageError::OperationalTimestamp {
1238 operation: "create recording",
1239 source,
1240 })?;
1241 let metadata = RecordingMetadata::running(
1242 stored_time,
1243 builder.user_metadata,
1244 declarations,
1245 created_at_utc,
1246 );
1247 metadata.validate(&metadata_path)?;
1248 Ok(Self {
1249 root: builder.root,
1250 metadata_path,
1251 spec: builder.spec,
1252 metadata,
1253 streams,
1254 })
1255 }
1256}
1257
1258struct PreparedStateStream {
1260 name: String,
1261 encoder: JsonStateRecordEncoder,
1262 sampling_interval: SamplingInterval,
1263 writer: StateStreamStorageConfig,
1264}
1265
1266struct ScheduledStateStream {
1268 encoder: JsonStateRecordEncoder,
1269 sampling_interval: SamplingInterval,
1270 last_recorded_iteration: Option<u64>,
1271}
1272
1273pub(crate) struct RecordingManifest {
1279 root: PathBuf,
1280 path: PathBuf,
1281 metadata: Mutex<RecordingMetadata>,
1282}
1283
1284impl RecordingManifest {
1285 fn new(root: PathBuf, path: PathBuf, metadata: RecordingMetadata) -> Self {
1287 Self {
1288 root,
1289 path,
1290 metadata: Mutex::new(metadata),
1291 }
1292 }
1293
1294 pub(crate) fn prepare_chunk(
1296 &self,
1297 stream: &str,
1298 descriptor: jsonl_format::ChunkMetadata,
1299 ) -> Result<(), StorageError> {
1300 let mut current = lock_metadata(&self.metadata);
1301 if !matches!(current.status, RecordingStatus::Running) {
1302 return Err(StorageError::RecordingFinished);
1303 }
1304 let mut candidate = current.clone();
1305 let declaration =
1306 candidate
1307 .stream_mut(stream)
1308 .ok_or_else(|| StorageError::UnknownStateStream {
1309 stream: stream.to_owned(),
1310 })?;
1311 let expected = u64::try_from(declaration.chunks.len()).map_err(|_| {
1312 StorageError::ByteCountOverflow {
1313 stream: stream.to_owned(),
1314 }
1315 })?;
1316 if descriptor.ordinal != expected {
1317 return Err(StorageError::InvalidMetadata {
1318 path: self.path.clone(),
1319 reason: format!(
1320 "stream `{stream}` prepared chunk ordinal {}, expected {expected}",
1321 descriptor.ordinal
1322 ),
1323 });
1324 }
1325 declaration.chunks.push(descriptor);
1326 commit_metadata(&self.root, &self.path, &candidate)?;
1327 *current = candidate;
1328 Ok(())
1329 }
1330
1331 fn transition_terminal(
1333 &self,
1334 status: RecordingStatus,
1335 finalized_at_utc: String,
1336 active_duration_ns: u64,
1337 terminal_metadata: Map<String, Value>,
1338 ) -> Result<(), StorageError> {
1339 let mut current = lock_metadata(&self.metadata);
1340 let mut candidate = current.clone();
1341 candidate.status = status;
1342 candidate.timing.finalized_at_utc = Some(finalized_at_utc);
1343 candidate.timing.active_duration_ns = candidate
1344 .timing
1345 .active_duration_ns
1346 .checked_add(active_duration_ns)
1347 .ok_or(StorageError::OperationalDurationOverflow)?;
1348 candidate.terminal_metadata = terminal_metadata;
1349 commit_metadata(&self.root, &self.path, &candidate)?;
1350 *current = candidate;
1351 Ok(())
1352 }
1353
1354 fn snapshot(&self) -> RecordingMetadata {
1356 lock_metadata(&self.metadata).clone()
1357 }
1358}
1359
1360struct RecordingLease {
1365 _directory: File,
1366}
1367
1368impl RecordingLease {
1369 fn acquire(root: &Path) -> Result<Self, StorageError> {
1371 let directory = File::open(root).map_err(|source| StorageError::Io {
1372 operation: "open output root for exclusive ownership",
1373 path: root.to_path_buf(),
1374 source,
1375 })?;
1376 match FileExt::try_lock_exclusive(&directory) {
1377 Ok(()) => Ok(Self {
1378 _directory: directory,
1379 }),
1380 Err(source) if source.kind() == std::io::ErrorKind::WouldBlock => {
1381 Err(StorageError::RecordingDirectoryInUse {
1382 path: root.to_path_buf(),
1383 })
1384 }
1385 Err(source) => Err(StorageError::Io {
1386 operation: "acquire exclusive output ownership",
1387 path: root.to_path_buf(),
1388 source,
1389 }),
1390 }
1391 }
1392}
1393
1394fn load_metadata(path: &Path) -> Result<RecordingMetadata, StorageError> {
1396 let bytes = fs::read(path).map_err(|source| StorageError::Io {
1397 operation: "read metadata for resume",
1398 path: path.to_path_buf(),
1399 source,
1400 })?;
1401 let metadata: RecordingMetadata =
1402 serde_json::from_slice(&bytes).map_err(|source| StorageError::Json {
1403 operation: "parse metadata for resume",
1404 path: path.to_path_buf(),
1405 source,
1406 })?;
1407 metadata.validate(path)?;
1408 Ok(metadata)
1409}
1410
1411fn ensure_resume_match(
1413 path: &Path,
1414 expected: &RecordingMetadata,
1415 existing: &RecordingMetadata,
1416) -> Result<(), StorageError> {
1417 let mut configuration = existing.clone();
1418 for stream in &mut configuration.streams {
1419 stream.chunks.clear();
1420 }
1421 configuration.status = RecordingStatus::Running;
1422 configuration.timing = expected.timing.clone();
1423 configuration.terminal_metadata.clear();
1424 if &configuration != expected {
1425 return Err(StorageError::RecordingConfigurationMismatch {
1426 path: path.to_path_buf(),
1427 reason: "builder time axis, user metadata, or stream declarations differ".to_owned(),
1428 });
1429 }
1430 Ok(())
1431}
1432
1433fn remove_stale_metadata_temp(root: &Path) -> Result<(), StorageError> {
1435 let path = root.join(METADATA_TEMP_FILE);
1436 match fs::remove_file(&path) {
1437 Ok(()) => File::open(root)
1438 .and_then(|directory| directory.sync_all())
1439 .map_err(|source| StorageError::Io {
1440 operation: "synchronize stale metadata cleanup",
1441 path: root.to_path_buf(),
1442 source,
1443 }),
1444 Err(source) if source.kind() == std::io::ErrorKind::NotFound => Ok(()),
1445 Err(source) => Err(StorageError::Io {
1446 operation: "remove stale temporary metadata",
1447 path,
1448 source,
1449 }),
1450 }
1451}
1452
1453fn lock_metadata(metadata: &Mutex<RecordingMetadata>) -> MutexGuard<'_, RecordingMetadata> {
1455 metadata
1456 .lock()
1457 .unwrap_or_else(|poisoned| poisoned.into_inner())
1458}
1459
1460fn ensure_absent(root: &Path) -> Result<(), StorageError> {
1462 match root.try_exists() {
1463 Ok(false) => Ok(()),
1464 Ok(true) => Err(StorageError::RecordingDirectoryExists {
1465 path: root.to_path_buf(),
1466 }),
1467 Err(source) => Err(StorageError::Io {
1468 operation: "inspect output root",
1469 path: root.to_path_buf(),
1470 source,
1471 }),
1472 }
1473}
1474
1475fn create_root(root: &Path) -> Result<(), StorageError> {
1477 if let Some(parent) = root.parent() {
1478 fs::create_dir_all(parent).map_err(|source| StorageError::Io {
1479 operation: "create recording parent directories",
1480 path: parent.to_path_buf(),
1481 source,
1482 })?;
1483 }
1484 match fs::create_dir(root) {
1485 Ok(()) => Ok(()),
1486 Err(source) if source.kind() == std::io::ErrorKind::AlreadyExists => {
1487 Err(StorageError::RecordingDirectoryExists {
1488 path: root.to_path_buf(),
1489 })
1490 }
1491 Err(source) => Err(StorageError::Io {
1492 operation: "create output root",
1493 path: root.to_path_buf(),
1494 source,
1495 }),
1496 }
1497}
1498
1499fn commit_metadata(
1507 root: &Path,
1508 metadata_path: &Path,
1509 metadata: &RecordingMetadata,
1510) -> Result<(), StorageError> {
1511 metadata.validate(metadata_path)?;
1512 let mut bytes = serde_json::to_vec_pretty(metadata).map_err(|source| StorageError::Json {
1513 operation: "serialize metadata",
1514 path: metadata_path.to_path_buf(),
1515 source,
1516 })?;
1517 bytes.push(b'\n');
1518
1519 let temporary_path = root.join(METADATA_TEMP_FILE);
1520 let result = write_and_replace_metadata(root, metadata_path, &temporary_path, &bytes);
1521 if result.is_err() {
1522 let _ = fs::remove_file(&temporary_path);
1523 }
1524 result
1525}
1526
1527fn write_and_replace_metadata(
1529 root: &Path,
1530 metadata_path: &Path,
1531 temporary_path: &Path,
1532 bytes: &[u8],
1533) -> Result<(), StorageError> {
1534 let mut temporary = OpenOptions::new()
1535 .write(true)
1536 .create_new(true)
1537 .open(temporary_path)
1538 .map_err(|source| StorageError::Io {
1539 operation: "create temporary metadata",
1540 path: temporary_path.to_path_buf(),
1541 source,
1542 })?;
1543 temporary
1544 .write_all(bytes)
1545 .map_err(|source| StorageError::Io {
1546 operation: "write temporary metadata",
1547 path: temporary_path.to_path_buf(),
1548 source,
1549 })?;
1550 temporary.sync_all().map_err(|source| StorageError::Io {
1551 operation: "sync temporary metadata",
1552 path: temporary_path.to_path_buf(),
1553 source,
1554 })?;
1555 drop(temporary);
1556
1557 fs::rename(temporary_path, metadata_path).map_err(|source| StorageError::Io {
1558 operation: "publish metadata",
1559 path: metadata_path.to_path_buf(),
1560 source,
1561 })?;
1562
1563 File::open(root)
1564 .and_then(|directory| directory.sync_all())
1565 .map_err(|source| StorageError::Io {
1566 operation: "sync output root",
1567 path: root.to_path_buf(),
1568 source,
1569 })
1570}