1use std::collections::{HashMap, HashSet};
58use std::fs::{self, File, OpenOptions};
59use std::io::Write;
60use std::num::NonZeroU64;
61use std::path::{Path, PathBuf};
62use std::sync::{Arc, Mutex, MutexGuard};
63use std::time::{Duration, Instant};
64
65use fs2::FileExt;
66use serde::{Deserialize, Serialize};
67use serde_json::{Map, Value};
68
69use crate::clock::{duration_nanoseconds, utc_now_rfc3339};
70use crate::configuration::TaskParameters;
71use crate::system_state::{StateSchemaSource, SystemState, SystemStateSchema};
72
73mod error;
74mod json_payload_decoder;
75mod json_state_record_encoder;
76mod jsonl_format;
77mod queued_state_writer;
78mod resume;
79mod stored_state_series_reader;
80
81pub use error::StorageError;
82pub use json_payload_decoder::{
83 JsonPayloadDecoder, JsonPayloadDecoderRegistry, JsonStringDecoder, JsonVecF64Decoder,
84};
85pub use stored_state_series_reader::StoredStateSeriesReader;
86
87use json_state_record_encoder::JsonStateRecordEncoder;
88use jsonl_format::{
89 RecordingMetadata, RecordingStatus, StateFieldMetadata, StateStreamMetadata,
90 TimeAxisMetadata as StoredTimeAxis,
91};
92use queued_state_writer::{RecoveredStateStream, StateStreamStorageConfig, StateWriterWorker};
93
94const METADATA_FILE: &str = "metadata.json";
96
97const METADATA_TEMP_FILE: &str = ".metadata.json.tmp";
99
100#[derive(Clone, Debug, Eq, PartialEq)]
107pub struct TimeAxisMetadata {
108 iteration_name: String,
109 iteration_unit: Option<String>,
110 physical_time_name: Option<String>,
111 physical_time_unit: Option<String>,
112}
113
114impl TimeAxisMetadata {
115 pub fn new(iteration_name: impl Into<String>) -> Self {
121 Self {
122 iteration_name: iteration_name.into(),
123 iteration_unit: None,
124 physical_time_name: None,
125 physical_time_unit: None,
126 }
127 }
128
129 #[must_use]
131 pub fn with_iteration_unit(mut self, unit: impl Into<String>) -> Self {
132 self.iteration_unit = Some(unit.into());
133 self
134 }
135
136 #[must_use]
138 pub fn with_physical_time_name(mut self, name: impl Into<String>) -> Self {
139 self.physical_time_name = Some(name.into());
140 self
141 }
142
143 #[must_use]
148 pub fn with_physical_time_unit(mut self, unit: impl Into<String>) -> Self {
149 self.physical_time_unit = Some(unit.into());
150 self
151 }
152
153 #[must_use]
155 pub fn with_physical_axis(mut self, name: impl Into<String>, unit: impl Into<String>) -> Self {
156 self.physical_time_name = Some(name.into());
157 self.physical_time_unit = Some(unit.into());
158 self
159 }
160
161 fn into_stored(self) -> StoredTimeAxis {
163 StoredTimeAxis {
164 iteration_name: self.iteration_name,
165 iteration_unit: self.iteration_unit,
166 physical_time_name: self.physical_time_name,
167 physical_time_unit: self.physical_time_unit,
168 }
169 }
170}
171
172impl Default for TimeAxisMetadata {
173 fn default() -> Self {
176 Self::new("iteration")
177 }
178}
179
180#[derive(Clone, Debug, Eq, PartialEq)]
186pub struct RecordingTiming {
187 created_at_utc: String,
188 finalized_at_utc: String,
189 active_duration_ns: u64,
190 continuation_count: u64,
191}
192
193impl RecordingTiming {
194 fn from_stored(
196 timing: &jsonl_format::RecordingTiming,
197 metadata_path: &Path,
198 ) -> Result<Self, StorageError> {
199 let finalized_at_utc =
200 timing
201 .finalized_at_utc
202 .clone()
203 .ok_or_else(|| StorageError::InvalidMetadata {
204 path: metadata_path.to_path_buf(),
205 reason: "completed recording lacks finalized timestamp".to_owned(),
206 })?;
207 Ok(Self {
208 created_at_utc: timing.created_at_utc.clone(),
209 finalized_at_utc,
210 active_duration_ns: timing.active_duration_ns,
211 continuation_count: timing.continuation_count,
212 })
213 }
214
215 pub fn created_at_utc(&self) -> &str {
217 &self.created_at_utc
218 }
219
220 pub fn finalized_at_utc(&self) -> &str {
222 &self.finalized_at_utc
223 }
224
225 pub fn active_duration_ns(&self) -> u64 {
227 self.active_duration_ns
228 }
229
230 pub fn active_duration(&self) -> Duration {
232 Duration::from_nanos(self.active_duration_ns)
233 }
234
235 pub fn continuation_count(&self) -> u64 {
237 self.continuation_count
238 }
239}
240
241#[derive(Clone, Debug, Eq, PartialEq)]
243pub struct CompletedStreamSummary {
244 name: String,
245 chunk_count: u64,
246 record_count: u64,
247 encoded_bytes: u64,
248 first_iteration: Option<u64>,
249 last_iteration: Option<u64>,
250}
251
252impl CompletedStreamSummary {
253 pub fn name(&self) -> &str {
255 &self.name
256 }
257
258 pub fn chunk_count(&self) -> u64 {
260 self.chunk_count
261 }
262
263 pub fn record_count(&self) -> u64 {
265 self.record_count
266 }
267
268 pub fn encoded_bytes(&self) -> u64 {
270 self.encoded_bytes
271 }
272
273 pub fn first_iteration(&self) -> Option<u64> {
275 self.first_iteration
276 }
277
278 pub fn last_iteration(&self) -> Option<u64> {
280 self.last_iteration
281 }
282}
283
284#[derive(Clone, Debug, Eq, PartialEq)]
289pub struct CompletedRecording {
290 directory: PathBuf,
291 timing: RecordingTiming,
292 terminal_metadata: Map<String, Value>,
293 streams: Vec<CompletedStreamSummary>,
294}
295
296impl CompletedRecording {
297 pub fn directory(&self) -> &Path {
299 &self.directory
300 }
301
302 pub fn timing(&self) -> &RecordingTiming {
304 &self.timing
305 }
306
307 pub fn terminal_metadata(&self) -> &Map<String, Value> {
309 &self.terminal_metadata
310 }
311
312 pub fn stream_summaries(&self) -> &[CompletedStreamSummary] {
314 &self.streams
315 }
316
317 pub fn stream_summary(&self, name: &str) -> Option<&CompletedStreamSummary> {
319 self.streams.iter().find(|stream| stream.name == name)
320 }
321}
322
323#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
334#[serde(rename_all = "snake_case")]
335pub enum SamplingInterval {
336 Iterations(NonZeroU64),
338}
339
340#[derive(Deserialize)]
341#[serde(untagged)]
342enum SamplingIntervalInput {
343 Iterations(NonZeroU64),
345 Tagged { iterations: NonZeroU64 },
347}
348
349impl<'de> Deserialize<'de> for SamplingInterval {
350 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
351 where
352 D: serde::Deserializer<'de>,
353 {
354 match SamplingIntervalInput::deserialize(deserializer)? {
355 SamplingIntervalInput::Iterations(interval)
356 | SamplingIntervalInput::Tagged {
357 iterations: interval,
358 } => Ok(Self::Iterations(interval)),
359 }
360 }
361}
362
363impl SamplingInterval {
364 pub const fn iterations(interval: u64) -> Option<Self> {
366 match NonZeroU64::new(interval) {
367 Some(interval) => Some(Self::Iterations(interval)),
368 None => None,
369 }
370 }
371
372 const fn includes(self, iteration: u64) -> bool {
374 match self {
375 Self::Iterations(interval) => iteration.is_multiple_of(interval.get()),
376 }
377 }
378}
379
380#[derive(Clone, Copy, Debug, Eq, PartialEq, Deserialize, Serialize)]
382#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
383pub enum StateStreamLayout {
384 Chunked { target_bytes: NonZeroU64 },
386 IndividualFiles,
388}
389
390#[derive(Clone, Copy, Debug, Eq, PartialEq, Deserialize, Serialize)]
392#[serde(deny_unknown_fields)]
393pub struct StateStreamStorage {
394 layout: StateStreamLayout,
395 storage_queue_bytes: NonZeroU64,
396}
397
398impl StateStreamStorage {
399 pub const fn chunked(target_bytes: NonZeroU64, storage_queue_bytes: NonZeroU64) -> Self {
401 Self {
402 layout: StateStreamLayout::Chunked { target_bytes },
403 storage_queue_bytes,
404 }
405 }
406
407 pub const fn individual_files(storage_queue_bytes: NonZeroU64) -> Self {
409 Self {
410 layout: StateStreamLayout::IndividualFiles,
411 storage_queue_bytes,
412 }
413 }
414
415 pub const fn layout(self) -> StateStreamLayout {
416 self.layout
417 }
418
419 pub const fn storage_queue_bytes(self) -> NonZeroU64 {
421 self.storage_queue_bytes
422 }
423}
424
425#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
432#[serde(deny_unknown_fields)]
433pub struct StateStreamConfig {
434 name: String,
435 #[serde(default, skip_serializing_if = "Option::is_none")]
436 directory: Option<String>,
437 sampling_interval: SamplingInterval,
438 fields: Vec<String>,
439 #[serde(default, skip_serializing_if = "Option::is_none")]
440 storage: Option<StateStreamStorage>,
441}
442
443impl StateStreamConfig {
444 pub fn new<I, K>(
452 name: impl Into<String>,
453 fields: I,
454 sampling_interval: SamplingInterval,
455 storage: Option<StateStreamStorage>,
456 ) -> Self
457 where
458 I: IntoIterator<Item = K>,
459 K: Into<String>,
460 {
461 let name = name.into();
462 Self {
463 directory: None,
464 name,
465 sampling_interval,
466 fields: fields.into_iter().map(Into::into).collect(),
467 storage,
468 }
469 }
470
471 #[must_use]
476 pub fn with_relative_directory(mut self, directory: impl Into<String>) -> Self {
477 self.directory = Some(directory.into());
478 self
479 }
480
481 pub fn name(&self) -> &str {
483 &self.name
484 }
485
486 pub fn relative_directory(&self) -> &str {
488 self.directory.as_deref().unwrap_or(&self.name)
489 }
490
491 pub const fn sampling_interval(&self) -> SamplingInterval {
493 self.sampling_interval
494 }
495
496 pub fn fields(&self) -> &[String] {
498 &self.fields
499 }
500
501 pub const fn storage(&self) -> Option<StateStreamStorage> {
503 self.storage
504 }
505}
506
507#[derive(Debug)]
514pub struct SystemStateWriterBuilder {
515 root: PathBuf,
516 spec: SystemStateSchema,
517 time: TimeAxisMetadata,
518 user_metadata: Map<String, Value>,
519 shared_stream_storage: Option<StateStreamStorage>,
520 streams: Vec<StateStreamConfig>,
521}
522
523impl SystemStateWriterBuilder {
524 pub fn new<S>(root: impl Into<PathBuf>, source: &S) -> Self
530 where
531 S: StateSchemaSource + ?Sized,
532 {
533 Self {
534 root: root.into(),
535 spec: source.state_schema().clone(),
536 time: TimeAxisMetadata::default(),
537 user_metadata: Map::new(),
538 shared_stream_storage: None,
539 streams: Vec::new(),
540 }
541 }
542
543 #[must_use]
545 pub fn with_time_axis_metadata(mut self, time: TimeAxisMetadata) -> Self {
546 self.time = time;
547 self
548 }
549
550 #[must_use]
556 pub fn with_user_metadata(mut self, metadata: Map<String, Value>) -> Self {
557 self.user_metadata.extend(metadata);
558 self
559 }
560
561 #[must_use]
567 pub fn with_shared_stream_storage(mut self, storage: StateStreamStorage) -> Self {
568 self.shared_stream_storage = Some(storage);
569 self
570 }
571
572 #[must_use]
580 pub fn with_task_parameters(mut self, parameters: &TaskParameters) -> Self {
581 self.user_metadata
582 .extend(parameters.resolved_object().clone());
583 self.user_metadata.insert(
584 "task_ordinal".to_owned(),
585 Value::from(parameters.task_ordinal()),
586 );
587 self
588 }
589
590 #[must_use]
595 pub fn add_state_stream(mut self, stream: StateStreamConfig) -> Self {
596 self.streams.push(stream);
597 self
598 }
599
600 pub fn create_new_recording(self) -> Result<SystemStateWriter, StorageError> {
612 SystemStateWriter::create_new_recording(self)
613 }
614
615 pub fn open_or_resume_from_latest_checkpoint(
623 self,
624 decoders: JsonPayloadDecoderRegistry,
625 ) -> Result<(SystemStateWriter, Option<SystemState>), StorageError> {
626 match self.root.try_exists() {
627 Ok(false) => Self::create_new_recording(self).map(|writer| (writer, None)),
628 Ok(true) => SystemStateWriter::continue_recording(
629 self,
630 Some(CheckpointRequest::LatestComplete(decoders)),
631 ),
632 Err(source) => Err(StorageError::Io {
633 operation: "inspect recording root for automatic resume",
634 path: self.root.clone(),
635 source,
636 }),
637 }
638 }
639
640 pub fn continue_existing_recording(self) -> Result<SystemStateWriter, StorageError> {
648 SystemStateWriter::continue_recording(self, None).map(|(writer, _)| writer)
649 }
650
651 pub fn continue_recording_from_latest_checkpoint(
660 self,
661 stream: &str,
662 decoders: JsonPayloadDecoderRegistry,
663 ) -> Result<(SystemStateWriter, SystemState), StorageError> {
664 let (writer, state) = SystemStateWriter::continue_recording(
665 self,
666 Some(CheckpointRequest::Named(stream.to_owned(), decoders)),
667 )?;
668 Ok((
669 writer,
670 state.expect("checkpoint-aware resume always reconstructs one state"),
671 ))
672 }
673}
674
675enum CheckpointRequest {
676 Named(String, JsonPayloadDecoderRegistry),
677 LatestComplete(JsonPayloadDecoderRegistry),
678}
679
680pub struct SystemStateWriter {
687 root: PathBuf,
688 stream_order: Vec<String>,
689 manifest: Arc<RecordingManifest>,
690 streams: HashMap<String, ScheduledStateStream>,
691 writer: Option<StateWriterWorker>,
692 session_started: Instant,
693 _lease: RecordingLease,
696}
697
698impl SystemStateWriter {
699 pub fn builder<S>(root: impl Into<PathBuf>, source: &S) -> SystemStateWriterBuilder
701 where
702 S: StateSchemaSource + ?Sized,
703 {
704 SystemStateWriterBuilder::new(root, source)
705 }
706
707 pub fn recording_directory(&self) -> &Path {
709 &self.root
710 }
711
712 pub fn stream_names(&self) -> impl ExactSizeIterator<Item = &str> {
714 self.stream_order.iter().map(String::as_str)
715 }
716
717 pub fn observe_state(&mut self, state: &SystemState) -> Result<(), StorageError> {
731 let iteration = state.simulation_time().iteration();
732 let writer = self
733 .writer
734 .as_ref()
735 .expect("an active recording owns its writer worker");
736 for name in &self.stream_order {
737 let stream = self
738 .streams
739 .get_mut(name)
740 .expect("stream order contains every configured stream");
741 if !stream.sampling_interval.includes(iteration)
742 || stream.last_recorded_iteration == Some(iteration)
743 {
744 continue;
745 }
746 let record = stream.encoder.encode(state)?;
747 writer.submit_record(name, record)?;
748 stream.last_recorded_iteration = Some(iteration);
749 }
750 Ok(())
751 }
752
753 pub fn flush_stream_to_storage(&self, stream: &str) -> Result<(), StorageError> {
760 if !self.streams.contains_key(stream) {
761 return Err(StorageError::UnknownStateStream {
762 stream: stream.to_owned(),
763 });
764 }
765 self.writer
766 .as_ref()
767 .expect("an active recording owns its writer worker")
768 .flush_state_stream(stream)
769 }
770
771 pub fn complete_recording(self) -> Result<CompletedRecording, StorageError> {
779 self.complete_recording_with_terminal_metadata(Map::new())
780 }
781
782 pub fn complete_recording_with_terminal_metadata(
788 mut self,
789 terminal_metadata: Map<String, Value>,
790 ) -> Result<CompletedRecording, StorageError> {
791 if let Err(error) = self.finish_writer() {
792 let _ = self.transition_terminal(
793 RecordingStatus::Failed {
794 message: error.to_string(),
795 },
796 Map::new(),
797 );
798 return Err(error);
799 }
800 self.transition_terminal(RecordingStatus::Complete, terminal_metadata)?;
801 self.completed_recording()
802 }
803
804 pub fn complete_recording_with_final_state(
811 mut self,
812 state: &SystemState,
813 ) -> Result<CompletedRecording, StorageError> {
814 self.record_final_state(state)?;
815 self.complete_recording()
816 }
817
818 pub fn complete_recording_with_final_state_and_terminal_metadata(
821 mut self,
822 state: &SystemState,
823 terminal_metadata: Map<String, Value>,
824 ) -> Result<CompletedRecording, StorageError> {
825 self.record_final_state(state)?;
826 self.complete_recording_with_terminal_metadata(terminal_metadata)
827 }
828
829 fn record_final_state(&mut self, state: &SystemState) -> Result<(), StorageError> {
831 let iteration = state.simulation_time().iteration();
832 let writer = self
833 .writer
834 .as_ref()
835 .expect("an active recording owns its writer worker");
836 for name in &self.stream_order {
837 let stream = self
838 .streams
839 .get_mut(name)
840 .expect("stream order contains every configured stream");
841 if stream.last_recorded_iteration == Some(iteration) {
842 continue;
843 }
844 let record = stream.encoder.encode(state)?;
845 writer.submit_record(name, record)?;
846 stream.last_recorded_iteration = Some(iteration);
847 }
848 Ok(())
849 }
850
851 pub fn mark_recording_failed(self, message: impl Into<String>) -> Result<(), StorageError> {
863 self.mark_recording_failed_with_terminal_metadata(message, Map::new())
864 }
865
866 pub fn mark_recording_failed_with_terminal_metadata(
868 mut self,
869 message: impl Into<String>,
870 terminal_metadata: Map<String, Value>,
871 ) -> Result<(), StorageError> {
872 let message = message.into();
873 if message.trim().is_empty() {
874 return Err(StorageError::InvalidConfiguration {
875 setting: "failure_message",
876 reason: "failed run message must not be empty".to_owned(),
877 });
878 }
879
880 if let Err(error) = self.finish_writer() {
881 let _ = self.transition_terminal(
882 RecordingStatus::Failed {
883 message: error.to_string(),
884 },
885 Map::new(),
886 );
887 return Err(error);
888 }
889 self.transition_terminal(RecordingStatus::Failed { message }, terminal_metadata)
890 }
891
892 fn create_new_recording(builder: SystemStateWriterBuilder) -> Result<Self, StorageError> {
894 ensure_absent(&builder.root)?;
895 let prepared = PreparedRecording::from_builder(builder)?;
896 create_root(&prepared.root)?;
897 let lease = RecordingLease::acquire(&prepared.root)?;
898 for stream in &prepared.streams {
899 stream.writer.create_directory()?;
900 }
901 commit_metadata(&prepared.root, &prepared.metadata_path, &prepared.metadata)?;
902 let manifest = Arc::new(RecordingManifest::new(
903 prepared.root.clone(),
904 prepared.metadata_path.clone(),
905 prepared.metadata,
906 ));
907 Self::start_new_prepared(prepared.root, prepared.streams, manifest, lease)
908 }
909
910 fn continue_recording(
912 builder: SystemStateWriterBuilder,
913 checkpoint: Option<CheckpointRequest>,
914 ) -> Result<(Self, Option<SystemState>), StorageError> {
915 let prepared = PreparedRecording::from_builder(builder)?;
916 let lease = RecordingLease::acquire(&prepared.root)?;
917 remove_stale_metadata_temp(&prepared.root)?;
918 let mut existing = load_metadata(&prepared.metadata_path)?;
919 if !matches!(existing.status, RecordingStatus::Running) {
920 return Err(StorageError::RecordingNotContinuable {
921 path: prepared.metadata_path,
922 });
923 }
924 ensure_resume_match(&prepared.metadata_path, &prepared.metadata, &existing)?;
925
926 if checkpoint.is_some() {
927 for stream in &prepared.streams {
928 let declaration = existing
929 .stream(&stream.name)
930 .expect("matched metadata contains every prepared stream");
931 StateWriterWorker::recover_state_stream(&stream.writer, declaration)?;
932 }
933 }
934
935 let state = if let Some(checkpoint) = checkpoint {
936 let (checkpoint_stream, decoders) = match checkpoint {
937 CheckpointRequest::Named(stream, decoders) => (stream, decoders),
938 CheckpointRequest::LatestComplete(decoders) => {
939 let stream = existing
940 .streams
941 .iter()
942 .filter(|stream| {
943 stored_state_series_reader::is_complete_checkpoint_stream(
944 stream,
945 &prepared.spec,
946 ) && !stream.chunks.is_empty()
947 })
948 .max_by_key(|stream| stream.chunks.last().map(|chunk| chunk.last_iteration))
949 .map(|stream| stream.name.clone())
950 .ok_or(StorageError::NoCompleteCheckpoint)?;
951 (stream, decoders)
952 }
953 };
954 let declaration = existing.stream(&checkpoint_stream).ok_or_else(|| {
955 StorageError::UnknownStateStream {
956 stream: checkpoint_stream.clone(),
957 }
958 })?;
959 let state = stored_state_series_reader::decode_resume_state(
960 &prepared.root,
961 &prepared.metadata_path,
962 declaration,
963 &prepared.spec,
964 &decoders,
965 )?;
966 resume::prepare_rewind_after_checkpoint(
967 &prepared.root,
968 &prepared.metadata_path,
969 &mut existing,
970 state.simulation_time().iteration(),
971 )?;
972 commit_metadata(&prepared.root, &prepared.metadata_path, &existing)?;
973 Some(state)
974 } else {
975 None
976 };
977
978 let mut recovered = Vec::with_capacity(prepared.streams.len());
979 for stream in prepared.streams {
980 let declaration = existing
981 .stream(&stream.name)
982 .expect("matched metadata contains every prepared stream");
983 let seed = StateWriterWorker::recover_state_stream(&stream.writer, declaration)?;
984 recovered.push((stream, seed));
985 }
986
987 existing.timing.continuation_count = existing
988 .timing
989 .continuation_count
990 .checked_add(1)
991 .ok_or_else(|| StorageError::InvalidMetadata {
992 path: prepared.metadata_path.clone(),
993 reason: "timing.continuation_count overflowed".to_owned(),
994 })?;
995 commit_metadata(&prepared.root, &prepared.metadata_path, &existing)?;
996 let manifest = Arc::new(RecordingManifest::new(
997 prepared.root.clone(),
998 prepared.metadata_path.clone(),
999 existing,
1000 ));
1001
1002 let output = Self::start_resumed_prepared(prepared.root, recovered, manifest, lease)?;
1003 Ok((output, state))
1004 }
1005
1006 fn start_new_prepared(
1008 root: PathBuf,
1009 streams: Vec<PreparedStateStream>,
1010 manifest: Arc<RecordingManifest>,
1011 lease: RecordingLease,
1012 ) -> Result<Self, StorageError> {
1013 let mut scheduled = HashMap::with_capacity(streams.len());
1014 let mut configs = Vec::with_capacity(streams.len());
1015 let mut stream_order = Vec::with_capacity(streams.len());
1016 for prepared in streams {
1017 let name = prepared.name;
1018 stream_order.push(name.clone());
1019 scheduled.insert(
1020 name,
1021 ScheduledStateStream {
1022 encoder: prepared.encoder,
1023 sampling_interval: prepared.sampling_interval,
1024 last_recorded_iteration: None,
1025 },
1026 );
1027 configs.push(prepared.writer);
1028 }
1029 let writer = StateWriterWorker::start_new_recording(configs, Arc::clone(&manifest))?;
1030 Ok(Self {
1031 root,
1032 stream_order,
1033 manifest,
1034 streams: scheduled,
1035 writer: Some(writer),
1036 session_started: Instant::now(),
1037 _lease: lease,
1038 })
1039 }
1040
1041 fn start_resumed_prepared(
1043 root: PathBuf,
1044 streams: Vec<(PreparedStateStream, RecoveredStateStream)>,
1045 manifest: Arc<RecordingManifest>,
1046 lease: RecordingLease,
1047 ) -> Result<Self, StorageError> {
1048 let mut scheduled = HashMap::with_capacity(streams.len());
1049 let mut recovered_streams = Vec::with_capacity(streams.len());
1050 let mut stream_order = Vec::with_capacity(streams.len());
1051 for (prepared, seed) in streams {
1052 let name = prepared.name;
1053 stream_order.push(name.clone());
1054 scheduled.insert(
1055 name,
1056 ScheduledStateStream {
1057 encoder: prepared.encoder,
1058 sampling_interval: prepared.sampling_interval,
1059 last_recorded_iteration: seed.last_iteration(),
1060 },
1061 );
1062 recovered_streams.push((prepared.writer, seed));
1063 }
1064 let writer = StateWriterWorker::continue_recovered_recording(
1065 recovered_streams,
1066 Arc::clone(&manifest),
1067 )?;
1068 Ok(Self {
1069 root,
1070 stream_order,
1071 manifest,
1072 streams: scheduled,
1073 writer: Some(writer),
1074 session_started: Instant::now(),
1075 _lease: lease,
1076 })
1077 }
1078
1079 fn finish_writer(&mut self) -> Result<(), StorageError> {
1081 let Some(writer) = self.writer.take() else {
1082 return Ok(());
1083 };
1084 writer.finish_recording()
1085 }
1086
1087 fn transition_terminal(
1089 &self,
1090 status: RecordingStatus,
1091 terminal_metadata: Map<String, Value>,
1092 ) -> Result<(), StorageError> {
1093 let finalized_at_utc =
1094 utc_now_rfc3339().map_err(|source| StorageError::OperationalTimestamp {
1095 operation: "finalize recording",
1096 source,
1097 })?;
1098 let active_duration_ns = duration_nanoseconds(self.session_started.elapsed())
1099 .ok_or(StorageError::OperationalDurationOverflow)?;
1100 self.manifest.transition_terminal(
1101 status,
1102 finalized_at_utc,
1103 active_duration_ns,
1104 terminal_metadata,
1105 )
1106 }
1107
1108 fn completed_recording(&self) -> Result<CompletedRecording, StorageError> {
1110 let metadata = self.manifest.snapshot();
1111 let timing = RecordingTiming::from_stored(&metadata.timing, &self.manifest.path)?;
1112 let streams = metadata
1113 .streams
1114 .iter()
1115 .map(completed_stream_summary)
1116 .collect::<Result<Vec<_>, _>>()?;
1117 Ok(CompletedRecording {
1118 directory: self.root.clone(),
1119 timing,
1120 terminal_metadata: metadata.terminal_metadata,
1121 streams,
1122 })
1123 }
1124}
1125
1126fn completed_stream_summary(
1128 stream: &StateStreamMetadata,
1129) -> Result<CompletedStreamSummary, StorageError> {
1130 let overflow = || StorageError::ByteCountOverflow {
1131 stream: stream.name.clone(),
1132 };
1133 let chunk_count = u64::try_from(stream.chunks.len()).map_err(|_| overflow())?;
1134 let record_count = stream
1135 .chunks
1136 .iter()
1137 .try_fold(0_u64, |total, chunk| total.checked_add(chunk.records))
1138 .ok_or_else(&overflow)?;
1139 let encoded_bytes = stream
1140 .chunks
1141 .iter()
1142 .try_fold(0_u64, |total, chunk| total.checked_add(chunk.bytes))
1143 .ok_or_else(overflow)?;
1144 Ok(CompletedStreamSummary {
1145 name: stream.name.clone(),
1146 chunk_count,
1147 record_count,
1148 encoded_bytes,
1149 first_iteration: stream.chunks.first().map(|chunk| chunk.first_iteration),
1150 last_iteration: stream.chunks.last().map(|chunk| chunk.last_iteration),
1151 })
1152}
1153
1154struct PreparedRecording {
1156 root: PathBuf,
1157 metadata_path: PathBuf,
1158 spec: SystemStateSchema,
1159 metadata: RecordingMetadata,
1160 streams: Vec<PreparedStateStream>,
1161}
1162
1163impl PreparedRecording {
1164 fn from_builder(builder: SystemStateWriterBuilder) -> Result<Self, StorageError> {
1166 let metadata_path = builder.root.join(METADATA_FILE);
1167 let stored_time = builder.time.into_stored();
1168 let mut names = HashSet::with_capacity(builder.streams.len());
1169 let mut directories = HashSet::with_capacity(builder.streams.len());
1170 let mut streams = Vec::with_capacity(builder.streams.len());
1171 let mut declarations = Vec::with_capacity(builder.streams.len());
1172
1173 for config in builder.streams {
1174 if !names.insert(config.name.clone()) {
1175 return Err(StorageError::DuplicateStateStream {
1176 stream: config.name,
1177 });
1178 }
1179 let directory = config.relative_directory().to_owned();
1180 if !directories.insert(directory.clone()) {
1181 return Err(StorageError::InvalidConfiguration {
1182 setting: "stream.directory",
1183 reason: format!("multiple streams use relative directory `{}`", directory),
1184 });
1185 }
1186
1187 let storage = config
1188 .storage
1189 .or(builder.shared_stream_storage)
1190 .ok_or_else(|| StorageError::InvalidConfiguration {
1191 setting: "stream.storage",
1192 reason: format!(
1193 "stream `{}` has no explicit storage and the writer has no shared storage",
1194 config.name
1195 ),
1196 })?;
1197 let encoder = JsonStateRecordEncoder::new(&config.name, &builder.spec, &config.fields)?;
1198 let fields = encoder
1199 .fields()
1200 .map(|name| {
1201 let field = builder
1202 .spec
1203 .field_schema(name)
1204 .expect("encoder fields were validated against this specification");
1205 StateFieldMetadata {
1206 name: name.to_owned(),
1207 description: field.description().map(str::to_owned),
1208 }
1209 })
1210 .collect::<Vec<_>>();
1211 declarations.push(StateStreamMetadata {
1212 name: config.name.clone(),
1213 directory: directory.clone(),
1214 sampling_interval: config.sampling_interval,
1215 fields,
1216 storage,
1217 chunks: Vec::new(),
1218 });
1219 streams.push(PreparedStateStream {
1220 name: config.name.clone(),
1221 encoder,
1222 sampling_interval: config.sampling_interval,
1223 writer: StateStreamStorageConfig::new(
1224 &config.name,
1225 builder.root.join(&directory),
1226 storage,
1227 )?,
1228 });
1229 }
1230
1231 let created_at_utc =
1232 utc_now_rfc3339().map_err(|source| StorageError::OperationalTimestamp {
1233 operation: "create recording",
1234 source,
1235 })?;
1236 let metadata = RecordingMetadata::running(
1237 stored_time,
1238 builder.user_metadata,
1239 declarations,
1240 created_at_utc,
1241 );
1242 metadata.validate(&metadata_path)?;
1243 Ok(Self {
1244 root: builder.root,
1245 metadata_path,
1246 spec: builder.spec,
1247 metadata,
1248 streams,
1249 })
1250 }
1251}
1252
1253struct PreparedStateStream {
1255 name: String,
1256 encoder: JsonStateRecordEncoder,
1257 sampling_interval: SamplingInterval,
1258 writer: StateStreamStorageConfig,
1259}
1260
1261struct ScheduledStateStream {
1263 encoder: JsonStateRecordEncoder,
1264 sampling_interval: SamplingInterval,
1265 last_recorded_iteration: Option<u64>,
1266}
1267
1268pub(crate) struct RecordingManifest {
1274 root: PathBuf,
1275 path: PathBuf,
1276 metadata: Mutex<RecordingMetadata>,
1277}
1278
1279impl RecordingManifest {
1280 fn new(root: PathBuf, path: PathBuf, metadata: RecordingMetadata) -> Self {
1282 Self {
1283 root,
1284 path,
1285 metadata: Mutex::new(metadata),
1286 }
1287 }
1288
1289 pub(crate) fn prepare_chunk(
1291 &self,
1292 stream: &str,
1293 descriptor: jsonl_format::ChunkMetadata,
1294 ) -> Result<(), StorageError> {
1295 let mut current = lock_metadata(&self.metadata);
1296 if !matches!(current.status, RecordingStatus::Running) {
1297 return Err(StorageError::RecordingFinished);
1298 }
1299 let mut candidate = current.clone();
1300 let declaration =
1301 candidate
1302 .stream_mut(stream)
1303 .ok_or_else(|| StorageError::UnknownStateStream {
1304 stream: stream.to_owned(),
1305 })?;
1306 let expected = u64::try_from(declaration.chunks.len()).map_err(|_| {
1307 StorageError::ByteCountOverflow {
1308 stream: stream.to_owned(),
1309 }
1310 })?;
1311 if descriptor.ordinal != expected {
1312 return Err(StorageError::InvalidMetadata {
1313 path: self.path.clone(),
1314 reason: format!(
1315 "stream `{stream}` prepared chunk ordinal {}, expected {expected}",
1316 descriptor.ordinal
1317 ),
1318 });
1319 }
1320 declaration.chunks.push(descriptor);
1321 commit_metadata(&self.root, &self.path, &candidate)?;
1322 *current = candidate;
1323 Ok(())
1324 }
1325
1326 fn transition_terminal(
1328 &self,
1329 status: RecordingStatus,
1330 finalized_at_utc: String,
1331 active_duration_ns: u64,
1332 terminal_metadata: Map<String, Value>,
1333 ) -> Result<(), StorageError> {
1334 let mut current = lock_metadata(&self.metadata);
1335 let mut candidate = current.clone();
1336 candidate.status = status;
1337 candidate.timing.finalized_at_utc = Some(finalized_at_utc);
1338 candidate.timing.active_duration_ns = candidate
1339 .timing
1340 .active_duration_ns
1341 .checked_add(active_duration_ns)
1342 .ok_or(StorageError::OperationalDurationOverflow)?;
1343 candidate.terminal_metadata = terminal_metadata;
1344 commit_metadata(&self.root, &self.path, &candidate)?;
1345 *current = candidate;
1346 Ok(())
1347 }
1348
1349 fn snapshot(&self) -> RecordingMetadata {
1351 lock_metadata(&self.metadata).clone()
1352 }
1353}
1354
1355struct RecordingLease {
1360 _directory: File,
1361}
1362
1363impl RecordingLease {
1364 fn acquire(root: &Path) -> Result<Self, StorageError> {
1366 let directory = File::open(root).map_err(|source| StorageError::Io {
1367 operation: "open output root for exclusive ownership",
1368 path: root.to_path_buf(),
1369 source,
1370 })?;
1371 match FileExt::try_lock_exclusive(&directory) {
1372 Ok(()) => Ok(Self {
1373 _directory: directory,
1374 }),
1375 Err(source) if source.kind() == std::io::ErrorKind::WouldBlock => {
1376 Err(StorageError::RecordingDirectoryInUse {
1377 path: root.to_path_buf(),
1378 })
1379 }
1380 Err(source) => Err(StorageError::Io {
1381 operation: "acquire exclusive output ownership",
1382 path: root.to_path_buf(),
1383 source,
1384 }),
1385 }
1386 }
1387}
1388
1389fn load_metadata(path: &Path) -> Result<RecordingMetadata, StorageError> {
1391 let bytes = fs::read(path).map_err(|source| StorageError::Io {
1392 operation: "read metadata for resume",
1393 path: path.to_path_buf(),
1394 source,
1395 })?;
1396 let metadata: RecordingMetadata =
1397 serde_json::from_slice(&bytes).map_err(|source| StorageError::Json {
1398 operation: "parse metadata for resume",
1399 path: path.to_path_buf(),
1400 source,
1401 })?;
1402 metadata.validate(path)?;
1403 Ok(metadata)
1404}
1405
1406fn ensure_resume_match(
1408 path: &Path,
1409 expected: &RecordingMetadata,
1410 existing: &RecordingMetadata,
1411) -> Result<(), StorageError> {
1412 let mut configuration = existing.clone();
1413 for stream in &mut configuration.streams {
1414 stream.chunks.clear();
1415 }
1416 configuration.status = RecordingStatus::Running;
1417 configuration.timing = expected.timing.clone();
1418 configuration.terminal_metadata.clear();
1419 if &configuration != expected {
1420 return Err(StorageError::RecordingConfigurationMismatch {
1421 path: path.to_path_buf(),
1422 reason: "builder time axis, user metadata, or stream declarations differ".to_owned(),
1423 });
1424 }
1425 Ok(())
1426}
1427
1428fn remove_stale_metadata_temp(root: &Path) -> Result<(), StorageError> {
1430 let path = root.join(METADATA_TEMP_FILE);
1431 match fs::remove_file(&path) {
1432 Ok(()) => File::open(root)
1433 .and_then(|directory| directory.sync_all())
1434 .map_err(|source| StorageError::Io {
1435 operation: "synchronize stale metadata cleanup",
1436 path: root.to_path_buf(),
1437 source,
1438 }),
1439 Err(source) if source.kind() == std::io::ErrorKind::NotFound => Ok(()),
1440 Err(source) => Err(StorageError::Io {
1441 operation: "remove stale temporary metadata",
1442 path,
1443 source,
1444 }),
1445 }
1446}
1447
1448fn lock_metadata(metadata: &Mutex<RecordingMetadata>) -> MutexGuard<'_, RecordingMetadata> {
1450 metadata
1451 .lock()
1452 .unwrap_or_else(|poisoned| poisoned.into_inner())
1453}
1454
1455fn ensure_absent(root: &Path) -> Result<(), StorageError> {
1457 match root.try_exists() {
1458 Ok(false) => Ok(()),
1459 Ok(true) => Err(StorageError::RecordingDirectoryExists {
1460 path: root.to_path_buf(),
1461 }),
1462 Err(source) => Err(StorageError::Io {
1463 operation: "inspect output root",
1464 path: root.to_path_buf(),
1465 source,
1466 }),
1467 }
1468}
1469
1470fn create_root(root: &Path) -> Result<(), StorageError> {
1472 if let Some(parent) = root.parent() {
1473 fs::create_dir_all(parent).map_err(|source| StorageError::Io {
1474 operation: "create recording parent directories",
1475 path: parent.to_path_buf(),
1476 source,
1477 })?;
1478 }
1479 match fs::create_dir(root) {
1480 Ok(()) => Ok(()),
1481 Err(source) if source.kind() == std::io::ErrorKind::AlreadyExists => {
1482 Err(StorageError::RecordingDirectoryExists {
1483 path: root.to_path_buf(),
1484 })
1485 }
1486 Err(source) => Err(StorageError::Io {
1487 operation: "create output root",
1488 path: root.to_path_buf(),
1489 source,
1490 }),
1491 }
1492}
1493
1494fn commit_metadata(
1502 root: &Path,
1503 metadata_path: &Path,
1504 metadata: &RecordingMetadata,
1505) -> Result<(), StorageError> {
1506 metadata.validate(metadata_path)?;
1507 let mut bytes = serde_json::to_vec_pretty(metadata).map_err(|source| StorageError::Json {
1508 operation: "serialize metadata",
1509 path: metadata_path.to_path_buf(),
1510 source,
1511 })?;
1512 bytes.push(b'\n');
1513
1514 let temporary_path = root.join(METADATA_TEMP_FILE);
1515 let result = write_and_replace_metadata(root, metadata_path, &temporary_path, &bytes);
1516 if result.is_err() {
1517 let _ = fs::remove_file(&temporary_path);
1518 }
1519 result
1520}
1521
1522fn write_and_replace_metadata(
1524 root: &Path,
1525 metadata_path: &Path,
1526 temporary_path: &Path,
1527 bytes: &[u8],
1528) -> Result<(), StorageError> {
1529 let mut temporary = OpenOptions::new()
1530 .write(true)
1531 .create_new(true)
1532 .open(temporary_path)
1533 .map_err(|source| StorageError::Io {
1534 operation: "create temporary metadata",
1535 path: temporary_path.to_path_buf(),
1536 source,
1537 })?;
1538 temporary
1539 .write_all(bytes)
1540 .map_err(|source| StorageError::Io {
1541 operation: "write temporary metadata",
1542 path: temporary_path.to_path_buf(),
1543 source,
1544 })?;
1545 temporary.sync_all().map_err(|source| StorageError::Io {
1546 operation: "sync temporary metadata",
1547 path: temporary_path.to_path_buf(),
1548 source,
1549 })?;
1550 drop(temporary);
1551
1552 fs::rename(temporary_path, metadata_path).map_err(|source| StorageError::Io {
1553 operation: "publish metadata",
1554 path: metadata_path.to_path_buf(),
1555 source,
1556 })?;
1557
1558 File::open(root)
1559 .and_then(|directory| directory.sync_all())
1560 .map_err(|source| StorageError::Io {
1561 operation: "sync output root",
1562 path: root.to_path_buf(),
1563 source,
1564 })
1565}