1use std::collections::{HashMap, HashSet};
51use std::fs::{self, File, OpenOptions};
52use std::io::Write;
53use std::num::NonZeroU64;
54use std::path::{Path, PathBuf};
55use std::sync::{Arc, Mutex, MutexGuard};
56use std::time::{Duration, Instant};
57
58use fs2::FileExt;
59use serde::{Deserialize, Serialize};
60use serde_json::{Map, Value};
61
62use crate::clock::{duration_nanoseconds, utc_now_rfc3339};
63use crate::configuration::TaskParameters;
64use crate::system_state::{SystemState, SystemStateSchema};
65
66mod error;
67mod json_payload_decoder;
68mod json_state_record_encoder;
69mod jsonl_format;
70mod queued_state_writer;
71mod stored_state_series_reader;
72
73pub use error::StorageError;
74pub use json_payload_decoder::{
75 JsonPayloadDecoder, JsonPayloadDecoderRegistry, JsonStringDecoder, JsonVecF64Decoder,
76};
77pub use stored_state_series_reader::StoredStateSeriesReader;
78
79use json_state_record_encoder::JsonStateRecordEncoder;
80use jsonl_format::{
81 RecordingMetadata, RecordingStatus, StateFieldMetadata, StateStreamMetadata,
82 TimeAxisMetadata as StoredTimeAxis,
83};
84use queued_state_writer::{RecoveredStateStream, StateStreamStorageConfig, StateWriterWorker};
85
86const METADATA_FILE: &str = "metadata.json";
88
89const METADATA_TEMP_FILE: &str = ".metadata.json.tmp";
91
92#[derive(Clone, Debug, Eq, PartialEq)]
99pub struct TimeAxisMetadata {
100 iteration_name: String,
101 iteration_unit: Option<String>,
102 physical_time_name: Option<String>,
103 physical_time_unit: Option<String>,
104}
105
106impl TimeAxisMetadata {
107 pub fn new(iteration_name: impl Into<String>) -> Self {
113 Self {
114 iteration_name: iteration_name.into(),
115 iteration_unit: None,
116 physical_time_name: None,
117 physical_time_unit: None,
118 }
119 }
120
121 #[must_use]
123 pub fn with_iteration_unit(mut self, unit: impl Into<String>) -> Self {
124 self.iteration_unit = Some(unit.into());
125 self
126 }
127
128 #[must_use]
130 pub fn with_physical_time_name(mut self, name: impl Into<String>) -> Self {
131 self.physical_time_name = Some(name.into());
132 self
133 }
134
135 #[must_use]
140 pub fn with_physical_time_unit(mut self, unit: impl Into<String>) -> Self {
141 self.physical_time_unit = Some(unit.into());
142 self
143 }
144
145 #[must_use]
147 pub fn with_physical_axis(mut self, name: impl Into<String>, unit: impl Into<String>) -> Self {
148 self.physical_time_name = Some(name.into());
149 self.physical_time_unit = Some(unit.into());
150 self
151 }
152
153 fn into_stored(self) -> StoredTimeAxis {
155 StoredTimeAxis {
156 iteration_name: self.iteration_name,
157 iteration_unit: self.iteration_unit,
158 physical_time_name: self.physical_time_name,
159 physical_time_unit: self.physical_time_unit,
160 }
161 }
162}
163
164impl Default for TimeAxisMetadata {
165 fn default() -> Self {
168 Self::new("iteration")
169 }
170}
171
172#[derive(Clone, Debug, Eq, PartialEq)]
178pub struct RecordingTiming {
179 created_at_utc: String,
180 finalized_at_utc: String,
181 active_duration_ns: u64,
182 continuation_count: u64,
183}
184
185impl RecordingTiming {
186 fn from_stored(
188 timing: &jsonl_format::RecordingTiming,
189 metadata_path: &Path,
190 ) -> Result<Self, StorageError> {
191 let finalized_at_utc =
192 timing
193 .finalized_at_utc
194 .clone()
195 .ok_or_else(|| StorageError::InvalidMetadata {
196 path: metadata_path.to_path_buf(),
197 reason: "completed recording lacks finalized timestamp".to_owned(),
198 })?;
199 Ok(Self {
200 created_at_utc: timing.created_at_utc.clone(),
201 finalized_at_utc,
202 active_duration_ns: timing.active_duration_ns,
203 continuation_count: timing.continuation_count,
204 })
205 }
206
207 pub fn created_at_utc(&self) -> &str {
209 &self.created_at_utc
210 }
211
212 pub fn finalized_at_utc(&self) -> &str {
214 &self.finalized_at_utc
215 }
216
217 pub fn active_duration_ns(&self) -> u64 {
219 self.active_duration_ns
220 }
221
222 pub fn active_duration(&self) -> Duration {
224 Duration::from_nanos(self.active_duration_ns)
225 }
226
227 pub fn continuation_count(&self) -> u64 {
229 self.continuation_count
230 }
231}
232
233#[derive(Clone, Debug, Eq, PartialEq)]
235pub struct CompletedStreamSummary {
236 name: String,
237 chunk_count: u64,
238 record_count: u64,
239 encoded_bytes: u64,
240 first_iteration: Option<u64>,
241 last_iteration: Option<u64>,
242}
243
244impl CompletedStreamSummary {
245 pub fn name(&self) -> &str {
247 &self.name
248 }
249
250 pub fn chunk_count(&self) -> u64 {
252 self.chunk_count
253 }
254
255 pub fn record_count(&self) -> u64 {
257 self.record_count
258 }
259
260 pub fn encoded_bytes(&self) -> u64 {
262 self.encoded_bytes
263 }
264
265 pub fn first_iteration(&self) -> Option<u64> {
267 self.first_iteration
268 }
269
270 pub fn last_iteration(&self) -> Option<u64> {
272 self.last_iteration
273 }
274}
275
276#[derive(Clone, Debug, Eq, PartialEq)]
281pub struct CompletedRecording {
282 directory: PathBuf,
283 timing: RecordingTiming,
284 terminal_metadata: Map<String, Value>,
285 streams: Vec<CompletedStreamSummary>,
286}
287
288impl CompletedRecording {
289 pub fn directory(&self) -> &Path {
291 &self.directory
292 }
293
294 pub fn timing(&self) -> &RecordingTiming {
296 &self.timing
297 }
298
299 pub fn terminal_metadata(&self) -> &Map<String, Value> {
301 &self.terminal_metadata
302 }
303
304 pub fn stream_summaries(&self) -> &[CompletedStreamSummary] {
306 &self.streams
307 }
308
309 pub fn stream_summary(&self, name: &str) -> Option<&CompletedStreamSummary> {
311 self.streams.iter().find(|stream| stream.name == name)
312 }
313}
314
315#[derive(Clone, Copy, Debug, Eq, PartialEq, Deserialize, Serialize)]
322#[serde(rename_all = "snake_case")]
323pub enum SamplingInterval {
324 Iterations(NonZeroU64),
326}
327
328impl SamplingInterval {
329 pub const fn iterations(interval: u64) -> Option<Self> {
331 match NonZeroU64::new(interval) {
332 Some(interval) => Some(Self::Iterations(interval)),
333 None => None,
334 }
335 }
336
337 const fn includes(self, iteration: u64) -> bool {
339 match self {
340 Self::Iterations(interval) => iteration % interval.get() == 0,
341 }
342 }
343}
344
345#[derive(Clone, Debug, Eq, PartialEq)]
354pub struct StateStreamConfig {
355 name: String,
356 directory: String,
357 sampling_interval: SamplingInterval,
358 fields: Vec<String>,
359 storage_limits: Option<(NonZeroU64, NonZeroU64)>,
360}
361
362impl StateStreamConfig {
363 pub fn new<I, K>(
371 name: impl Into<String>,
372 fields: I,
373 sampling_interval: SamplingInterval,
374 max_chunk_bytes: NonZeroU64,
375 queue_bytes: NonZeroU64,
376 ) -> Self
377 where
378 I: IntoIterator<Item = K>,
379 K: Into<String>,
380 {
381 let name = name.into();
382 Self {
383 directory: name.clone(),
384 name,
385 sampling_interval,
386 fields: fields.into_iter().map(Into::into).collect(),
387 storage_limits: Some((max_chunk_bytes, queue_bytes)),
388 }
389 }
390
391 fn sampled<I, K>(
393 name: impl Into<String>,
394 fields: I,
395 sampling_interval: SamplingInterval,
396 ) -> Self
397 where
398 I: IntoIterator<Item = K>,
399 K: Into<String>,
400 {
401 let name = name.into();
402 Self {
403 directory: name.clone(),
404 name,
405 sampling_interval,
406 fields: fields.into_iter().map(Into::into).collect(),
407 storage_limits: None,
408 }
409 }
410
411 #[must_use]
416 pub fn with_relative_directory(mut self, directory: impl Into<String>) -> Self {
417 self.directory = directory.into();
418 self
419 }
420}
421
422#[derive(Debug)]
429pub struct SystemStateWriterBuilder {
430 root: PathBuf,
431 spec: SystemStateSchema,
432 time: TimeAxisMetadata,
433 user_metadata: Map<String, Value>,
434 shared_stream_limits: Option<(NonZeroU64, NonZeroU64)>,
435 streams: Vec<StateStreamConfig>,
436}
437
438impl SystemStateWriterBuilder {
439 pub fn new(root: impl Into<PathBuf>, spec: &SystemStateSchema) -> Self {
444 Self {
445 root: root.into(),
446 spec: spec.clone(),
447 time: TimeAxisMetadata::default(),
448 user_metadata: Map::new(),
449 shared_stream_limits: None,
450 streams: Vec::new(),
451 }
452 }
453
454 #[must_use]
456 pub fn with_time_axis_metadata(mut self, time: TimeAxisMetadata) -> Self {
457 self.time = time;
458 self
459 }
460
461 #[must_use]
467 pub fn with_user_metadata(mut self, metadata: Map<String, Value>) -> Self {
468 self.user_metadata = metadata;
469 self
470 }
471
472 #[must_use]
479 pub fn with_shared_stream_limits(
480 mut self,
481 max_chunk_bytes: NonZeroU64,
482 queue_bytes: NonZeroU64,
483 ) -> Self {
484 self.shared_stream_limits = Some((max_chunk_bytes, queue_bytes));
485 self
486 }
487
488 #[must_use]
494 pub fn with_task_parameters(mut self, parameters: &TaskParameters) -> Self {
495 self.user_metadata = parameters
496 .iter()
497 .map(|(key, value)| (key.to_owned(), value.clone()))
498 .collect();
499 self.user_metadata.insert(
500 "task_ordinal".to_owned(),
501 Value::from(parameters.task_ordinal()),
502 );
503 self
504 }
505
506 #[must_use]
511 pub fn add_state_stream(mut self, stream: StateStreamConfig) -> Self {
512 self.streams.push(stream);
513 self
514 }
515
516 #[must_use]
523 pub fn add_sampled_state_stream<I, K>(
524 mut self,
525 name: impl Into<String>,
526 fields: I,
527 sampling_interval: SamplingInterval,
528 ) -> Self
529 where
530 I: IntoIterator<Item = K>,
531 K: Into<String>,
532 {
533 self.streams
534 .push(StateStreamConfig::sampled(name, fields, sampling_interval));
535 self
536 }
537
538 pub fn create_new_recording(self) -> Result<SystemStateWriter, StorageError> {
550 SystemStateWriter::create_new_recording(self)
551 }
552
553 pub fn continue_existing_recording(self) -> Result<SystemStateWriter, StorageError> {
559 SystemStateWriter::continue_recording(self, None).map(|(writer, _)| writer)
560 }
561
562 pub fn continue_recording_from_latest_checkpoint(
568 self,
569 stream: &str,
570 decoders: JsonPayloadDecoderRegistry,
571 ) -> Result<(SystemStateWriter, SystemState), StorageError> {
572 let (writer, state) =
573 SystemStateWriter::continue_recording(self, Some((stream, decoders)))?;
574 Ok((
575 writer,
576 state.expect("checkpoint-aware resume always reconstructs one state"),
577 ))
578 }
579}
580
581pub struct SystemStateWriter {
588 root: PathBuf,
589 stream_order: Vec<String>,
590 manifest: Arc<RecordingManifest>,
591 streams: HashMap<String, ScheduledStateStream>,
592 writer: Option<StateWriterWorker>,
593 session_started: Instant,
594 _lease: RecordingLease,
597}
598
599impl SystemStateWriter {
600 pub fn builder(root: impl Into<PathBuf>, spec: &SystemStateSchema) -> SystemStateWriterBuilder {
602 SystemStateWriterBuilder::new(root, spec)
603 }
604
605 pub fn recording_directory(&self) -> &Path {
607 &self.root
608 }
609
610 pub fn stream_names(&self) -> impl ExactSizeIterator<Item = &str> {
612 self.stream_order.iter().map(String::as_str)
613 }
614
615 pub fn observe_state(&mut self, state: &SystemState) -> Result<(), StorageError> {
629 let iteration = state.simulation_time().iteration();
630 let writer = self
631 .writer
632 .as_ref()
633 .expect("an active recording owns its writer worker");
634 for name in &self.stream_order {
635 let stream = self
636 .streams
637 .get_mut(name)
638 .expect("stream order contains every configured stream");
639 if !stream.sampling_interval.includes(iteration)
640 || stream.last_recorded_iteration == Some(iteration)
641 {
642 continue;
643 }
644 let record = stream.encoder.encode(state)?;
645 writer.submit_record(name, record)?;
646 stream.last_recorded_iteration = Some(iteration);
647 }
648 Ok(())
649 }
650
651 pub fn flush_stream_to_storage(&self, stream: &str) -> Result<(), StorageError> {
658 if !self.streams.contains_key(stream) {
659 return Err(StorageError::UnknownStateStream {
660 stream: stream.to_owned(),
661 });
662 }
663 self.writer
664 .as_ref()
665 .expect("an active recording owns its writer worker")
666 .flush_state_stream(stream)
667 }
668
669 pub fn complete_recording(self) -> Result<CompletedRecording, StorageError> {
677 self.complete_recording_with_terminal_metadata(Map::new())
678 }
679
680 pub fn complete_recording_with_terminal_metadata(
686 mut self,
687 terminal_metadata: Map<String, Value>,
688 ) -> Result<CompletedRecording, StorageError> {
689 if let Err(error) = self.finish_writer() {
690 let _ = self.transition_terminal(
691 RecordingStatus::Failed {
692 message: error.to_string(),
693 },
694 Map::new(),
695 );
696 return Err(error);
697 }
698 self.transition_terminal(RecordingStatus::Complete, terminal_metadata)?;
699 self.completed_recording()
700 }
701
702 pub fn complete_recording_with_final_state(
709 mut self,
710 state: &SystemState,
711 ) -> Result<CompletedRecording, StorageError> {
712 self.record_final_state(state)?;
713 self.complete_recording()
714 }
715
716 pub fn complete_recording_with_final_state_and_terminal_metadata(
719 mut self,
720 state: &SystemState,
721 terminal_metadata: Map<String, Value>,
722 ) -> Result<CompletedRecording, StorageError> {
723 self.record_final_state(state)?;
724 self.complete_recording_with_terminal_metadata(terminal_metadata)
725 }
726
727 fn record_final_state(&mut self, state: &SystemState) -> Result<(), StorageError> {
729 let iteration = state.simulation_time().iteration();
730 let writer = self
731 .writer
732 .as_ref()
733 .expect("an active recording owns its writer worker");
734 for name in &self.stream_order {
735 let stream = self
736 .streams
737 .get_mut(name)
738 .expect("stream order contains every configured stream");
739 if stream.last_recorded_iteration == Some(iteration) {
740 continue;
741 }
742 let record = stream.encoder.encode(state)?;
743 writer.submit_record(name, record)?;
744 stream.last_recorded_iteration = Some(iteration);
745 }
746 Ok(())
747 }
748
749 pub fn mark_recording_failed(self, message: impl Into<String>) -> Result<(), StorageError> {
761 self.mark_recording_failed_with_terminal_metadata(message, Map::new())
762 }
763
764 pub fn mark_recording_failed_with_terminal_metadata(
766 mut self,
767 message: impl Into<String>,
768 terminal_metadata: Map<String, Value>,
769 ) -> Result<(), StorageError> {
770 let message = message.into();
771 if message.trim().is_empty() {
772 return Err(StorageError::InvalidConfiguration {
773 setting: "failure_message",
774 reason: "failed run message must not be empty".to_owned(),
775 });
776 }
777
778 if let Err(error) = self.finish_writer() {
779 let _ = self.transition_terminal(
780 RecordingStatus::Failed {
781 message: error.to_string(),
782 },
783 Map::new(),
784 );
785 return Err(error);
786 }
787 self.transition_terminal(RecordingStatus::Failed { message }, terminal_metadata)
788 }
789
790 fn create_new_recording(builder: SystemStateWriterBuilder) -> Result<Self, StorageError> {
792 ensure_absent(&builder.root)?;
793 let prepared = PreparedRecording::from_builder(builder)?;
794 create_root(&prepared.root)?;
795 let lease = RecordingLease::acquire(&prepared.root)?;
796 for stream in &prepared.streams {
797 stream.writer.create_directory()?;
798 }
799 commit_metadata(&prepared.root, &prepared.metadata_path, &prepared.metadata)?;
800 let manifest = Arc::new(RecordingManifest::new(
801 prepared.root.clone(),
802 prepared.metadata_path.clone(),
803 prepared.metadata,
804 ));
805 Self::start_new_prepared(prepared.root, prepared.streams, manifest, lease)
806 }
807
808 fn continue_recording(
810 builder: SystemStateWriterBuilder,
811 checkpoint: Option<(&str, JsonPayloadDecoderRegistry)>,
812 ) -> Result<(Self, Option<SystemState>), StorageError> {
813 let prepared = PreparedRecording::from_builder(builder)?;
814 let lease = RecordingLease::acquire(&prepared.root)?;
815 remove_stale_metadata_temp(&prepared.root)?;
816 let mut existing = load_metadata(&prepared.metadata_path)?;
817 if !matches!(existing.status, RecordingStatus::Running) {
818 return Err(StorageError::RecordingNotContinuable {
819 path: prepared.metadata_path,
820 });
821 }
822 ensure_resume_match(&prepared.metadata_path, &prepared.metadata, &existing)?;
823
824 let mut recovered = Vec::with_capacity(prepared.streams.len());
825 for stream in prepared.streams {
826 let declaration = existing
827 .stream(&stream.name)
828 .expect("matched metadata contains every prepared stream");
829 let seed = StateWriterWorker::recover_state_stream(&stream.writer, declaration)?;
830 recovered.push((stream, seed));
831 }
832
833 let state = if let Some((checkpoint_stream, decoders)) = checkpoint {
834 let declaration = existing.stream(checkpoint_stream).ok_or_else(|| {
835 StorageError::UnknownStateStream {
836 stream: checkpoint_stream.to_owned(),
837 }
838 })?;
839 let seed = recovered
840 .iter()
841 .find(|(stream, _)| stream.name == checkpoint_stream)
842 .map(|(_, seed)| seed)
843 .expect("matched stream has one recovered seed");
844 Some(stored_state_series_reader::decode_resume_state(
845 &prepared.root,
846 &prepared.metadata_path,
847 declaration,
848 &prepared.spec,
849 &decoders,
850 seed.latest_open_record(),
851 )?)
852 } else {
853 None
854 };
855
856 existing.timing.continuation_count = existing
857 .timing
858 .continuation_count
859 .checked_add(1)
860 .ok_or_else(|| StorageError::InvalidMetadata {
861 path: prepared.metadata_path.clone(),
862 reason: "timing.continuation_count overflowed".to_owned(),
863 })?;
864 commit_metadata(&prepared.root, &prepared.metadata_path, &existing)?;
865 let manifest = Arc::new(RecordingManifest::new(
866 prepared.root.clone(),
867 prepared.metadata_path.clone(),
868 existing,
869 ));
870
871 let output = Self::start_resumed_prepared(prepared.root, recovered, manifest, lease)?;
872 Ok((output, state))
873 }
874
875 fn start_new_prepared(
877 root: PathBuf,
878 streams: Vec<PreparedStateStream>,
879 manifest: Arc<RecordingManifest>,
880 lease: RecordingLease,
881 ) -> Result<Self, StorageError> {
882 let mut scheduled = HashMap::with_capacity(streams.len());
883 let mut configs = Vec::with_capacity(streams.len());
884 let mut stream_order = Vec::with_capacity(streams.len());
885 for prepared in streams {
886 let name = prepared.name;
887 stream_order.push(name.clone());
888 scheduled.insert(
889 name,
890 ScheduledStateStream {
891 encoder: prepared.encoder,
892 sampling_interval: prepared.sampling_interval,
893 last_recorded_iteration: None,
894 },
895 );
896 configs.push(prepared.writer);
897 }
898 let writer = StateWriterWorker::start_new_recording(configs, Arc::clone(&manifest))?;
899 Ok(Self {
900 root,
901 stream_order,
902 manifest,
903 streams: scheduled,
904 writer: Some(writer),
905 session_started: Instant::now(),
906 _lease: lease,
907 })
908 }
909
910 fn start_resumed_prepared(
912 root: PathBuf,
913 streams: Vec<(PreparedStateStream, RecoveredStateStream)>,
914 manifest: Arc<RecordingManifest>,
915 lease: RecordingLease,
916 ) -> Result<Self, StorageError> {
917 let mut scheduled = HashMap::with_capacity(streams.len());
918 let mut recovered_streams = Vec::with_capacity(streams.len());
919 let mut stream_order = Vec::with_capacity(streams.len());
920 for (prepared, seed) in streams {
921 let name = prepared.name;
922 stream_order.push(name.clone());
923 scheduled.insert(
924 name,
925 ScheduledStateStream {
926 encoder: prepared.encoder,
927 sampling_interval: prepared.sampling_interval,
928 last_recorded_iteration: seed.last_iteration(),
929 },
930 );
931 recovered_streams.push((prepared.writer, seed));
932 }
933 let writer = StateWriterWorker::continue_recovered_recording(
934 recovered_streams,
935 Arc::clone(&manifest),
936 )?;
937 Ok(Self {
938 root,
939 stream_order,
940 manifest,
941 streams: scheduled,
942 writer: Some(writer),
943 session_started: Instant::now(),
944 _lease: lease,
945 })
946 }
947
948 fn finish_writer(&mut self) -> Result<(), StorageError> {
950 let Some(writer) = self.writer.take() else {
951 return Ok(());
952 };
953 writer.finish_recording()
954 }
955
956 fn transition_terminal(
958 &self,
959 status: RecordingStatus,
960 terminal_metadata: Map<String, Value>,
961 ) -> Result<(), StorageError> {
962 let finalized_at_utc =
963 utc_now_rfc3339().map_err(|source| StorageError::OperationalTimestamp {
964 operation: "finalize recording",
965 source,
966 })?;
967 let active_duration_ns = duration_nanoseconds(self.session_started.elapsed())
968 .ok_or(StorageError::OperationalDurationOverflow)?;
969 self.manifest.transition_terminal(
970 status,
971 finalized_at_utc,
972 active_duration_ns,
973 terminal_metadata,
974 )
975 }
976
977 fn completed_recording(&self) -> Result<CompletedRecording, StorageError> {
979 let metadata = self.manifest.snapshot();
980 let timing = RecordingTiming::from_stored(&metadata.timing, &self.manifest.path)?;
981 let streams = metadata
982 .streams
983 .iter()
984 .map(completed_stream_summary)
985 .collect::<Result<Vec<_>, _>>()?;
986 Ok(CompletedRecording {
987 directory: self.root.clone(),
988 timing,
989 terminal_metadata: metadata.terminal_metadata,
990 streams,
991 })
992 }
993}
994
995fn completed_stream_summary(
997 stream: &StateStreamMetadata,
998) -> Result<CompletedStreamSummary, StorageError> {
999 let overflow = || StorageError::ByteCountOverflow {
1000 stream: stream.name.clone(),
1001 };
1002 let chunk_count = u64::try_from(stream.chunks.len()).map_err(|_| overflow())?;
1003 let record_count = stream
1004 .chunks
1005 .iter()
1006 .try_fold(0_u64, |total, chunk| total.checked_add(chunk.records))
1007 .ok_or_else(&overflow)?;
1008 let encoded_bytes = stream
1009 .chunks
1010 .iter()
1011 .try_fold(0_u64, |total, chunk| total.checked_add(chunk.bytes))
1012 .ok_or_else(overflow)?;
1013 Ok(CompletedStreamSummary {
1014 name: stream.name.clone(),
1015 chunk_count,
1016 record_count,
1017 encoded_bytes,
1018 first_iteration: stream.chunks.first().map(|chunk| chunk.first_iteration),
1019 last_iteration: stream.chunks.last().map(|chunk| chunk.last_iteration),
1020 })
1021}
1022
1023struct PreparedRecording {
1025 root: PathBuf,
1026 metadata_path: PathBuf,
1027 spec: SystemStateSchema,
1028 metadata: RecordingMetadata,
1029 streams: Vec<PreparedStateStream>,
1030}
1031
1032impl PreparedRecording {
1033 fn from_builder(builder: SystemStateWriterBuilder) -> Result<Self, StorageError> {
1035 let metadata_path = builder.root.join(METADATA_FILE);
1036 let stored_time = builder.time.into_stored();
1037 let mut names = HashSet::with_capacity(builder.streams.len());
1038 let mut directories = HashSet::with_capacity(builder.streams.len());
1039 let mut streams = Vec::with_capacity(builder.streams.len());
1040 let mut declarations = Vec::with_capacity(builder.streams.len());
1041
1042 for config in builder.streams {
1043 if !names.insert(config.name.clone()) {
1044 return Err(StorageError::DuplicateStateStream {
1045 stream: config.name,
1046 });
1047 }
1048 if !directories.insert(config.directory.clone()) {
1049 return Err(StorageError::InvalidConfiguration {
1050 setting: "stream.directory",
1051 reason: format!(
1052 "multiple streams use relative directory `{}`",
1053 config.directory
1054 ),
1055 });
1056 }
1057
1058 let (max_chunk_bytes, queue_bytes) = config
1059 .storage_limits
1060 .or(builder.shared_stream_limits)
1061 .ok_or_else(|| StorageError::InvalidConfiguration {
1062 setting: "stream.storage_limits",
1063 reason: format!(
1064 "stream `{}` has no explicit limits and the writer has no shared limits",
1065 config.name
1066 ),
1067 })?;
1068 let encoder = JsonStateRecordEncoder::new(&config.name, &builder.spec, &config.fields)?;
1069 let fields = encoder
1070 .fields()
1071 .map(|name| {
1072 let field = builder
1073 .spec
1074 .field_schema(name)
1075 .expect("encoder fields were validated against this specification");
1076 StateFieldMetadata {
1077 name: name.to_owned(),
1078 description: field.description().map(str::to_owned),
1079 }
1080 })
1081 .collect::<Vec<_>>();
1082 declarations.push(StateStreamMetadata {
1083 name: config.name.clone(),
1084 directory: config.directory.clone(),
1085 sampling_interval: config.sampling_interval,
1086 fields,
1087 max_chunk_bytes: max_chunk_bytes.get(),
1088 queue_bytes: queue_bytes.get(),
1089 chunks: Vec::new(),
1090 });
1091 streams.push(PreparedStateStream {
1092 name: config.name.clone(),
1093 encoder,
1094 sampling_interval: config.sampling_interval,
1095 writer: StateStreamStorageConfig::new(
1096 &config.name,
1097 builder.root.join(&config.directory),
1098 max_chunk_bytes,
1099 queue_bytes,
1100 )?,
1101 });
1102 }
1103
1104 let created_at_utc =
1105 utc_now_rfc3339().map_err(|source| StorageError::OperationalTimestamp {
1106 operation: "create recording",
1107 source,
1108 })?;
1109 let metadata = RecordingMetadata::running(
1110 stored_time,
1111 builder.user_metadata,
1112 declarations,
1113 created_at_utc,
1114 );
1115 metadata.validate(&metadata_path)?;
1116 Ok(Self {
1117 root: builder.root,
1118 metadata_path,
1119 spec: builder.spec,
1120 metadata,
1121 streams,
1122 })
1123 }
1124}
1125
1126struct PreparedStateStream {
1128 name: String,
1129 encoder: JsonStateRecordEncoder,
1130 sampling_interval: SamplingInterval,
1131 writer: StateStreamStorageConfig,
1132}
1133
1134struct ScheduledStateStream {
1136 encoder: JsonStateRecordEncoder,
1137 sampling_interval: SamplingInterval,
1138 last_recorded_iteration: Option<u64>,
1139}
1140
1141pub(crate) struct RecordingManifest {
1147 root: PathBuf,
1148 path: PathBuf,
1149 metadata: Mutex<RecordingMetadata>,
1150}
1151
1152impl RecordingManifest {
1153 fn new(root: PathBuf, path: PathBuf, metadata: RecordingMetadata) -> Self {
1155 Self {
1156 root,
1157 path,
1158 metadata: Mutex::new(metadata),
1159 }
1160 }
1161
1162 pub(crate) fn prepare_chunk(
1164 &self,
1165 stream: &str,
1166 descriptor: jsonl_format::ChunkMetadata,
1167 ) -> Result<(), StorageError> {
1168 let mut current = lock_metadata(&self.metadata);
1169 if !matches!(current.status, RecordingStatus::Running) {
1170 return Err(StorageError::RecordingFinished);
1171 }
1172 let mut candidate = current.clone();
1173 let declaration =
1174 candidate
1175 .stream_mut(stream)
1176 .ok_or_else(|| StorageError::UnknownStateStream {
1177 stream: stream.to_owned(),
1178 })?;
1179 let expected = u64::try_from(declaration.chunks.len()).map_err(|_| {
1180 StorageError::ByteCountOverflow {
1181 stream: stream.to_owned(),
1182 }
1183 })?;
1184 if descriptor.ordinal != expected {
1185 return Err(StorageError::InvalidMetadata {
1186 path: self.path.clone(),
1187 reason: format!(
1188 "stream `{stream}` prepared chunk ordinal {}, expected {expected}",
1189 descriptor.ordinal
1190 ),
1191 });
1192 }
1193 declaration.chunks.push(descriptor);
1194 commit_metadata(&self.root, &self.path, &candidate)?;
1195 *current = candidate;
1196 Ok(())
1197 }
1198
1199 fn transition_terminal(
1201 &self,
1202 status: RecordingStatus,
1203 finalized_at_utc: String,
1204 active_duration_ns: u64,
1205 terminal_metadata: Map<String, Value>,
1206 ) -> Result<(), StorageError> {
1207 let mut current = lock_metadata(&self.metadata);
1208 let mut candidate = current.clone();
1209 candidate.status = status;
1210 candidate.timing.finalized_at_utc = Some(finalized_at_utc);
1211 candidate.timing.active_duration_ns = candidate
1212 .timing
1213 .active_duration_ns
1214 .checked_add(active_duration_ns)
1215 .ok_or(StorageError::OperationalDurationOverflow)?;
1216 candidate.terminal_metadata = terminal_metadata;
1217 commit_metadata(&self.root, &self.path, &candidate)?;
1218 *current = candidate;
1219 Ok(())
1220 }
1221
1222 fn snapshot(&self) -> RecordingMetadata {
1224 lock_metadata(&self.metadata).clone()
1225 }
1226}
1227
1228struct RecordingLease {
1233 _directory: File,
1234}
1235
1236impl RecordingLease {
1237 fn acquire(root: &Path) -> Result<Self, StorageError> {
1239 let directory = File::open(root).map_err(|source| StorageError::Io {
1240 operation: "open output root for exclusive ownership",
1241 path: root.to_path_buf(),
1242 source,
1243 })?;
1244 match FileExt::try_lock_exclusive(&directory) {
1245 Ok(()) => Ok(Self {
1246 _directory: directory,
1247 }),
1248 Err(source) if source.kind() == std::io::ErrorKind::WouldBlock => {
1249 Err(StorageError::RecordingDirectoryInUse {
1250 path: root.to_path_buf(),
1251 })
1252 }
1253 Err(source) => Err(StorageError::Io {
1254 operation: "acquire exclusive output ownership",
1255 path: root.to_path_buf(),
1256 source,
1257 }),
1258 }
1259 }
1260}
1261
1262fn load_metadata(path: &Path) -> Result<RecordingMetadata, StorageError> {
1264 let bytes = fs::read(path).map_err(|source| StorageError::Io {
1265 operation: "read metadata for resume",
1266 path: path.to_path_buf(),
1267 source,
1268 })?;
1269 let metadata: RecordingMetadata =
1270 serde_json::from_slice(&bytes).map_err(|source| StorageError::Json {
1271 operation: "parse metadata for resume",
1272 path: path.to_path_buf(),
1273 source,
1274 })?;
1275 metadata.validate(path)?;
1276 Ok(metadata)
1277}
1278
1279fn ensure_resume_match(
1281 path: &Path,
1282 expected: &RecordingMetadata,
1283 existing: &RecordingMetadata,
1284) -> Result<(), StorageError> {
1285 let mut configuration = existing.clone();
1286 for stream in &mut configuration.streams {
1287 stream.chunks.clear();
1288 }
1289 configuration.status = RecordingStatus::Running;
1290 configuration.timing = expected.timing.clone();
1291 configuration.terminal_metadata.clear();
1292 if &configuration != expected {
1293 return Err(StorageError::RecordingConfigurationMismatch {
1294 path: path.to_path_buf(),
1295 reason: "builder time axis, user metadata, or stream declarations differ".to_owned(),
1296 });
1297 }
1298 Ok(())
1299}
1300
1301fn remove_stale_metadata_temp(root: &Path) -> Result<(), StorageError> {
1303 let path = root.join(METADATA_TEMP_FILE);
1304 match fs::remove_file(&path) {
1305 Ok(()) => File::open(root)
1306 .and_then(|directory| directory.sync_all())
1307 .map_err(|source| StorageError::Io {
1308 operation: "synchronize stale metadata cleanup",
1309 path: root.to_path_buf(),
1310 source,
1311 }),
1312 Err(source) if source.kind() == std::io::ErrorKind::NotFound => Ok(()),
1313 Err(source) => Err(StorageError::Io {
1314 operation: "remove stale temporary metadata",
1315 path,
1316 source,
1317 }),
1318 }
1319}
1320
1321fn lock_metadata(metadata: &Mutex<RecordingMetadata>) -> MutexGuard<'_, RecordingMetadata> {
1323 metadata
1324 .lock()
1325 .unwrap_or_else(|poisoned| poisoned.into_inner())
1326}
1327
1328fn ensure_absent(root: &Path) -> Result<(), StorageError> {
1330 match root.try_exists() {
1331 Ok(false) => Ok(()),
1332 Ok(true) => Err(StorageError::RecordingDirectoryExists {
1333 path: root.to_path_buf(),
1334 }),
1335 Err(source) => Err(StorageError::Io {
1336 operation: "inspect output root",
1337 path: root.to_path_buf(),
1338 source,
1339 }),
1340 }
1341}
1342
1343fn create_root(root: &Path) -> Result<(), StorageError> {
1345 match fs::create_dir(root) {
1346 Ok(()) => Ok(()),
1347 Err(source) if source.kind() == std::io::ErrorKind::AlreadyExists => {
1348 Err(StorageError::RecordingDirectoryExists {
1349 path: root.to_path_buf(),
1350 })
1351 }
1352 Err(source) => Err(StorageError::Io {
1353 operation: "create output root",
1354 path: root.to_path_buf(),
1355 source,
1356 }),
1357 }
1358}
1359
1360fn commit_metadata(
1368 root: &Path,
1369 metadata_path: &Path,
1370 metadata: &RecordingMetadata,
1371) -> Result<(), StorageError> {
1372 metadata.validate(metadata_path)?;
1373 let mut bytes = serde_json::to_vec_pretty(metadata).map_err(|source| StorageError::Json {
1374 operation: "serialize metadata",
1375 path: metadata_path.to_path_buf(),
1376 source,
1377 })?;
1378 bytes.push(b'\n');
1379
1380 let temporary_path = root.join(METADATA_TEMP_FILE);
1381 let result = write_and_replace_metadata(root, metadata_path, &temporary_path, &bytes);
1382 if result.is_err() {
1383 let _ = fs::remove_file(&temporary_path);
1384 }
1385 result
1386}
1387
1388fn write_and_replace_metadata(
1390 root: &Path,
1391 metadata_path: &Path,
1392 temporary_path: &Path,
1393 bytes: &[u8],
1394) -> Result<(), StorageError> {
1395 let mut temporary = OpenOptions::new()
1396 .write(true)
1397 .create_new(true)
1398 .open(temporary_path)
1399 .map_err(|source| StorageError::Io {
1400 operation: "create temporary metadata",
1401 path: temporary_path.to_path_buf(),
1402 source,
1403 })?;
1404 temporary
1405 .write_all(bytes)
1406 .map_err(|source| StorageError::Io {
1407 operation: "write temporary metadata",
1408 path: temporary_path.to_path_buf(),
1409 source,
1410 })?;
1411 temporary.sync_all().map_err(|source| StorageError::Io {
1412 operation: "sync temporary metadata",
1413 path: temporary_path.to_path_buf(),
1414 source,
1415 })?;
1416 drop(temporary);
1417
1418 fs::rename(temporary_path, metadata_path).map_err(|source| StorageError::Io {
1419 operation: "publish metadata",
1420 path: metadata_path.to_path_buf(),
1421 source,
1422 })?;
1423
1424 File::open(root)
1425 .and_then(|directory| directory.sync_all())
1426 .map_err(|source| StorageError::Io {
1427 operation: "sync output root",
1428 path: root.to_path_buf(),
1429 source,
1430 })
1431}