1use std::collections::{HashMap, HashSet};
54use std::fs::{self, File, OpenOptions};
55use std::io::Write;
56use std::num::NonZeroU64;
57use std::path::{Path, PathBuf};
58use std::sync::{Arc, Mutex, MutexGuard};
59use std::time::{Duration, Instant};
60
61use fs2::FileExt;
62use serde::{Deserialize, Serialize};
63use serde_json::{Map, Value};
64
65use crate::clock::{duration_nanoseconds, utc_now_rfc3339};
66use crate::configuration::TaskParameters;
67use crate::system_state::{StateSchemaSource, SystemState, SystemStateSchema};
68
69mod error;
70mod json_payload_decoder;
71mod json_state_record_encoder;
72mod jsonl_format;
73mod queued_state_writer;
74mod stored_state_series_reader;
75
76pub use error::StorageError;
77pub use json_payload_decoder::{
78 JsonPayloadDecoder, JsonPayloadDecoderRegistry, JsonStringDecoder, JsonVecF64Decoder,
79};
80pub use stored_state_series_reader::StoredStateSeriesReader;
81
82use json_state_record_encoder::JsonStateRecordEncoder;
83use jsonl_format::{
84 RecordingMetadata, RecordingStatus, StateFieldMetadata, StateStreamMetadata,
85 TimeAxisMetadata as StoredTimeAxis,
86};
87use queued_state_writer::{RecoveredStateStream, StateStreamStorageConfig, StateWriterWorker};
88
89const METADATA_FILE: &str = "metadata.json";
91
92const METADATA_TEMP_FILE: &str = ".metadata.json.tmp";
94
95#[derive(Clone, Debug, Eq, PartialEq)]
102pub struct TimeAxisMetadata {
103 iteration_name: String,
104 iteration_unit: Option<String>,
105 physical_time_name: Option<String>,
106 physical_time_unit: Option<String>,
107}
108
109impl TimeAxisMetadata {
110 pub fn new(iteration_name: impl Into<String>) -> Self {
116 Self {
117 iteration_name: iteration_name.into(),
118 iteration_unit: None,
119 physical_time_name: None,
120 physical_time_unit: None,
121 }
122 }
123
124 #[must_use]
126 pub fn with_iteration_unit(mut self, unit: impl Into<String>) -> Self {
127 self.iteration_unit = Some(unit.into());
128 self
129 }
130
131 #[must_use]
133 pub fn with_physical_time_name(mut self, name: impl Into<String>) -> Self {
134 self.physical_time_name = Some(name.into());
135 self
136 }
137
138 #[must_use]
143 pub fn with_physical_time_unit(mut self, unit: impl Into<String>) -> Self {
144 self.physical_time_unit = Some(unit.into());
145 self
146 }
147
148 #[must_use]
150 pub fn with_physical_axis(mut self, name: impl Into<String>, unit: impl Into<String>) -> Self {
151 self.physical_time_name = Some(name.into());
152 self.physical_time_unit = Some(unit.into());
153 self
154 }
155
156 fn into_stored(self) -> StoredTimeAxis {
158 StoredTimeAxis {
159 iteration_name: self.iteration_name,
160 iteration_unit: self.iteration_unit,
161 physical_time_name: self.physical_time_name,
162 physical_time_unit: self.physical_time_unit,
163 }
164 }
165}
166
167impl Default for TimeAxisMetadata {
168 fn default() -> Self {
171 Self::new("iteration")
172 }
173}
174
175#[derive(Clone, Debug, Eq, PartialEq)]
181pub struct RecordingTiming {
182 created_at_utc: String,
183 finalized_at_utc: String,
184 active_duration_ns: u64,
185 continuation_count: u64,
186}
187
188impl RecordingTiming {
189 fn from_stored(
191 timing: &jsonl_format::RecordingTiming,
192 metadata_path: &Path,
193 ) -> Result<Self, StorageError> {
194 let finalized_at_utc =
195 timing
196 .finalized_at_utc
197 .clone()
198 .ok_or_else(|| StorageError::InvalidMetadata {
199 path: metadata_path.to_path_buf(),
200 reason: "completed recording lacks finalized timestamp".to_owned(),
201 })?;
202 Ok(Self {
203 created_at_utc: timing.created_at_utc.clone(),
204 finalized_at_utc,
205 active_duration_ns: timing.active_duration_ns,
206 continuation_count: timing.continuation_count,
207 })
208 }
209
210 pub fn created_at_utc(&self) -> &str {
212 &self.created_at_utc
213 }
214
215 pub fn finalized_at_utc(&self) -> &str {
217 &self.finalized_at_utc
218 }
219
220 pub fn active_duration_ns(&self) -> u64 {
222 self.active_duration_ns
223 }
224
225 pub fn active_duration(&self) -> Duration {
227 Duration::from_nanos(self.active_duration_ns)
228 }
229
230 pub fn continuation_count(&self) -> u64 {
232 self.continuation_count
233 }
234}
235
236#[derive(Clone, Debug, Eq, PartialEq)]
238pub struct CompletedStreamSummary {
239 name: String,
240 chunk_count: u64,
241 record_count: u64,
242 encoded_bytes: u64,
243 first_iteration: Option<u64>,
244 last_iteration: Option<u64>,
245}
246
247impl CompletedStreamSummary {
248 pub fn name(&self) -> &str {
250 &self.name
251 }
252
253 pub fn chunk_count(&self) -> u64 {
255 self.chunk_count
256 }
257
258 pub fn record_count(&self) -> u64 {
260 self.record_count
261 }
262
263 pub fn encoded_bytes(&self) -> u64 {
265 self.encoded_bytes
266 }
267
268 pub fn first_iteration(&self) -> Option<u64> {
270 self.first_iteration
271 }
272
273 pub fn last_iteration(&self) -> Option<u64> {
275 self.last_iteration
276 }
277}
278
279#[derive(Clone, Debug, Eq, PartialEq)]
284pub struct CompletedRecording {
285 directory: PathBuf,
286 timing: RecordingTiming,
287 terminal_metadata: Map<String, Value>,
288 streams: Vec<CompletedStreamSummary>,
289}
290
291impl CompletedRecording {
292 pub fn directory(&self) -> &Path {
294 &self.directory
295 }
296
297 pub fn timing(&self) -> &RecordingTiming {
299 &self.timing
300 }
301
302 pub fn terminal_metadata(&self) -> &Map<String, Value> {
304 &self.terminal_metadata
305 }
306
307 pub fn stream_summaries(&self) -> &[CompletedStreamSummary] {
309 &self.streams
310 }
311
312 pub fn stream_summary(&self, name: &str) -> Option<&CompletedStreamSummary> {
314 self.streams.iter().find(|stream| stream.name == name)
315 }
316}
317
318#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
329#[serde(rename_all = "snake_case")]
330pub enum SamplingInterval {
331 Iterations(NonZeroU64),
333}
334
335#[derive(Deserialize)]
336#[serde(untagged)]
337enum SamplingIntervalInput {
338 Iterations(NonZeroU64),
340 Tagged { iterations: NonZeroU64 },
342}
343
344impl<'de> Deserialize<'de> for SamplingInterval {
345 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
346 where
347 D: serde::Deserializer<'de>,
348 {
349 match SamplingIntervalInput::deserialize(deserializer)? {
350 SamplingIntervalInput::Iterations(interval)
351 | SamplingIntervalInput::Tagged {
352 iterations: interval,
353 } => Ok(Self::Iterations(interval)),
354 }
355 }
356}
357
358impl SamplingInterval {
359 pub const fn iterations(interval: u64) -> Option<Self> {
361 match NonZeroU64::new(interval) {
362 Some(interval) => Some(Self::Iterations(interval)),
363 None => None,
364 }
365 }
366
367 const fn includes(self, iteration: u64) -> bool {
369 match self {
370 Self::Iterations(interval) => iteration.is_multiple_of(interval.get()),
371 }
372 }
373}
374
375#[derive(Clone, Copy, Debug, Eq, PartialEq, Deserialize, Serialize)]
377#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
378pub enum StateStreamLayout {
379 Chunked { target_bytes: NonZeroU64 },
381 IndividualFiles,
383}
384
385#[derive(Clone, Copy, Debug, Eq, PartialEq, Deserialize, Serialize)]
387#[serde(deny_unknown_fields)]
388pub struct StateStreamStorage {
389 layout: StateStreamLayout,
390 queue_bytes: NonZeroU64,
391}
392
393impl StateStreamStorage {
394 pub const fn chunked(target_bytes: NonZeroU64, queue_bytes: NonZeroU64) -> Self {
396 Self {
397 layout: StateStreamLayout::Chunked { target_bytes },
398 queue_bytes,
399 }
400 }
401
402 pub const fn individual_files(queue_bytes: NonZeroU64) -> Self {
404 Self {
405 layout: StateStreamLayout::IndividualFiles,
406 queue_bytes,
407 }
408 }
409
410 pub const fn layout(self) -> StateStreamLayout {
411 self.layout
412 }
413
414 pub const fn queue_bytes(self) -> NonZeroU64 {
415 self.queue_bytes
416 }
417}
418
419#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
426#[serde(deny_unknown_fields)]
427pub struct StateStreamConfig {
428 name: String,
429 #[serde(default, skip_serializing_if = "Option::is_none")]
430 directory: Option<String>,
431 sampling_interval: SamplingInterval,
432 fields: Vec<String>,
433 #[serde(default, skip_serializing_if = "Option::is_none")]
434 storage: Option<StateStreamStorage>,
435}
436
437impl StateStreamConfig {
438 pub fn new<I, K>(
446 name: impl Into<String>,
447 fields: I,
448 sampling_interval: SamplingInterval,
449 storage: Option<StateStreamStorage>,
450 ) -> Self
451 where
452 I: IntoIterator<Item = K>,
453 K: Into<String>,
454 {
455 let name = name.into();
456 Self {
457 directory: None,
458 name,
459 sampling_interval,
460 fields: fields.into_iter().map(Into::into).collect(),
461 storage,
462 }
463 }
464
465 #[must_use]
470 pub fn with_relative_directory(mut self, directory: impl Into<String>) -> Self {
471 self.directory = Some(directory.into());
472 self
473 }
474
475 pub fn name(&self) -> &str {
477 &self.name
478 }
479
480 pub fn relative_directory(&self) -> &str {
482 self.directory.as_deref().unwrap_or(&self.name)
483 }
484
485 pub const fn sampling_interval(&self) -> SamplingInterval {
487 self.sampling_interval
488 }
489
490 pub fn fields(&self) -> &[String] {
492 &self.fields
493 }
494
495 pub const fn storage(&self) -> Option<StateStreamStorage> {
497 self.storage
498 }
499}
500
501#[derive(Debug)]
508pub struct SystemStateWriterBuilder {
509 root: PathBuf,
510 spec: SystemStateSchema,
511 time: TimeAxisMetadata,
512 user_metadata: Map<String, Value>,
513 shared_stream_storage: Option<StateStreamStorage>,
514 streams: Vec<StateStreamConfig>,
515}
516
517impl SystemStateWriterBuilder {
518 pub fn new<S>(root: impl Into<PathBuf>, source: &S) -> Self
524 where
525 S: StateSchemaSource + ?Sized,
526 {
527 Self {
528 root: root.into(),
529 spec: source.state_schema().clone(),
530 time: TimeAxisMetadata::default(),
531 user_metadata: Map::new(),
532 shared_stream_storage: None,
533 streams: Vec::new(),
534 }
535 }
536
537 #[must_use]
539 pub fn with_time_axis_metadata(mut self, time: TimeAxisMetadata) -> Self {
540 self.time = time;
541 self
542 }
543
544 #[must_use]
550 pub fn with_user_metadata(mut self, metadata: Map<String, Value>) -> Self {
551 self.user_metadata.extend(metadata);
552 self
553 }
554
555 #[must_use]
561 pub fn with_shared_stream_storage(mut self, storage: StateStreamStorage) -> Self {
562 self.shared_stream_storage = Some(storage);
563 self
564 }
565
566 #[must_use]
574 pub fn with_task_parameters(mut self, parameters: &TaskParameters) -> Self {
575 self.user_metadata.extend(
576 parameters
577 .iter()
578 .map(|(key, value)| (key.to_owned(), value.clone())),
579 );
580 self.user_metadata.insert(
581 "task_ordinal".to_owned(),
582 Value::from(parameters.task_ordinal()),
583 );
584 self
585 }
586
587 #[must_use]
592 pub fn add_state_stream(mut self, stream: StateStreamConfig) -> Self {
593 self.streams.push(stream);
594 self
595 }
596
597 pub fn create_new_recording(self) -> Result<SystemStateWriter, StorageError> {
609 SystemStateWriter::create_new_recording(self)
610 }
611
612 pub fn continue_existing_recording(self) -> Result<SystemStateWriter, StorageError> {
620 SystemStateWriter::continue_recording(self, None).map(|(writer, _)| writer)
621 }
622
623 pub fn continue_recording_from_latest_checkpoint(
631 self,
632 stream: &str,
633 decoders: JsonPayloadDecoderRegistry,
634 ) -> Result<(SystemStateWriter, SystemState), StorageError> {
635 let (writer, state) =
636 SystemStateWriter::continue_recording(self, Some((stream, decoders)))?;
637 Ok((
638 writer,
639 state.expect("checkpoint-aware resume always reconstructs one state"),
640 ))
641 }
642}
643
644pub struct SystemStateWriter {
651 root: PathBuf,
652 stream_order: Vec<String>,
653 manifest: Arc<RecordingManifest>,
654 streams: HashMap<String, ScheduledStateStream>,
655 writer: Option<StateWriterWorker>,
656 session_started: Instant,
657 _lease: RecordingLease,
660}
661
662impl SystemStateWriter {
663 pub fn builder<S>(root: impl Into<PathBuf>, source: &S) -> SystemStateWriterBuilder
665 where
666 S: StateSchemaSource + ?Sized,
667 {
668 SystemStateWriterBuilder::new(root, source)
669 }
670
671 pub fn recording_directory(&self) -> &Path {
673 &self.root
674 }
675
676 pub fn stream_names(&self) -> impl ExactSizeIterator<Item = &str> {
678 self.stream_order.iter().map(String::as_str)
679 }
680
681 pub fn observe_state(&mut self, state: &SystemState) -> Result<(), StorageError> {
695 let iteration = state.simulation_time().iteration();
696 let writer = self
697 .writer
698 .as_ref()
699 .expect("an active recording owns its writer worker");
700 for name in &self.stream_order {
701 let stream = self
702 .streams
703 .get_mut(name)
704 .expect("stream order contains every configured stream");
705 if !stream.sampling_interval.includes(iteration)
706 || stream.last_recorded_iteration == Some(iteration)
707 {
708 continue;
709 }
710 let record = stream.encoder.encode(state)?;
711 writer.submit_record(name, record)?;
712 stream.last_recorded_iteration = Some(iteration);
713 }
714 Ok(())
715 }
716
717 pub fn flush_stream_to_storage(&self, stream: &str) -> Result<(), StorageError> {
724 if !self.streams.contains_key(stream) {
725 return Err(StorageError::UnknownStateStream {
726 stream: stream.to_owned(),
727 });
728 }
729 self.writer
730 .as_ref()
731 .expect("an active recording owns its writer worker")
732 .flush_state_stream(stream)
733 }
734
735 pub fn complete_recording(self) -> Result<CompletedRecording, StorageError> {
743 self.complete_recording_with_terminal_metadata(Map::new())
744 }
745
746 pub fn complete_recording_with_terminal_metadata(
752 mut self,
753 terminal_metadata: Map<String, Value>,
754 ) -> Result<CompletedRecording, StorageError> {
755 if let Err(error) = self.finish_writer() {
756 let _ = self.transition_terminal(
757 RecordingStatus::Failed {
758 message: error.to_string(),
759 },
760 Map::new(),
761 );
762 return Err(error);
763 }
764 self.transition_terminal(RecordingStatus::Complete, terminal_metadata)?;
765 self.completed_recording()
766 }
767
768 pub fn complete_recording_with_final_state(
775 mut self,
776 state: &SystemState,
777 ) -> Result<CompletedRecording, StorageError> {
778 self.record_final_state(state)?;
779 self.complete_recording()
780 }
781
782 pub fn complete_recording_with_final_state_and_terminal_metadata(
785 mut self,
786 state: &SystemState,
787 terminal_metadata: Map<String, Value>,
788 ) -> Result<CompletedRecording, StorageError> {
789 self.record_final_state(state)?;
790 self.complete_recording_with_terminal_metadata(terminal_metadata)
791 }
792
793 fn record_final_state(&mut self, state: &SystemState) -> Result<(), StorageError> {
795 let iteration = state.simulation_time().iteration();
796 let writer = self
797 .writer
798 .as_ref()
799 .expect("an active recording owns its writer worker");
800 for name in &self.stream_order {
801 let stream = self
802 .streams
803 .get_mut(name)
804 .expect("stream order contains every configured stream");
805 if stream.last_recorded_iteration == Some(iteration) {
806 continue;
807 }
808 let record = stream.encoder.encode(state)?;
809 writer.submit_record(name, record)?;
810 stream.last_recorded_iteration = Some(iteration);
811 }
812 Ok(())
813 }
814
815 pub fn mark_recording_failed(self, message: impl Into<String>) -> Result<(), StorageError> {
827 self.mark_recording_failed_with_terminal_metadata(message, Map::new())
828 }
829
830 pub fn mark_recording_failed_with_terminal_metadata(
832 mut self,
833 message: impl Into<String>,
834 terminal_metadata: Map<String, Value>,
835 ) -> Result<(), StorageError> {
836 let message = message.into();
837 if message.trim().is_empty() {
838 return Err(StorageError::InvalidConfiguration {
839 setting: "failure_message",
840 reason: "failed run message must not be empty".to_owned(),
841 });
842 }
843
844 if let Err(error) = self.finish_writer() {
845 let _ = self.transition_terminal(
846 RecordingStatus::Failed {
847 message: error.to_string(),
848 },
849 Map::new(),
850 );
851 return Err(error);
852 }
853 self.transition_terminal(RecordingStatus::Failed { message }, terminal_metadata)
854 }
855
856 fn create_new_recording(builder: SystemStateWriterBuilder) -> Result<Self, StorageError> {
858 ensure_absent(&builder.root)?;
859 let prepared = PreparedRecording::from_builder(builder)?;
860 create_root(&prepared.root)?;
861 let lease = RecordingLease::acquire(&prepared.root)?;
862 for stream in &prepared.streams {
863 stream.writer.create_directory()?;
864 }
865 commit_metadata(&prepared.root, &prepared.metadata_path, &prepared.metadata)?;
866 let manifest = Arc::new(RecordingManifest::new(
867 prepared.root.clone(),
868 prepared.metadata_path.clone(),
869 prepared.metadata,
870 ));
871 Self::start_new_prepared(prepared.root, prepared.streams, manifest, lease)
872 }
873
874 fn continue_recording(
876 builder: SystemStateWriterBuilder,
877 checkpoint: Option<(&str, JsonPayloadDecoderRegistry)>,
878 ) -> Result<(Self, Option<SystemState>), StorageError> {
879 let prepared = PreparedRecording::from_builder(builder)?;
880 let lease = RecordingLease::acquire(&prepared.root)?;
881 remove_stale_metadata_temp(&prepared.root)?;
882 let mut existing = load_metadata(&prepared.metadata_path)?;
883 if !matches!(existing.status, RecordingStatus::Running) {
884 return Err(StorageError::RecordingNotContinuable {
885 path: prepared.metadata_path,
886 });
887 }
888 ensure_resume_match(&prepared.metadata_path, &prepared.metadata, &existing)?;
889
890 let mut recovered = Vec::with_capacity(prepared.streams.len());
891 for stream in prepared.streams {
892 let declaration = existing
893 .stream(&stream.name)
894 .expect("matched metadata contains every prepared stream");
895 let seed = StateWriterWorker::recover_state_stream(&stream.writer, declaration)?;
896 recovered.push((stream, seed));
897 }
898
899 let state = if let Some((checkpoint_stream, decoders)) = checkpoint {
900 let declaration = existing.stream(checkpoint_stream).ok_or_else(|| {
901 StorageError::UnknownStateStream {
902 stream: checkpoint_stream.to_owned(),
903 }
904 })?;
905 Some(stored_state_series_reader::decode_resume_state(
906 &prepared.root,
907 &prepared.metadata_path,
908 declaration,
909 &prepared.spec,
910 &decoders,
911 )?)
912 } else {
913 None
914 };
915
916 existing.timing.continuation_count = existing
917 .timing
918 .continuation_count
919 .checked_add(1)
920 .ok_or_else(|| StorageError::InvalidMetadata {
921 path: prepared.metadata_path.clone(),
922 reason: "timing.continuation_count overflowed".to_owned(),
923 })?;
924 commit_metadata(&prepared.root, &prepared.metadata_path, &existing)?;
925 let manifest = Arc::new(RecordingManifest::new(
926 prepared.root.clone(),
927 prepared.metadata_path.clone(),
928 existing,
929 ));
930
931 let output = Self::start_resumed_prepared(prepared.root, recovered, manifest, lease)?;
932 Ok((output, state))
933 }
934
935 fn start_new_prepared(
937 root: PathBuf,
938 streams: Vec<PreparedStateStream>,
939 manifest: Arc<RecordingManifest>,
940 lease: RecordingLease,
941 ) -> Result<Self, StorageError> {
942 let mut scheduled = HashMap::with_capacity(streams.len());
943 let mut configs = Vec::with_capacity(streams.len());
944 let mut stream_order = Vec::with_capacity(streams.len());
945 for prepared in streams {
946 let name = prepared.name;
947 stream_order.push(name.clone());
948 scheduled.insert(
949 name,
950 ScheduledStateStream {
951 encoder: prepared.encoder,
952 sampling_interval: prepared.sampling_interval,
953 last_recorded_iteration: None,
954 },
955 );
956 configs.push(prepared.writer);
957 }
958 let writer = StateWriterWorker::start_new_recording(configs, Arc::clone(&manifest))?;
959 Ok(Self {
960 root,
961 stream_order,
962 manifest,
963 streams: scheduled,
964 writer: Some(writer),
965 session_started: Instant::now(),
966 _lease: lease,
967 })
968 }
969
970 fn start_resumed_prepared(
972 root: PathBuf,
973 streams: Vec<(PreparedStateStream, RecoveredStateStream)>,
974 manifest: Arc<RecordingManifest>,
975 lease: RecordingLease,
976 ) -> Result<Self, StorageError> {
977 let mut scheduled = HashMap::with_capacity(streams.len());
978 let mut recovered_streams = Vec::with_capacity(streams.len());
979 let mut stream_order = Vec::with_capacity(streams.len());
980 for (prepared, seed) in streams {
981 let name = prepared.name;
982 stream_order.push(name.clone());
983 scheduled.insert(
984 name,
985 ScheduledStateStream {
986 encoder: prepared.encoder,
987 sampling_interval: prepared.sampling_interval,
988 last_recorded_iteration: seed.last_iteration(),
989 },
990 );
991 recovered_streams.push((prepared.writer, seed));
992 }
993 let writer = StateWriterWorker::continue_recovered_recording(
994 recovered_streams,
995 Arc::clone(&manifest),
996 )?;
997 Ok(Self {
998 root,
999 stream_order,
1000 manifest,
1001 streams: scheduled,
1002 writer: Some(writer),
1003 session_started: Instant::now(),
1004 _lease: lease,
1005 })
1006 }
1007
1008 fn finish_writer(&mut self) -> Result<(), StorageError> {
1010 let Some(writer) = self.writer.take() else {
1011 return Ok(());
1012 };
1013 writer.finish_recording()
1014 }
1015
1016 fn transition_terminal(
1018 &self,
1019 status: RecordingStatus,
1020 terminal_metadata: Map<String, Value>,
1021 ) -> Result<(), StorageError> {
1022 let finalized_at_utc =
1023 utc_now_rfc3339().map_err(|source| StorageError::OperationalTimestamp {
1024 operation: "finalize recording",
1025 source,
1026 })?;
1027 let active_duration_ns = duration_nanoseconds(self.session_started.elapsed())
1028 .ok_or(StorageError::OperationalDurationOverflow)?;
1029 self.manifest.transition_terminal(
1030 status,
1031 finalized_at_utc,
1032 active_duration_ns,
1033 terminal_metadata,
1034 )
1035 }
1036
1037 fn completed_recording(&self) -> Result<CompletedRecording, StorageError> {
1039 let metadata = self.manifest.snapshot();
1040 let timing = RecordingTiming::from_stored(&metadata.timing, &self.manifest.path)?;
1041 let streams = metadata
1042 .streams
1043 .iter()
1044 .map(completed_stream_summary)
1045 .collect::<Result<Vec<_>, _>>()?;
1046 Ok(CompletedRecording {
1047 directory: self.root.clone(),
1048 timing,
1049 terminal_metadata: metadata.terminal_metadata,
1050 streams,
1051 })
1052 }
1053}
1054
1055fn completed_stream_summary(
1057 stream: &StateStreamMetadata,
1058) -> Result<CompletedStreamSummary, StorageError> {
1059 let overflow = || StorageError::ByteCountOverflow {
1060 stream: stream.name.clone(),
1061 };
1062 let chunk_count = u64::try_from(stream.chunks.len()).map_err(|_| overflow())?;
1063 let record_count = stream
1064 .chunks
1065 .iter()
1066 .try_fold(0_u64, |total, chunk| total.checked_add(chunk.records))
1067 .ok_or_else(&overflow)?;
1068 let encoded_bytes = stream
1069 .chunks
1070 .iter()
1071 .try_fold(0_u64, |total, chunk| total.checked_add(chunk.bytes))
1072 .ok_or_else(overflow)?;
1073 Ok(CompletedStreamSummary {
1074 name: stream.name.clone(),
1075 chunk_count,
1076 record_count,
1077 encoded_bytes,
1078 first_iteration: stream.chunks.first().map(|chunk| chunk.first_iteration),
1079 last_iteration: stream.chunks.last().map(|chunk| chunk.last_iteration),
1080 })
1081}
1082
1083struct PreparedRecording {
1085 root: PathBuf,
1086 metadata_path: PathBuf,
1087 spec: SystemStateSchema,
1088 metadata: RecordingMetadata,
1089 streams: Vec<PreparedStateStream>,
1090}
1091
1092impl PreparedRecording {
1093 fn from_builder(builder: SystemStateWriterBuilder) -> Result<Self, StorageError> {
1095 let metadata_path = builder.root.join(METADATA_FILE);
1096 let stored_time = builder.time.into_stored();
1097 let mut names = HashSet::with_capacity(builder.streams.len());
1098 let mut directories = HashSet::with_capacity(builder.streams.len());
1099 let mut streams = Vec::with_capacity(builder.streams.len());
1100 let mut declarations = Vec::with_capacity(builder.streams.len());
1101
1102 for config in builder.streams {
1103 if !names.insert(config.name.clone()) {
1104 return Err(StorageError::DuplicateStateStream {
1105 stream: config.name,
1106 });
1107 }
1108 let directory = config.relative_directory().to_owned();
1109 if !directories.insert(directory.clone()) {
1110 return Err(StorageError::InvalidConfiguration {
1111 setting: "stream.directory",
1112 reason: format!("multiple streams use relative directory `{}`", directory),
1113 });
1114 }
1115
1116 let storage = config
1117 .storage
1118 .or(builder.shared_stream_storage)
1119 .ok_or_else(|| StorageError::InvalidConfiguration {
1120 setting: "stream.storage",
1121 reason: format!(
1122 "stream `{}` has no explicit storage and the writer has no shared storage",
1123 config.name
1124 ),
1125 })?;
1126 let encoder = JsonStateRecordEncoder::new(&config.name, &builder.spec, &config.fields)?;
1127 let fields = encoder
1128 .fields()
1129 .map(|name| {
1130 let field = builder
1131 .spec
1132 .field_schema(name)
1133 .expect("encoder fields were validated against this specification");
1134 StateFieldMetadata {
1135 name: name.to_owned(),
1136 description: field.description().map(str::to_owned),
1137 }
1138 })
1139 .collect::<Vec<_>>();
1140 declarations.push(StateStreamMetadata {
1141 name: config.name.clone(),
1142 directory: directory.clone(),
1143 sampling_interval: config.sampling_interval,
1144 fields,
1145 storage,
1146 chunks: Vec::new(),
1147 });
1148 streams.push(PreparedStateStream {
1149 name: config.name.clone(),
1150 encoder,
1151 sampling_interval: config.sampling_interval,
1152 writer: StateStreamStorageConfig::new(
1153 &config.name,
1154 builder.root.join(&directory),
1155 storage,
1156 )?,
1157 });
1158 }
1159
1160 let created_at_utc =
1161 utc_now_rfc3339().map_err(|source| StorageError::OperationalTimestamp {
1162 operation: "create recording",
1163 source,
1164 })?;
1165 let metadata = RecordingMetadata::running(
1166 stored_time,
1167 builder.user_metadata,
1168 declarations,
1169 created_at_utc,
1170 );
1171 metadata.validate(&metadata_path)?;
1172 Ok(Self {
1173 root: builder.root,
1174 metadata_path,
1175 spec: builder.spec,
1176 metadata,
1177 streams,
1178 })
1179 }
1180}
1181
1182struct PreparedStateStream {
1184 name: String,
1185 encoder: JsonStateRecordEncoder,
1186 sampling_interval: SamplingInterval,
1187 writer: StateStreamStorageConfig,
1188}
1189
1190struct ScheduledStateStream {
1192 encoder: JsonStateRecordEncoder,
1193 sampling_interval: SamplingInterval,
1194 last_recorded_iteration: Option<u64>,
1195}
1196
1197pub(crate) struct RecordingManifest {
1203 root: PathBuf,
1204 path: PathBuf,
1205 metadata: Mutex<RecordingMetadata>,
1206}
1207
1208impl RecordingManifest {
1209 fn new(root: PathBuf, path: PathBuf, metadata: RecordingMetadata) -> Self {
1211 Self {
1212 root,
1213 path,
1214 metadata: Mutex::new(metadata),
1215 }
1216 }
1217
1218 pub(crate) fn prepare_chunk(
1220 &self,
1221 stream: &str,
1222 descriptor: jsonl_format::ChunkMetadata,
1223 ) -> Result<(), StorageError> {
1224 let mut current = lock_metadata(&self.metadata);
1225 if !matches!(current.status, RecordingStatus::Running) {
1226 return Err(StorageError::RecordingFinished);
1227 }
1228 let mut candidate = current.clone();
1229 let declaration =
1230 candidate
1231 .stream_mut(stream)
1232 .ok_or_else(|| StorageError::UnknownStateStream {
1233 stream: stream.to_owned(),
1234 })?;
1235 let expected = u64::try_from(declaration.chunks.len()).map_err(|_| {
1236 StorageError::ByteCountOverflow {
1237 stream: stream.to_owned(),
1238 }
1239 })?;
1240 if descriptor.ordinal != expected {
1241 return Err(StorageError::InvalidMetadata {
1242 path: self.path.clone(),
1243 reason: format!(
1244 "stream `{stream}` prepared chunk ordinal {}, expected {expected}",
1245 descriptor.ordinal
1246 ),
1247 });
1248 }
1249 declaration.chunks.push(descriptor);
1250 commit_metadata(&self.root, &self.path, &candidate)?;
1251 *current = candidate;
1252 Ok(())
1253 }
1254
1255 fn transition_terminal(
1257 &self,
1258 status: RecordingStatus,
1259 finalized_at_utc: String,
1260 active_duration_ns: u64,
1261 terminal_metadata: Map<String, Value>,
1262 ) -> Result<(), StorageError> {
1263 let mut current = lock_metadata(&self.metadata);
1264 let mut candidate = current.clone();
1265 candidate.status = status;
1266 candidate.timing.finalized_at_utc = Some(finalized_at_utc);
1267 candidate.timing.active_duration_ns = candidate
1268 .timing
1269 .active_duration_ns
1270 .checked_add(active_duration_ns)
1271 .ok_or(StorageError::OperationalDurationOverflow)?;
1272 candidate.terminal_metadata = terminal_metadata;
1273 commit_metadata(&self.root, &self.path, &candidate)?;
1274 *current = candidate;
1275 Ok(())
1276 }
1277
1278 fn snapshot(&self) -> RecordingMetadata {
1280 lock_metadata(&self.metadata).clone()
1281 }
1282}
1283
1284struct RecordingLease {
1289 _directory: File,
1290}
1291
1292impl RecordingLease {
1293 fn acquire(root: &Path) -> Result<Self, StorageError> {
1295 let directory = File::open(root).map_err(|source| StorageError::Io {
1296 operation: "open output root for exclusive ownership",
1297 path: root.to_path_buf(),
1298 source,
1299 })?;
1300 match FileExt::try_lock_exclusive(&directory) {
1301 Ok(()) => Ok(Self {
1302 _directory: directory,
1303 }),
1304 Err(source) if source.kind() == std::io::ErrorKind::WouldBlock => {
1305 Err(StorageError::RecordingDirectoryInUse {
1306 path: root.to_path_buf(),
1307 })
1308 }
1309 Err(source) => Err(StorageError::Io {
1310 operation: "acquire exclusive output ownership",
1311 path: root.to_path_buf(),
1312 source,
1313 }),
1314 }
1315 }
1316}
1317
1318fn load_metadata(path: &Path) -> Result<RecordingMetadata, StorageError> {
1320 let bytes = fs::read(path).map_err(|source| StorageError::Io {
1321 operation: "read metadata for resume",
1322 path: path.to_path_buf(),
1323 source,
1324 })?;
1325 let metadata: RecordingMetadata =
1326 serde_json::from_slice(&bytes).map_err(|source| StorageError::Json {
1327 operation: "parse metadata for resume",
1328 path: path.to_path_buf(),
1329 source,
1330 })?;
1331 metadata.validate(path)?;
1332 Ok(metadata)
1333}
1334
1335fn ensure_resume_match(
1337 path: &Path,
1338 expected: &RecordingMetadata,
1339 existing: &RecordingMetadata,
1340) -> Result<(), StorageError> {
1341 let mut configuration = existing.clone();
1342 for stream in &mut configuration.streams {
1343 stream.chunks.clear();
1344 }
1345 configuration.status = RecordingStatus::Running;
1346 configuration.timing = expected.timing.clone();
1347 configuration.terminal_metadata.clear();
1348 if &configuration != expected {
1349 return Err(StorageError::RecordingConfigurationMismatch {
1350 path: path.to_path_buf(),
1351 reason: "builder time axis, user metadata, or stream declarations differ".to_owned(),
1352 });
1353 }
1354 Ok(())
1355}
1356
1357fn remove_stale_metadata_temp(root: &Path) -> Result<(), StorageError> {
1359 let path = root.join(METADATA_TEMP_FILE);
1360 match fs::remove_file(&path) {
1361 Ok(()) => File::open(root)
1362 .and_then(|directory| directory.sync_all())
1363 .map_err(|source| StorageError::Io {
1364 operation: "synchronize stale metadata cleanup",
1365 path: root.to_path_buf(),
1366 source,
1367 }),
1368 Err(source) if source.kind() == std::io::ErrorKind::NotFound => Ok(()),
1369 Err(source) => Err(StorageError::Io {
1370 operation: "remove stale temporary metadata",
1371 path,
1372 source,
1373 }),
1374 }
1375}
1376
1377fn lock_metadata(metadata: &Mutex<RecordingMetadata>) -> MutexGuard<'_, RecordingMetadata> {
1379 metadata
1380 .lock()
1381 .unwrap_or_else(|poisoned| poisoned.into_inner())
1382}
1383
1384fn ensure_absent(root: &Path) -> Result<(), StorageError> {
1386 match root.try_exists() {
1387 Ok(false) => Ok(()),
1388 Ok(true) => Err(StorageError::RecordingDirectoryExists {
1389 path: root.to_path_buf(),
1390 }),
1391 Err(source) => Err(StorageError::Io {
1392 operation: "inspect output root",
1393 path: root.to_path_buf(),
1394 source,
1395 }),
1396 }
1397}
1398
1399fn create_root(root: &Path) -> Result<(), StorageError> {
1401 if let Some(parent) = root.parent() {
1402 fs::create_dir_all(parent).map_err(|source| StorageError::Io {
1403 operation: "create recording parent directories",
1404 path: parent.to_path_buf(),
1405 source,
1406 })?;
1407 }
1408 match fs::create_dir(root) {
1409 Ok(()) => Ok(()),
1410 Err(source) if source.kind() == std::io::ErrorKind::AlreadyExists => {
1411 Err(StorageError::RecordingDirectoryExists {
1412 path: root.to_path_buf(),
1413 })
1414 }
1415 Err(source) => Err(StorageError::Io {
1416 operation: "create output root",
1417 path: root.to_path_buf(),
1418 source,
1419 }),
1420 }
1421}
1422
1423fn commit_metadata(
1431 root: &Path,
1432 metadata_path: &Path,
1433 metadata: &RecordingMetadata,
1434) -> Result<(), StorageError> {
1435 metadata.validate(metadata_path)?;
1436 let mut bytes = serde_json::to_vec_pretty(metadata).map_err(|source| StorageError::Json {
1437 operation: "serialize metadata",
1438 path: metadata_path.to_path_buf(),
1439 source,
1440 })?;
1441 bytes.push(b'\n');
1442
1443 let temporary_path = root.join(METADATA_TEMP_FILE);
1444 let result = write_and_replace_metadata(root, metadata_path, &temporary_path, &bytes);
1445 if result.is_err() {
1446 let _ = fs::remove_file(&temporary_path);
1447 }
1448 result
1449}
1450
1451fn write_and_replace_metadata(
1453 root: &Path,
1454 metadata_path: &Path,
1455 temporary_path: &Path,
1456 bytes: &[u8],
1457) -> Result<(), StorageError> {
1458 let mut temporary = OpenOptions::new()
1459 .write(true)
1460 .create_new(true)
1461 .open(temporary_path)
1462 .map_err(|source| StorageError::Io {
1463 operation: "create temporary metadata",
1464 path: temporary_path.to_path_buf(),
1465 source,
1466 })?;
1467 temporary
1468 .write_all(bytes)
1469 .map_err(|source| StorageError::Io {
1470 operation: "write temporary metadata",
1471 path: temporary_path.to_path_buf(),
1472 source,
1473 })?;
1474 temporary.sync_all().map_err(|source| StorageError::Io {
1475 operation: "sync temporary metadata",
1476 path: temporary_path.to_path_buf(),
1477 source,
1478 })?;
1479 drop(temporary);
1480
1481 fs::rename(temporary_path, metadata_path).map_err(|source| StorageError::Io {
1482 operation: "publish metadata",
1483 path: metadata_path.to_path_buf(),
1484 source,
1485 })?;
1486
1487 File::open(root)
1488 .and_then(|directory| directory.sync_all())
1489 .map_err(|source| StorageError::Io {
1490 operation: "sync output root",
1491 path: root.to_path_buf(),
1492 source,
1493 })
1494}