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 {
393 target_bytes: NonZeroU64,
395 },
396 IndividualFiles,
398}
399
400#[derive(Clone, Copy, Debug, Eq, PartialEq, Deserialize, Serialize)]
402#[serde(deny_unknown_fields)]
403pub struct StateStreamStorage {
404 layout: StateStreamLayout,
405 storage_queue_bytes: NonZeroU64,
406}
407
408impl StateStreamStorage {
409 pub const fn chunked(target_bytes: NonZeroU64, storage_queue_bytes: NonZeroU64) -> Self {
411 Self {
412 layout: StateStreamLayout::Chunked { target_bytes },
413 storage_queue_bytes,
414 }
415 }
416
417 pub const fn individual_files(storage_queue_bytes: NonZeroU64) -> Self {
419 Self {
420 layout: StateStreamLayout::IndividualFiles,
421 storage_queue_bytes,
422 }
423 }
424
425 pub const fn layout(self) -> StateStreamLayout {
427 self.layout
428 }
429
430 pub const fn storage_queue_bytes(self) -> NonZeroU64 {
432 self.storage_queue_bytes
433 }
434}
435
436#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
443#[serde(deny_unknown_fields)]
444pub struct StateStreamConfig {
445 name: String,
446 #[serde(default, skip_serializing_if = "Option::is_none")]
447 directory: Option<String>,
448 sampling_interval: SamplingInterval,
449 fields: Vec<String>,
450 #[serde(default, skip_serializing_if = "Option::is_none")]
451 storage: Option<StateStreamStorage>,
452}
453
454impl StateStreamConfig {
455 pub fn new<I, K>(
463 name: impl Into<String>,
464 fields: I,
465 sampling_interval: SamplingInterval,
466 storage: Option<StateStreamStorage>,
467 ) -> Self
468 where
469 I: IntoIterator<Item = K>,
470 K: Into<String>,
471 {
472 let name = name.into();
473 Self {
474 directory: None,
475 name,
476 sampling_interval,
477 fields: fields.into_iter().map(Into::into).collect(),
478 storage,
479 }
480 }
481
482 #[must_use]
487 pub fn with_relative_directory(mut self, directory: impl Into<String>) -> Self {
488 self.directory = Some(directory.into());
489 self
490 }
491
492 pub fn name(&self) -> &str {
494 &self.name
495 }
496
497 pub fn relative_directory(&self) -> &str {
499 self.directory.as_deref().unwrap_or(&self.name)
500 }
501
502 pub const fn sampling_interval(&self) -> SamplingInterval {
504 self.sampling_interval
505 }
506
507 pub fn fields(&self) -> &[String] {
509 &self.fields
510 }
511
512 pub const fn storage(&self) -> Option<StateStreamStorage> {
514 self.storage
515 }
516}
517
518#[derive(Debug)]
525pub struct SystemStateWriterBuilder {
526 root: PathBuf,
527 spec: SystemStateSchema,
528 time: TimeAxisMetadata,
529 user_metadata: Map<String, Value>,
530 shared_stream_storage: Option<StateStreamStorage>,
531 streams: Vec<StateStreamConfig>,
532}
533
534impl SystemStateWriterBuilder {
535 pub fn new<S>(root: impl Into<PathBuf>, source: &S) -> Self
541 where
542 S: StateSchemaSource + ?Sized,
543 {
544 Self {
545 root: root.into(),
546 spec: source.state_schema().clone(),
547 time: TimeAxisMetadata::default(),
548 user_metadata: Map::new(),
549 shared_stream_storage: None,
550 streams: Vec::new(),
551 }
552 }
553
554 #[must_use]
556 pub fn with_time_axis_metadata(mut self, time: TimeAxisMetadata) -> Self {
557 self.time = time;
558 self
559 }
560
561 #[must_use]
567 pub fn with_user_metadata(mut self, metadata: Map<String, Value>) -> Self {
568 self.user_metadata.extend(metadata);
569 self
570 }
571
572 #[must_use]
578 pub fn with_shared_stream_storage(mut self, storage: StateStreamStorage) -> Self {
579 self.shared_stream_storage = Some(storage);
580 self
581 }
582
583 #[must_use]
591 pub fn with_configuration(mut self, configuration: &ResolvedConfiguration) -> Self {
592 self.user_metadata
593 .extend(configuration.resolved_object().clone());
594 self.user_metadata
595 .insert("ordinal".to_owned(), Value::from(configuration.ordinal()));
596 self
597 }
598
599 #[must_use]
604 pub fn add_state_stream(mut self, stream: StateStreamConfig) -> Self {
605 self.streams.push(stream);
606 self
607 }
608
609 pub fn create_new_recording(self) -> Result<SystemStateWriter, StorageError> {
621 SystemStateWriter::create_new_recording(self)
622 }
623
624 pub fn open_or_resume_from_latest_checkpoint(
632 self,
633 decoders: JsonPayloadDecoderRegistry,
634 ) -> Result<(SystemStateWriter, Option<SystemState>), StorageError> {
635 match self.root.try_exists() {
636 Ok(false) => Self::create_new_recording(self).map(|writer| (writer, None)),
637 Ok(true) => SystemStateWriter::continue_recording(
638 self,
639 Some(CheckpointRequest::LatestComplete(decoders)),
640 ),
641 Err(source) => Err(StorageError::Io {
642 operation: "inspect recording root for automatic resume",
643 path: self.root.clone(),
644 source,
645 }),
646 }
647 }
648
649 pub fn continue_existing_recording(self) -> Result<SystemStateWriter, StorageError> {
657 SystemStateWriter::continue_recording(self, None).map(|(writer, _)| writer)
658 }
659
660 pub fn continue_recording_from_latest_checkpoint(
669 self,
670 stream: &str,
671 decoders: JsonPayloadDecoderRegistry,
672 ) -> Result<(SystemStateWriter, SystemState), StorageError> {
673 let (writer, state) = SystemStateWriter::continue_recording(
674 self,
675 Some(CheckpointRequest::Named(stream.to_owned(), decoders)),
676 )?;
677 Ok((
678 writer,
679 state.expect("checkpoint-aware resume always reconstructs one state"),
680 ))
681 }
682}
683
684enum CheckpointRequest {
685 Named(String, JsonPayloadDecoderRegistry),
686 LatestComplete(JsonPayloadDecoderRegistry),
687}
688
689pub struct SystemStateWriter {
696 root: PathBuf,
697 stream_order: Vec<String>,
698 manifest: Arc<RecordingManifest>,
699 streams: HashMap<String, ScheduledStateStream>,
700 writer: Option<StateWriterWorker>,
701 session_started: Instant,
702 _lease: RecordingLease,
705}
706
707impl SystemStateWriter {
708 pub fn builder<S>(root: impl Into<PathBuf>, source: &S) -> SystemStateWriterBuilder
710 where
711 S: StateSchemaSource + ?Sized,
712 {
713 SystemStateWriterBuilder::new(root, source)
714 }
715
716 pub fn recording_directory(&self) -> &Path {
718 &self.root
719 }
720
721 pub fn stream_names(&self) -> impl ExactSizeIterator<Item = &str> {
723 self.stream_order.iter().map(String::as_str)
724 }
725
726 pub fn observe_state(&mut self, state: &SystemState) -> Result<(), StorageError> {
740 let iteration = state.simulation_time().iteration();
741 let writer = self
742 .writer
743 .as_ref()
744 .expect("an active recording owns its writer worker");
745 for name in &self.stream_order {
746 let stream = self
747 .streams
748 .get_mut(name)
749 .expect("stream order contains every configured stream");
750 if !stream.sampling_interval.includes(iteration)
751 || stream.last_recorded_iteration == Some(iteration)
752 {
753 continue;
754 }
755 let record = stream.encoder.encode(state)?;
756 writer.submit_record(name, record)?;
757 stream.last_recorded_iteration = Some(iteration);
758 }
759 Ok(())
760 }
761
762 pub fn flush_stream_to_storage(&self, stream: &str) -> Result<(), StorageError> {
769 if !self.streams.contains_key(stream) {
770 return Err(StorageError::UnknownStateStream {
771 stream: stream.to_owned(),
772 });
773 }
774 self.writer
775 .as_ref()
776 .expect("an active recording owns its writer worker")
777 .flush_state_stream(stream)
778 }
779
780 pub fn complete_recording(self) -> Result<CompletedRecording, StorageError> {
788 self.complete_recording_with_terminal_metadata(Map::new())
789 }
790
791 pub fn complete_recording_with_terminal_metadata(
797 mut self,
798 terminal_metadata: Map<String, Value>,
799 ) -> Result<CompletedRecording, StorageError> {
800 if let Err(error) = self.finish_writer() {
801 let _ = self.transition_terminal(
802 RecordingStatus::Failed {
803 message: error.to_string(),
804 },
805 Map::new(),
806 );
807 return Err(error);
808 }
809 self.transition_terminal(RecordingStatus::Complete, terminal_metadata)?;
810 self.completed_recording()
811 }
812
813 pub fn complete_recording_with_final_state(
820 mut self,
821 state: &SystemState,
822 ) -> Result<CompletedRecording, StorageError> {
823 self.record_final_state(state)?;
824 self.complete_recording()
825 }
826
827 pub fn complete_recording_with_final_state_and_terminal_metadata(
830 mut self,
831 state: &SystemState,
832 terminal_metadata: Map<String, Value>,
833 ) -> Result<CompletedRecording, StorageError> {
834 self.record_final_state(state)?;
835 self.complete_recording_with_terminal_metadata(terminal_metadata)
836 }
837
838 fn record_final_state(&mut self, state: &SystemState) -> Result<(), StorageError> {
840 let iteration = state.simulation_time().iteration();
841 let writer = self
842 .writer
843 .as_ref()
844 .expect("an active recording owns its writer worker");
845 for name in &self.stream_order {
846 let stream = self
847 .streams
848 .get_mut(name)
849 .expect("stream order contains every configured stream");
850 if stream.last_recorded_iteration == Some(iteration) {
851 continue;
852 }
853 let record = stream.encoder.encode(state)?;
854 writer.submit_record(name, record)?;
855 stream.last_recorded_iteration = Some(iteration);
856 }
857 Ok(())
858 }
859
860 pub fn mark_recording_failed(self, message: impl Into<String>) -> Result<(), StorageError> {
872 self.mark_recording_failed_with_terminal_metadata(message, Map::new())
873 }
874
875 pub fn mark_recording_failed_with_terminal_metadata(
877 mut self,
878 message: impl Into<String>,
879 terminal_metadata: Map<String, Value>,
880 ) -> Result<(), StorageError> {
881 let message = message.into();
882 if message.trim().is_empty() {
883 return Err(StorageError::InvalidConfiguration {
884 setting: "failure_message",
885 reason: "failed run message must not be empty".to_owned(),
886 });
887 }
888
889 if let Err(error) = self.finish_writer() {
890 let _ = self.transition_terminal(
891 RecordingStatus::Failed {
892 message: error.to_string(),
893 },
894 Map::new(),
895 );
896 return Err(error);
897 }
898 self.transition_terminal(RecordingStatus::Failed { message }, terminal_metadata)
899 }
900
901 fn create_new_recording(builder: SystemStateWriterBuilder) -> Result<Self, StorageError> {
903 ensure_absent(&builder.root)?;
904 let prepared = PreparedRecording::from_builder(builder)?;
905 create_root(&prepared.root)?;
906 let lease = RecordingLease::acquire(&prepared.root)?;
907 for stream in &prepared.streams {
908 stream.writer.create_directory()?;
909 }
910 commit_metadata(&prepared.root, &prepared.metadata_path, &prepared.metadata)?;
911 let manifest = Arc::new(RecordingManifest::new(
912 prepared.root.clone(),
913 prepared.metadata_path.clone(),
914 prepared.metadata,
915 ));
916 Self::start_new_prepared(prepared.root, prepared.streams, manifest, lease)
917 }
918
919 fn continue_recording(
921 builder: SystemStateWriterBuilder,
922 checkpoint: Option<CheckpointRequest>,
923 ) -> Result<(Self, Option<SystemState>), StorageError> {
924 let prepared = PreparedRecording::from_builder(builder)?;
925 let lease = RecordingLease::acquire(&prepared.root)?;
926 remove_stale_metadata_temp(&prepared.root)?;
927 let mut existing = load_metadata(&prepared.metadata_path)?;
928 if !matches!(existing.status, RecordingStatus::Running) {
929 return Err(StorageError::RecordingNotContinuable {
930 path: prepared.metadata_path,
931 });
932 }
933 ensure_resume_match(&prepared.metadata_path, &prepared.metadata, &existing)?;
934
935 if checkpoint.is_some() {
936 for stream in &prepared.streams {
937 let declaration = existing
938 .stream(&stream.name)
939 .expect("matched metadata contains every prepared stream");
940 StateWriterWorker::recover_state_stream(&stream.writer, declaration)?;
941 }
942 }
943
944 let state = if let Some(checkpoint) = checkpoint {
945 let (checkpoint_stream, decoders) = match checkpoint {
946 CheckpointRequest::Named(stream, decoders) => (stream, decoders),
947 CheckpointRequest::LatestComplete(decoders) => {
948 let stream = existing
949 .streams
950 .iter()
951 .filter(|stream| {
952 stored_state_series_reader::is_complete_checkpoint_stream(
953 stream,
954 &prepared.spec,
955 ) && !stream.chunks.is_empty()
956 })
957 .max_by_key(|stream| stream.chunks.last().map(|chunk| chunk.last_iteration))
958 .map(|stream| stream.name.clone())
959 .ok_or(StorageError::NoCompleteCheckpoint)?;
960 (stream, decoders)
961 }
962 };
963 let declaration = existing.stream(&checkpoint_stream).ok_or_else(|| {
964 StorageError::UnknownStateStream {
965 stream: checkpoint_stream.clone(),
966 }
967 })?;
968 let state = stored_state_series_reader::decode_resume_state(
969 &prepared.root,
970 &prepared.metadata_path,
971 declaration,
972 &prepared.spec,
973 &decoders,
974 )?;
975 resume::prepare_rewind_after_checkpoint(
976 &prepared.root,
977 &prepared.metadata_path,
978 &mut existing,
979 state.simulation_time().iteration(),
980 )?;
981 commit_metadata(&prepared.root, &prepared.metadata_path, &existing)?;
982 Some(state)
983 } else {
984 None
985 };
986
987 let mut recovered = Vec::with_capacity(prepared.streams.len());
988 for stream in prepared.streams {
989 let declaration = existing
990 .stream(&stream.name)
991 .expect("matched metadata contains every prepared stream");
992 let seed = StateWriterWorker::recover_state_stream(&stream.writer, declaration)?;
993 recovered.push((stream, seed));
994 }
995
996 existing.timing.continuation_count = existing
997 .timing
998 .continuation_count
999 .checked_add(1)
1000 .ok_or_else(|| StorageError::InvalidMetadata {
1001 path: prepared.metadata_path.clone(),
1002 reason: "timing.continuation_count overflowed".to_owned(),
1003 })?;
1004 commit_metadata(&prepared.root, &prepared.metadata_path, &existing)?;
1005 let manifest = Arc::new(RecordingManifest::new(
1006 prepared.root.clone(),
1007 prepared.metadata_path.clone(),
1008 existing,
1009 ));
1010
1011 let output = Self::start_resumed_prepared(prepared.root, recovered, manifest, lease)?;
1012 Ok((output, state))
1013 }
1014
1015 fn start_new_prepared(
1017 root: PathBuf,
1018 streams: Vec<PreparedStateStream>,
1019 manifest: Arc<RecordingManifest>,
1020 lease: RecordingLease,
1021 ) -> Result<Self, StorageError> {
1022 let mut scheduled = HashMap::with_capacity(streams.len());
1023 let mut configs = Vec::with_capacity(streams.len());
1024 let mut stream_order = Vec::with_capacity(streams.len());
1025 for prepared in streams {
1026 let name = prepared.name;
1027 stream_order.push(name.clone());
1028 scheduled.insert(
1029 name,
1030 ScheduledStateStream {
1031 encoder: prepared.encoder,
1032 sampling_interval: prepared.sampling_interval,
1033 last_recorded_iteration: None,
1034 },
1035 );
1036 configs.push(prepared.writer);
1037 }
1038 let writer = StateWriterWorker::start_new_recording(configs, Arc::clone(&manifest))?;
1039 Ok(Self {
1040 root,
1041 stream_order,
1042 manifest,
1043 streams: scheduled,
1044 writer: Some(writer),
1045 session_started: Instant::now(),
1046 _lease: lease,
1047 })
1048 }
1049
1050 fn start_resumed_prepared(
1052 root: PathBuf,
1053 streams: Vec<(PreparedStateStream, RecoveredStateStream)>,
1054 manifest: Arc<RecordingManifest>,
1055 lease: RecordingLease,
1056 ) -> Result<Self, StorageError> {
1057 let mut scheduled = HashMap::with_capacity(streams.len());
1058 let mut recovered_streams = Vec::with_capacity(streams.len());
1059 let mut stream_order = Vec::with_capacity(streams.len());
1060 for (prepared, seed) in streams {
1061 let name = prepared.name;
1062 stream_order.push(name.clone());
1063 scheduled.insert(
1064 name,
1065 ScheduledStateStream {
1066 encoder: prepared.encoder,
1067 sampling_interval: prepared.sampling_interval,
1068 last_recorded_iteration: seed.last_iteration(),
1069 },
1070 );
1071 recovered_streams.push((prepared.writer, seed));
1072 }
1073 let writer = StateWriterWorker::continue_recovered_recording(
1074 recovered_streams,
1075 Arc::clone(&manifest),
1076 )?;
1077 Ok(Self {
1078 root,
1079 stream_order,
1080 manifest,
1081 streams: scheduled,
1082 writer: Some(writer),
1083 session_started: Instant::now(),
1084 _lease: lease,
1085 })
1086 }
1087
1088 fn finish_writer(&mut self) -> Result<(), StorageError> {
1090 let Some(writer) = self.writer.take() else {
1091 return Ok(());
1092 };
1093 writer.finish_recording()
1094 }
1095
1096 fn transition_terminal(
1098 &self,
1099 status: RecordingStatus,
1100 terminal_metadata: Map<String, Value>,
1101 ) -> Result<(), StorageError> {
1102 let finalized_at_utc =
1103 utc_now_rfc3339().map_err(|source| StorageError::OperationalTimestamp {
1104 operation: "finalize recording",
1105 source,
1106 })?;
1107 let active_duration_ns = duration_nanoseconds(self.session_started.elapsed())
1108 .ok_or(StorageError::OperationalDurationOverflow)?;
1109 self.manifest.transition_terminal(
1110 status,
1111 finalized_at_utc,
1112 active_duration_ns,
1113 terminal_metadata,
1114 )
1115 }
1116
1117 fn completed_recording(&self) -> Result<CompletedRecording, StorageError> {
1119 let metadata = self.manifest.snapshot();
1120 let timing = RecordingTiming::from_stored(&metadata.timing, &self.manifest.path)?;
1121 let streams = metadata
1122 .streams
1123 .iter()
1124 .map(completed_stream_summary)
1125 .collect::<Result<Vec<_>, _>>()?;
1126 Ok(CompletedRecording {
1127 directory: self.root.clone(),
1128 timing,
1129 terminal_metadata: metadata.terminal_metadata,
1130 streams,
1131 })
1132 }
1133}
1134
1135fn completed_stream_summary(
1137 stream: &StateStreamMetadata,
1138) -> Result<CompletedStreamSummary, StorageError> {
1139 let overflow = || StorageError::ByteCountOverflow {
1140 stream: stream.name.clone(),
1141 };
1142 let chunk_count = u64::try_from(stream.chunks.len()).map_err(|_| overflow())?;
1143 let record_count = stream
1144 .chunks
1145 .iter()
1146 .try_fold(0_u64, |total, chunk| total.checked_add(chunk.records))
1147 .ok_or_else(&overflow)?;
1148 let encoded_bytes = stream
1149 .chunks
1150 .iter()
1151 .try_fold(0_u64, |total, chunk| total.checked_add(chunk.bytes))
1152 .ok_or_else(overflow)?;
1153 Ok(CompletedStreamSummary {
1154 name: stream.name.clone(),
1155 chunk_count,
1156 record_count,
1157 encoded_bytes,
1158 first_iteration: stream.chunks.first().map(|chunk| chunk.first_iteration),
1159 last_iteration: stream.chunks.last().map(|chunk| chunk.last_iteration),
1160 })
1161}
1162
1163struct PreparedRecording {
1165 root: PathBuf,
1166 metadata_path: PathBuf,
1167 spec: SystemStateSchema,
1168 metadata: RecordingMetadata,
1169 streams: Vec<PreparedStateStream>,
1170}
1171
1172impl PreparedRecording {
1173 fn from_builder(builder: SystemStateWriterBuilder) -> Result<Self, StorageError> {
1175 let metadata_path = builder.root.join(METADATA_FILE);
1176 let stored_time = builder.time.into_stored();
1177 let mut names = HashSet::with_capacity(builder.streams.len());
1178 let mut directories = HashSet::with_capacity(builder.streams.len());
1179 let mut streams = Vec::with_capacity(builder.streams.len());
1180 let mut declarations = Vec::with_capacity(builder.streams.len());
1181
1182 for config in builder.streams {
1183 if !names.insert(config.name.clone()) {
1184 return Err(StorageError::DuplicateStateStream {
1185 stream: config.name,
1186 });
1187 }
1188 let directory = config.relative_directory().to_owned();
1189 if !directories.insert(directory.clone()) {
1190 return Err(StorageError::InvalidConfiguration {
1191 setting: "stream.directory",
1192 reason: format!("multiple streams use relative directory `{}`", directory),
1193 });
1194 }
1195
1196 let storage = config
1197 .storage
1198 .or(builder.shared_stream_storage)
1199 .ok_or_else(|| StorageError::InvalidConfiguration {
1200 setting: "stream.storage",
1201 reason: format!(
1202 "stream `{}` has no explicit storage and the writer has no shared storage",
1203 config.name
1204 ),
1205 })?;
1206 let encoder = JsonStateRecordEncoder::new(&config.name, &builder.spec, &config.fields)?;
1207 let fields = encoder
1208 .fields()
1209 .map(|name| {
1210 let field = builder
1211 .spec
1212 .field_schema(name)
1213 .expect("encoder fields were validated against this specification");
1214 StateFieldMetadata {
1215 name: name.to_owned(),
1216 description: field.description().map(str::to_owned),
1217 }
1218 })
1219 .collect::<Vec<_>>();
1220 declarations.push(StateStreamMetadata {
1221 name: config.name.clone(),
1222 directory: directory.clone(),
1223 sampling_interval: config.sampling_interval,
1224 fields,
1225 storage,
1226 chunks: Vec::new(),
1227 });
1228 streams.push(PreparedStateStream {
1229 name: config.name.clone(),
1230 encoder,
1231 sampling_interval: config.sampling_interval,
1232 writer: StateStreamStorageConfig::new(
1233 &config.name,
1234 builder.root.join(&directory),
1235 storage,
1236 )?,
1237 });
1238 }
1239
1240 let created_at_utc =
1241 utc_now_rfc3339().map_err(|source| StorageError::OperationalTimestamp {
1242 operation: "create recording",
1243 source,
1244 })?;
1245 let metadata = RecordingMetadata::running(
1246 stored_time,
1247 builder.user_metadata,
1248 declarations,
1249 created_at_utc,
1250 );
1251 metadata.validate(&metadata_path)?;
1252 Ok(Self {
1253 root: builder.root,
1254 metadata_path,
1255 spec: builder.spec,
1256 metadata,
1257 streams,
1258 })
1259 }
1260}
1261
1262struct PreparedStateStream {
1264 name: String,
1265 encoder: JsonStateRecordEncoder,
1266 sampling_interval: SamplingInterval,
1267 writer: StateStreamStorageConfig,
1268}
1269
1270struct ScheduledStateStream {
1272 encoder: JsonStateRecordEncoder,
1273 sampling_interval: SamplingInterval,
1274 last_recorded_iteration: Option<u64>,
1275}
1276
1277pub(crate) struct RecordingManifest {
1283 root: PathBuf,
1284 path: PathBuf,
1285 metadata: Mutex<RecordingMetadata>,
1286}
1287
1288impl RecordingManifest {
1289 fn new(root: PathBuf, path: PathBuf, metadata: RecordingMetadata) -> Self {
1291 Self {
1292 root,
1293 path,
1294 metadata: Mutex::new(metadata),
1295 }
1296 }
1297
1298 pub(crate) fn prepare_chunk(
1300 &self,
1301 stream: &str,
1302 descriptor: jsonl_format::ChunkMetadata,
1303 ) -> Result<(), StorageError> {
1304 let mut current = lock_metadata(&self.metadata);
1305 if !matches!(current.status, RecordingStatus::Running) {
1306 return Err(StorageError::RecordingFinished);
1307 }
1308 let mut candidate = current.clone();
1309 let declaration =
1310 candidate
1311 .stream_mut(stream)
1312 .ok_or_else(|| StorageError::UnknownStateStream {
1313 stream: stream.to_owned(),
1314 })?;
1315 let expected = u64::try_from(declaration.chunks.len()).map_err(|_| {
1316 StorageError::ByteCountOverflow {
1317 stream: stream.to_owned(),
1318 }
1319 })?;
1320 if descriptor.ordinal != expected {
1321 return Err(StorageError::InvalidMetadata {
1322 path: self.path.clone(),
1323 reason: format!(
1324 "stream `{stream}` prepared chunk ordinal {}, expected {expected}",
1325 descriptor.ordinal
1326 ),
1327 });
1328 }
1329 declaration.chunks.push(descriptor);
1330 commit_metadata(&self.root, &self.path, &candidate)?;
1331 *current = candidate;
1332 Ok(())
1333 }
1334
1335 fn transition_terminal(
1337 &self,
1338 status: RecordingStatus,
1339 finalized_at_utc: String,
1340 active_duration_ns: u64,
1341 terminal_metadata: Map<String, Value>,
1342 ) -> Result<(), StorageError> {
1343 let mut current = lock_metadata(&self.metadata);
1344 let mut candidate = current.clone();
1345 candidate.status = status;
1346 candidate.timing.finalized_at_utc = Some(finalized_at_utc);
1347 candidate.timing.active_duration_ns = candidate
1348 .timing
1349 .active_duration_ns
1350 .checked_add(active_duration_ns)
1351 .ok_or(StorageError::OperationalDurationOverflow)?;
1352 candidate.terminal_metadata = terminal_metadata;
1353 commit_metadata(&self.root, &self.path, &candidate)?;
1354 *current = candidate;
1355 Ok(())
1356 }
1357
1358 fn snapshot(&self) -> RecordingMetadata {
1360 lock_metadata(&self.metadata).clone()
1361 }
1362}
1363
1364struct RecordingLease {
1369 _directory: File,
1370}
1371
1372impl RecordingLease {
1373 fn acquire(root: &Path) -> Result<Self, StorageError> {
1375 let directory = File::open(root).map_err(|source| StorageError::Io {
1376 operation: "open output root for exclusive ownership",
1377 path: root.to_path_buf(),
1378 source,
1379 })?;
1380 match FileExt::try_lock_exclusive(&directory) {
1381 Ok(()) => Ok(Self {
1382 _directory: directory,
1383 }),
1384 Err(source) if source.kind() == std::io::ErrorKind::WouldBlock => {
1385 Err(StorageError::RecordingDirectoryInUse {
1386 path: root.to_path_buf(),
1387 })
1388 }
1389 Err(source) => Err(StorageError::Io {
1390 operation: "acquire exclusive output ownership",
1391 path: root.to_path_buf(),
1392 source,
1393 }),
1394 }
1395 }
1396}
1397
1398fn load_metadata(path: &Path) -> Result<RecordingMetadata, StorageError> {
1400 let bytes = fs::read(path).map_err(|source| StorageError::Io {
1401 operation: "read metadata for resume",
1402 path: path.to_path_buf(),
1403 source,
1404 })?;
1405 let metadata: RecordingMetadata =
1406 serde_json::from_slice(&bytes).map_err(|source| StorageError::Json {
1407 operation: "parse metadata for resume",
1408 path: path.to_path_buf(),
1409 source,
1410 })?;
1411 metadata.validate(path)?;
1412 Ok(metadata)
1413}
1414
1415fn ensure_resume_match(
1417 path: &Path,
1418 expected: &RecordingMetadata,
1419 existing: &RecordingMetadata,
1420) -> Result<(), StorageError> {
1421 let mut configuration = existing.clone();
1422 for stream in &mut configuration.streams {
1423 stream.chunks.clear();
1424 }
1425 configuration.status = RecordingStatus::Running;
1426 configuration.timing = expected.timing.clone();
1427 configuration.terminal_metadata.clear();
1428 if &configuration != expected {
1429 return Err(StorageError::RecordingConfigurationMismatch {
1430 path: path.to_path_buf(),
1431 reason: "builder time axis, user metadata, or stream declarations differ".to_owned(),
1432 });
1433 }
1434 Ok(())
1435}
1436
1437fn remove_stale_metadata_temp(root: &Path) -> Result<(), StorageError> {
1439 let path = root.join(METADATA_TEMP_FILE);
1440 match fs::remove_file(&path) {
1441 Ok(()) => File::open(root)
1442 .and_then(|directory| directory.sync_all())
1443 .map_err(|source| StorageError::Io {
1444 operation: "synchronize stale metadata cleanup",
1445 path: root.to_path_buf(),
1446 source,
1447 }),
1448 Err(source) if source.kind() == std::io::ErrorKind::NotFound => Ok(()),
1449 Err(source) => Err(StorageError::Io {
1450 operation: "remove stale temporary metadata",
1451 path,
1452 source,
1453 }),
1454 }
1455}
1456
1457fn lock_metadata(metadata: &Mutex<RecordingMetadata>) -> MutexGuard<'_, RecordingMetadata> {
1459 metadata
1460 .lock()
1461 .unwrap_or_else(|poisoned| poisoned.into_inner())
1462}
1463
1464fn ensure_absent(root: &Path) -> Result<(), StorageError> {
1466 match root.try_exists() {
1467 Ok(false) => Ok(()),
1468 Ok(true) => Err(StorageError::RecordingDirectoryExists {
1469 path: root.to_path_buf(),
1470 }),
1471 Err(source) => Err(StorageError::Io {
1472 operation: "inspect output root",
1473 path: root.to_path_buf(),
1474 source,
1475 }),
1476 }
1477}
1478
1479fn create_root(root: &Path) -> Result<(), StorageError> {
1481 if let Some(parent) = root.parent() {
1482 fs::create_dir_all(parent).map_err(|source| StorageError::Io {
1483 operation: "create recording parent directories",
1484 path: parent.to_path_buf(),
1485 source,
1486 })?;
1487 }
1488 match fs::create_dir(root) {
1489 Ok(()) => Ok(()),
1490 Err(source) if source.kind() == std::io::ErrorKind::AlreadyExists => {
1491 Err(StorageError::RecordingDirectoryExists {
1492 path: root.to_path_buf(),
1493 })
1494 }
1495 Err(source) => Err(StorageError::Io {
1496 operation: "create output root",
1497 path: root.to_path_buf(),
1498 source,
1499 }),
1500 }
1501}
1502
1503fn commit_metadata(
1511 root: &Path,
1512 metadata_path: &Path,
1513 metadata: &RecordingMetadata,
1514) -> Result<(), StorageError> {
1515 metadata.validate(metadata_path)?;
1516 let mut bytes = serde_json::to_vec_pretty(metadata).map_err(|source| StorageError::Json {
1517 operation: "serialize metadata",
1518 path: metadata_path.to_path_buf(),
1519 source,
1520 })?;
1521 bytes.push(b'\n');
1522
1523 let temporary_path = root.join(METADATA_TEMP_FILE);
1524 let result = write_and_replace_metadata(root, metadata_path, &temporary_path, &bytes);
1525 if result.is_err() {
1526 let _ = fs::remove_file(&temporary_path);
1527 }
1528 result
1529}
1530
1531fn write_and_replace_metadata(
1533 root: &Path,
1534 metadata_path: &Path,
1535 temporary_path: &Path,
1536 bytes: &[u8],
1537) -> Result<(), StorageError> {
1538 let mut temporary = OpenOptions::new()
1539 .write(true)
1540 .create_new(true)
1541 .open(temporary_path)
1542 .map_err(|source| StorageError::Io {
1543 operation: "create temporary metadata",
1544 path: temporary_path.to_path_buf(),
1545 source,
1546 })?;
1547 temporary
1548 .write_all(bytes)
1549 .map_err(|source| StorageError::Io {
1550 operation: "write temporary metadata",
1551 path: temporary_path.to_path_buf(),
1552 source,
1553 })?;
1554 temporary.sync_all().map_err(|source| StorageError::Io {
1555 operation: "sync temporary metadata",
1556 path: temporary_path.to_path_buf(),
1557 source,
1558 })?;
1559 drop(temporary);
1560
1561 fs::rename(temporary_path, metadata_path).map_err(|source| StorageError::Io {
1562 operation: "publish metadata",
1563 path: metadata_path.to_path_buf(),
1564 source,
1565 })?;
1566
1567 File::open(root)
1568 .and_then(|directory| directory.sync_all())
1569 .map_err(|source| StorageError::Io {
1570 operation: "sync output root",
1571 path: root.to_path_buf(),
1572 source,
1573 })
1574}