use std::collections::{HashMap, HashSet};
use std::fs::{self, File, OpenOptions};
use std::io::Write;
use std::num::NonZeroU64;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::{Duration, Instant};
use fs2::FileExt;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::clock::{duration_nanoseconds, utc_now_rfc3339};
use crate::configuration::ResolvedConfiguration;
use crate::system_state::{StateSchemaSource, SystemState, SystemStateSchema};
mod error;
mod json_payload_decoder;
mod json_state_record_encoder;
mod jsonl_format;
mod queued_state_writer;
mod resume;
mod stored_state_series_reader;
pub use error::StorageError;
pub use json_payload_decoder::{
JsonPayloadDecoder, JsonPayloadDecoderRegistry, JsonStringDecoder, JsonVecF64Decoder,
};
pub use stored_state_series_reader::StoredStateSeriesReader;
use json_state_record_encoder::JsonStateRecordEncoder;
use jsonl_format::{
RecordingMetadata, RecordingStatus, StateFieldMetadata, StateStreamMetadata,
TimeAxisMetadata as StoredTimeAxis,
};
use queued_state_writer::{RecoveredStateStream, StateStreamStorageConfig, StateWriterWorker};
const METADATA_FILE: &str = "metadata.json";
const METADATA_TEMP_FILE: &str = ".metadata.json.tmp";
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct TimeAxisMetadata {
iteration_name: String,
iteration_unit: Option<String>,
physical_time_name: Option<String>,
physical_time_unit: Option<String>,
}
impl TimeAxisMetadata {
pub fn new(iteration_name: impl Into<String>) -> Self {
Self {
iteration_name: iteration_name.into(),
iteration_unit: None,
physical_time_name: None,
physical_time_unit: None,
}
}
#[must_use]
pub fn with_iteration_unit(mut self, unit: impl Into<String>) -> Self {
self.iteration_unit = Some(unit.into());
self
}
#[must_use]
pub fn with_physical_time_name(mut self, name: impl Into<String>) -> Self {
self.physical_time_name = Some(name.into());
self
}
#[must_use]
pub fn with_physical_time_unit(mut self, unit: impl Into<String>) -> Self {
self.physical_time_unit = Some(unit.into());
self
}
#[must_use]
pub fn with_physical_axis(mut self, name: impl Into<String>, unit: impl Into<String>) -> Self {
self.physical_time_name = Some(name.into());
self.physical_time_unit = Some(unit.into());
self
}
fn into_stored(self) -> StoredTimeAxis {
StoredTimeAxis {
iteration_name: self.iteration_name,
iteration_unit: self.iteration_unit,
physical_time_name: self.physical_time_name,
physical_time_unit: self.physical_time_unit,
}
}
}
impl Default for TimeAxisMetadata {
fn default() -> Self {
Self::new("iteration")
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct RecordingTiming {
created_at_utc: String,
finalized_at_utc: String,
active_duration_ns: u64,
continuation_count: u64,
}
impl RecordingTiming {
fn from_stored(
timing: &jsonl_format::RecordingTiming,
metadata_path: &Path,
) -> Result<Self, StorageError> {
let finalized_at_utc =
timing
.finalized_at_utc
.clone()
.ok_or_else(|| StorageError::InvalidMetadata {
path: metadata_path.to_path_buf(),
reason: "completed recording lacks finalized timestamp".to_owned(),
})?;
Ok(Self {
created_at_utc: timing.created_at_utc.clone(),
finalized_at_utc,
active_duration_ns: timing.active_duration_ns,
continuation_count: timing.continuation_count,
})
}
pub fn created_at_utc(&self) -> &str {
&self.created_at_utc
}
pub fn finalized_at_utc(&self) -> &str {
&self.finalized_at_utc
}
pub fn active_duration_ns(&self) -> u64 {
self.active_duration_ns
}
pub fn active_duration(&self) -> Duration {
Duration::from_nanos(self.active_duration_ns)
}
pub fn continuation_count(&self) -> u64 {
self.continuation_count
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CompletedStreamSummary {
name: String,
chunk_count: u64,
record_count: u64,
encoded_bytes: u64,
first_iteration: Option<u64>,
last_iteration: Option<u64>,
}
impl CompletedStreamSummary {
pub fn name(&self) -> &str {
&self.name
}
pub fn chunk_count(&self) -> u64 {
self.chunk_count
}
pub fn record_count(&self) -> u64 {
self.record_count
}
pub fn encoded_bytes(&self) -> u64 {
self.encoded_bytes
}
pub fn first_iteration(&self) -> Option<u64> {
self.first_iteration
}
pub fn last_iteration(&self) -> Option<u64> {
self.last_iteration
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CompletedRecording {
directory: PathBuf,
timing: RecordingTiming,
terminal_metadata: Map<String, Value>,
streams: Vec<CompletedStreamSummary>,
}
impl CompletedRecording {
pub fn directory(&self) -> &Path {
&self.directory
}
pub fn timing(&self) -> &RecordingTiming {
&self.timing
}
pub fn terminal_metadata(&self) -> &Map<String, Value> {
&self.terminal_metadata
}
pub fn stream_summaries(&self) -> &[CompletedStreamSummary] {
&self.streams
}
pub fn stream_summary(&self, name: &str) -> Option<&CompletedStreamSummary> {
self.streams.iter().find(|stream| stream.name == name)
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum SamplingInterval {
Iterations(NonZeroU64),
}
#[derive(Deserialize)]
#[serde(untagged)]
enum SamplingIntervalInput {
Iterations(NonZeroU64),
Tagged { iterations: NonZeroU64 },
}
impl<'de> Deserialize<'de> for SamplingInterval {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
match SamplingIntervalInput::deserialize(deserializer)? {
SamplingIntervalInput::Iterations(interval)
| SamplingIntervalInput::Tagged {
iterations: interval,
} => Ok(Self::Iterations(interval)),
}
}
}
impl SamplingInterval {
pub const fn iterations(interval: u64) -> Option<Self> {
match NonZeroU64::new(interval) {
Some(interval) => Some(Self::Iterations(interval)),
None => None,
}
}
const fn includes(self, iteration: u64) -> bool {
match self {
Self::Iterations(interval) => iteration.is_multiple_of(interval.get()),
}
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq, Deserialize, Serialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
pub enum StateStreamLayout {
Chunked {
target_bytes: NonZeroU64,
},
IndividualFiles,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub struct StateStreamStorage {
layout: StateStreamLayout,
storage_queue_bytes: NonZeroU64,
}
impl StateStreamStorage {
pub const fn chunked(target_bytes: NonZeroU64, storage_queue_bytes: NonZeroU64) -> Self {
Self {
layout: StateStreamLayout::Chunked { target_bytes },
storage_queue_bytes,
}
}
pub const fn individual_files(storage_queue_bytes: NonZeroU64) -> Self {
Self {
layout: StateStreamLayout::IndividualFiles,
storage_queue_bytes,
}
}
pub const fn layout(self) -> StateStreamLayout {
self.layout
}
pub const fn storage_queue_bytes(self) -> NonZeroU64 {
self.storage_queue_bytes
}
}
#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub struct StateStreamConfig {
name: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
directory: Option<String>,
sampling_interval: SamplingInterval,
fields: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
storage: Option<StateStreamStorage>,
}
impl StateStreamConfig {
pub fn new<I, K>(
name: impl Into<String>,
fields: I,
sampling_interval: SamplingInterval,
storage: Option<StateStreamStorage>,
) -> Self
where
I: IntoIterator<Item = K>,
K: Into<String>,
{
let name = name.into();
Self {
directory: None,
name,
sampling_interval,
fields: fields.into_iter().map(Into::into).collect(),
storage,
}
}
#[must_use]
pub fn with_relative_directory(mut self, directory: impl Into<String>) -> Self {
self.directory = Some(directory.into());
self
}
pub fn name(&self) -> &str {
&self.name
}
pub fn relative_directory(&self) -> &str {
self.directory.as_deref().unwrap_or(&self.name)
}
pub const fn sampling_interval(&self) -> SamplingInterval {
self.sampling_interval
}
pub fn fields(&self) -> &[String] {
&self.fields
}
pub const fn storage(&self) -> Option<StateStreamStorage> {
self.storage
}
}
#[derive(Debug)]
pub struct SystemStateWriterBuilder {
root: PathBuf,
spec: SystemStateSchema,
time: TimeAxisMetadata,
user_metadata: Map<String, Value>,
shared_stream_storage: Option<StateStreamStorage>,
streams: Vec<StateStreamConfig>,
}
impl SystemStateWriterBuilder {
pub fn new<S>(root: impl Into<PathBuf>, source: &S) -> Self
where
S: StateSchemaSource + ?Sized,
{
Self {
root: root.into(),
spec: source.state_schema().clone(),
time: TimeAxisMetadata::default(),
user_metadata: Map::new(),
shared_stream_storage: None,
streams: Vec::new(),
}
}
#[must_use]
pub fn with_time_axis_metadata(mut self, time: TimeAxisMetadata) -> Self {
self.time = time;
self
}
#[must_use]
pub fn with_user_metadata(mut self, metadata: Map<String, Value>) -> Self {
self.user_metadata.extend(metadata);
self
}
#[must_use]
pub fn with_shared_stream_storage(mut self, storage: StateStreamStorage) -> Self {
self.shared_stream_storage = Some(storage);
self
}
#[must_use]
pub fn with_configuration(mut self, configuration: &ResolvedConfiguration) -> Self {
self.user_metadata
.extend(configuration.resolved_object().clone());
self.user_metadata
.insert("ordinal".to_owned(), Value::from(configuration.ordinal()));
self
}
#[must_use]
pub fn add_state_stream(mut self, stream: StateStreamConfig) -> Self {
self.streams.push(stream);
self
}
pub fn create_new_recording(self) -> Result<SystemStateWriter, StorageError> {
SystemStateWriter::create_new_recording(self)
}
pub fn open_or_resume_from_latest_checkpoint(
self,
decoders: JsonPayloadDecoderRegistry,
) -> Result<(SystemStateWriter, Option<SystemState>), StorageError> {
match self.root.try_exists() {
Ok(false) => Self::create_new_recording(self).map(|writer| (writer, None)),
Ok(true) => SystemStateWriter::continue_recording(
self,
Some(CheckpointRequest::LatestComplete(decoders)),
),
Err(source) => Err(StorageError::Io {
operation: "inspect recording root for automatic resume",
path: self.root.clone(),
source,
}),
}
}
pub fn continue_existing_recording(self) -> Result<SystemStateWriter, StorageError> {
SystemStateWriter::continue_recording(self, None).map(|(writer, _)| writer)
}
pub fn continue_recording_from_latest_checkpoint(
self,
stream: &str,
decoders: JsonPayloadDecoderRegistry,
) -> Result<(SystemStateWriter, SystemState), StorageError> {
let (writer, state) = SystemStateWriter::continue_recording(
self,
Some(CheckpointRequest::Named(stream.to_owned(), decoders)),
)?;
Ok((
writer,
state.expect("checkpoint-aware resume always reconstructs one state"),
))
}
}
enum CheckpointRequest {
Named(String, JsonPayloadDecoderRegistry),
LatestComplete(JsonPayloadDecoderRegistry),
}
pub struct SystemStateWriter {
root: PathBuf,
stream_order: Vec<String>,
manifest: Arc<RecordingManifest>,
streams: HashMap<String, ScheduledStateStream>,
writer: Option<StateWriterWorker>,
session_started: Instant,
_lease: RecordingLease,
}
impl SystemStateWriter {
pub fn builder<S>(root: impl Into<PathBuf>, source: &S) -> SystemStateWriterBuilder
where
S: StateSchemaSource + ?Sized,
{
SystemStateWriterBuilder::new(root, source)
}
pub fn recording_directory(&self) -> &Path {
&self.root
}
pub fn stream_names(&self) -> impl ExactSizeIterator<Item = &str> {
self.stream_order.iter().map(String::as_str)
}
pub fn observe_state(&mut self, state: &SystemState) -> Result<(), StorageError> {
let iteration = state.simulation_time().iteration();
let writer = self
.writer
.as_ref()
.expect("an active recording owns its writer worker");
for name in &self.stream_order {
let stream = self
.streams
.get_mut(name)
.expect("stream order contains every configured stream");
if !stream.sampling_interval.includes(iteration)
|| stream.last_recorded_iteration == Some(iteration)
{
continue;
}
let record = stream.encoder.encode(state)?;
writer.submit_record(name, record)?;
stream.last_recorded_iteration = Some(iteration);
}
Ok(())
}
pub fn flush_stream_to_storage(&self, stream: &str) -> Result<(), StorageError> {
if !self.streams.contains_key(stream) {
return Err(StorageError::UnknownStateStream {
stream: stream.to_owned(),
});
}
self.writer
.as_ref()
.expect("an active recording owns its writer worker")
.flush_state_stream(stream)
}
pub fn complete_recording(self) -> Result<CompletedRecording, StorageError> {
self.complete_recording_with_terminal_metadata(Map::new())
}
pub fn complete_recording_with_terminal_metadata(
mut self,
terminal_metadata: Map<String, Value>,
) -> Result<CompletedRecording, StorageError> {
if let Err(error) = self.finish_writer() {
let _ = self.transition_terminal(
RecordingStatus::Failed {
message: error.to_string(),
},
Map::new(),
);
return Err(error);
}
self.transition_terminal(RecordingStatus::Complete, terminal_metadata)?;
self.completed_recording()
}
pub fn complete_recording_with_final_state(
mut self,
state: &SystemState,
) -> Result<CompletedRecording, StorageError> {
self.record_final_state(state)?;
self.complete_recording()
}
pub fn complete_recording_with_final_state_and_terminal_metadata(
mut self,
state: &SystemState,
terminal_metadata: Map<String, Value>,
) -> Result<CompletedRecording, StorageError> {
self.record_final_state(state)?;
self.complete_recording_with_terminal_metadata(terminal_metadata)
}
fn record_final_state(&mut self, state: &SystemState) -> Result<(), StorageError> {
let iteration = state.simulation_time().iteration();
let writer = self
.writer
.as_ref()
.expect("an active recording owns its writer worker");
for name in &self.stream_order {
let stream = self
.streams
.get_mut(name)
.expect("stream order contains every configured stream");
if stream.last_recorded_iteration == Some(iteration) {
continue;
}
let record = stream.encoder.encode(state)?;
writer.submit_record(name, record)?;
stream.last_recorded_iteration = Some(iteration);
}
Ok(())
}
pub fn mark_recording_failed(self, message: impl Into<String>) -> Result<(), StorageError> {
self.mark_recording_failed_with_terminal_metadata(message, Map::new())
}
pub fn mark_recording_failed_with_terminal_metadata(
mut self,
message: impl Into<String>,
terminal_metadata: Map<String, Value>,
) -> Result<(), StorageError> {
let message = message.into();
if message.trim().is_empty() {
return Err(StorageError::InvalidConfiguration {
setting: "failure_message",
reason: "failed run message must not be empty".to_owned(),
});
}
if let Err(error) = self.finish_writer() {
let _ = self.transition_terminal(
RecordingStatus::Failed {
message: error.to_string(),
},
Map::new(),
);
return Err(error);
}
self.transition_terminal(RecordingStatus::Failed { message }, terminal_metadata)
}
fn create_new_recording(builder: SystemStateWriterBuilder) -> Result<Self, StorageError> {
ensure_absent(&builder.root)?;
let prepared = PreparedRecording::from_builder(builder)?;
create_root(&prepared.root)?;
let lease = RecordingLease::acquire(&prepared.root)?;
for stream in &prepared.streams {
stream.writer.create_directory()?;
}
commit_metadata(&prepared.root, &prepared.metadata_path, &prepared.metadata)?;
let manifest = Arc::new(RecordingManifest::new(
prepared.root.clone(),
prepared.metadata_path.clone(),
prepared.metadata,
));
Self::start_new_prepared(prepared.root, prepared.streams, manifest, lease)
}
fn continue_recording(
builder: SystemStateWriterBuilder,
checkpoint: Option<CheckpointRequest>,
) -> Result<(Self, Option<SystemState>), StorageError> {
let prepared = PreparedRecording::from_builder(builder)?;
let lease = RecordingLease::acquire(&prepared.root)?;
remove_stale_metadata_temp(&prepared.root)?;
let mut existing = load_metadata(&prepared.metadata_path)?;
if !matches!(existing.status, RecordingStatus::Running) {
return Err(StorageError::RecordingNotContinuable {
path: prepared.metadata_path,
});
}
ensure_resume_match(&prepared.metadata_path, &prepared.metadata, &existing)?;
if checkpoint.is_some() {
for stream in &prepared.streams {
let declaration = existing
.stream(&stream.name)
.expect("matched metadata contains every prepared stream");
StateWriterWorker::recover_state_stream(&stream.writer, declaration)?;
}
}
let state = if let Some(checkpoint) = checkpoint {
let (checkpoint_stream, decoders) = match checkpoint {
CheckpointRequest::Named(stream, decoders) => (stream, decoders),
CheckpointRequest::LatestComplete(decoders) => {
let stream = existing
.streams
.iter()
.filter(|stream| {
stored_state_series_reader::is_complete_checkpoint_stream(
stream,
&prepared.spec,
) && !stream.chunks.is_empty()
})
.max_by_key(|stream| stream.chunks.last().map(|chunk| chunk.last_iteration))
.map(|stream| stream.name.clone())
.ok_or(StorageError::NoCompleteCheckpoint)?;
(stream, decoders)
}
};
let declaration = existing.stream(&checkpoint_stream).ok_or_else(|| {
StorageError::UnknownStateStream {
stream: checkpoint_stream.clone(),
}
})?;
let state = stored_state_series_reader::decode_resume_state(
&prepared.root,
&prepared.metadata_path,
declaration,
&prepared.spec,
&decoders,
)?;
resume::prepare_rewind_after_checkpoint(
&prepared.root,
&prepared.metadata_path,
&mut existing,
state.simulation_time().iteration(),
)?;
commit_metadata(&prepared.root, &prepared.metadata_path, &existing)?;
Some(state)
} else {
None
};
let mut recovered = Vec::with_capacity(prepared.streams.len());
for stream in prepared.streams {
let declaration = existing
.stream(&stream.name)
.expect("matched metadata contains every prepared stream");
let seed = StateWriterWorker::recover_state_stream(&stream.writer, declaration)?;
recovered.push((stream, seed));
}
existing.timing.continuation_count = existing
.timing
.continuation_count
.checked_add(1)
.ok_or_else(|| StorageError::InvalidMetadata {
path: prepared.metadata_path.clone(),
reason: "timing.continuation_count overflowed".to_owned(),
})?;
commit_metadata(&prepared.root, &prepared.metadata_path, &existing)?;
let manifest = Arc::new(RecordingManifest::new(
prepared.root.clone(),
prepared.metadata_path.clone(),
existing,
));
let output = Self::start_resumed_prepared(prepared.root, recovered, manifest, lease)?;
Ok((output, state))
}
fn start_new_prepared(
root: PathBuf,
streams: Vec<PreparedStateStream>,
manifest: Arc<RecordingManifest>,
lease: RecordingLease,
) -> Result<Self, StorageError> {
let mut scheduled = HashMap::with_capacity(streams.len());
let mut configs = Vec::with_capacity(streams.len());
let mut stream_order = Vec::with_capacity(streams.len());
for prepared in streams {
let name = prepared.name;
stream_order.push(name.clone());
scheduled.insert(
name,
ScheduledStateStream {
encoder: prepared.encoder,
sampling_interval: prepared.sampling_interval,
last_recorded_iteration: None,
},
);
configs.push(prepared.writer);
}
let writer = StateWriterWorker::start_new_recording(configs, Arc::clone(&manifest))?;
Ok(Self {
root,
stream_order,
manifest,
streams: scheduled,
writer: Some(writer),
session_started: Instant::now(),
_lease: lease,
})
}
fn start_resumed_prepared(
root: PathBuf,
streams: Vec<(PreparedStateStream, RecoveredStateStream)>,
manifest: Arc<RecordingManifest>,
lease: RecordingLease,
) -> Result<Self, StorageError> {
let mut scheduled = HashMap::with_capacity(streams.len());
let mut recovered_streams = Vec::with_capacity(streams.len());
let mut stream_order = Vec::with_capacity(streams.len());
for (prepared, seed) in streams {
let name = prepared.name;
stream_order.push(name.clone());
scheduled.insert(
name,
ScheduledStateStream {
encoder: prepared.encoder,
sampling_interval: prepared.sampling_interval,
last_recorded_iteration: seed.last_iteration(),
},
);
recovered_streams.push((prepared.writer, seed));
}
let writer = StateWriterWorker::continue_recovered_recording(
recovered_streams,
Arc::clone(&manifest),
)?;
Ok(Self {
root,
stream_order,
manifest,
streams: scheduled,
writer: Some(writer),
session_started: Instant::now(),
_lease: lease,
})
}
fn finish_writer(&mut self) -> Result<(), StorageError> {
let Some(writer) = self.writer.take() else {
return Ok(());
};
writer.finish_recording()
}
fn transition_terminal(
&self,
status: RecordingStatus,
terminal_metadata: Map<String, Value>,
) -> Result<(), StorageError> {
let finalized_at_utc =
utc_now_rfc3339().map_err(|source| StorageError::OperationalTimestamp {
operation: "finalize recording",
source,
})?;
let active_duration_ns = duration_nanoseconds(self.session_started.elapsed())
.ok_or(StorageError::OperationalDurationOverflow)?;
self.manifest.transition_terminal(
status,
finalized_at_utc,
active_duration_ns,
terminal_metadata,
)
}
fn completed_recording(&self) -> Result<CompletedRecording, StorageError> {
let metadata = self.manifest.snapshot();
let timing = RecordingTiming::from_stored(&metadata.timing, &self.manifest.path)?;
let streams = metadata
.streams
.iter()
.map(completed_stream_summary)
.collect::<Result<Vec<_>, _>>()?;
Ok(CompletedRecording {
directory: self.root.clone(),
timing,
terminal_metadata: metadata.terminal_metadata,
streams,
})
}
}
fn completed_stream_summary(
stream: &StateStreamMetadata,
) -> Result<CompletedStreamSummary, StorageError> {
let overflow = || StorageError::ByteCountOverflow {
stream: stream.name.clone(),
};
let chunk_count = u64::try_from(stream.chunks.len()).map_err(|_| overflow())?;
let record_count = stream
.chunks
.iter()
.try_fold(0_u64, |total, chunk| total.checked_add(chunk.records))
.ok_or_else(&overflow)?;
let encoded_bytes = stream
.chunks
.iter()
.try_fold(0_u64, |total, chunk| total.checked_add(chunk.bytes))
.ok_or_else(overflow)?;
Ok(CompletedStreamSummary {
name: stream.name.clone(),
chunk_count,
record_count,
encoded_bytes,
first_iteration: stream.chunks.first().map(|chunk| chunk.first_iteration),
last_iteration: stream.chunks.last().map(|chunk| chunk.last_iteration),
})
}
struct PreparedRecording {
root: PathBuf,
metadata_path: PathBuf,
spec: SystemStateSchema,
metadata: RecordingMetadata,
streams: Vec<PreparedStateStream>,
}
impl PreparedRecording {
fn from_builder(builder: SystemStateWriterBuilder) -> Result<Self, StorageError> {
let metadata_path = builder.root.join(METADATA_FILE);
let stored_time = builder.time.into_stored();
let mut names = HashSet::with_capacity(builder.streams.len());
let mut directories = HashSet::with_capacity(builder.streams.len());
let mut streams = Vec::with_capacity(builder.streams.len());
let mut declarations = Vec::with_capacity(builder.streams.len());
for config in builder.streams {
if !names.insert(config.name.clone()) {
return Err(StorageError::DuplicateStateStream {
stream: config.name,
});
}
let directory = config.relative_directory().to_owned();
if !directories.insert(directory.clone()) {
return Err(StorageError::InvalidConfiguration {
setting: "stream.directory",
reason: format!("multiple streams use relative directory `{}`", directory),
});
}
let storage = config
.storage
.or(builder.shared_stream_storage)
.ok_or_else(|| StorageError::InvalidConfiguration {
setting: "stream.storage",
reason: format!(
"stream `{}` has no explicit storage and the writer has no shared storage",
config.name
),
})?;
let encoder = JsonStateRecordEncoder::new(&config.name, &builder.spec, &config.fields)?;
let fields = encoder
.fields()
.map(|name| {
let field = builder
.spec
.field_schema(name)
.expect("encoder fields were validated against this specification");
StateFieldMetadata {
name: name.to_owned(),
description: field.description().map(str::to_owned),
}
})
.collect::<Vec<_>>();
declarations.push(StateStreamMetadata {
name: config.name.clone(),
directory: directory.clone(),
sampling_interval: config.sampling_interval,
fields,
storage,
chunks: Vec::new(),
});
streams.push(PreparedStateStream {
name: config.name.clone(),
encoder,
sampling_interval: config.sampling_interval,
writer: StateStreamStorageConfig::new(
&config.name,
builder.root.join(&directory),
storage,
)?,
});
}
let created_at_utc =
utc_now_rfc3339().map_err(|source| StorageError::OperationalTimestamp {
operation: "create recording",
source,
})?;
let metadata = RecordingMetadata::running(
stored_time,
builder.user_metadata,
declarations,
created_at_utc,
);
metadata.validate(&metadata_path)?;
Ok(Self {
root: builder.root,
metadata_path,
spec: builder.spec,
metadata,
streams,
})
}
}
struct PreparedStateStream {
name: String,
encoder: JsonStateRecordEncoder,
sampling_interval: SamplingInterval,
writer: StateStreamStorageConfig,
}
struct ScheduledStateStream {
encoder: JsonStateRecordEncoder,
sampling_interval: SamplingInterval,
last_recorded_iteration: Option<u64>,
}
pub(crate) struct RecordingManifest {
root: PathBuf,
path: PathBuf,
metadata: Mutex<RecordingMetadata>,
}
impl RecordingManifest {
fn new(root: PathBuf, path: PathBuf, metadata: RecordingMetadata) -> Self {
Self {
root,
path,
metadata: Mutex::new(metadata),
}
}
pub(crate) fn prepare_chunk(
&self,
stream: &str,
descriptor: jsonl_format::ChunkMetadata,
) -> Result<(), StorageError> {
let mut current = lock_metadata(&self.metadata);
if !matches!(current.status, RecordingStatus::Running) {
return Err(StorageError::RecordingFinished);
}
let mut candidate = current.clone();
let declaration =
candidate
.stream_mut(stream)
.ok_or_else(|| StorageError::UnknownStateStream {
stream: stream.to_owned(),
})?;
let expected = u64::try_from(declaration.chunks.len()).map_err(|_| {
StorageError::ByteCountOverflow {
stream: stream.to_owned(),
}
})?;
if descriptor.ordinal != expected {
return Err(StorageError::InvalidMetadata {
path: self.path.clone(),
reason: format!(
"stream `{stream}` prepared chunk ordinal {}, expected {expected}",
descriptor.ordinal
),
});
}
declaration.chunks.push(descriptor);
commit_metadata(&self.root, &self.path, &candidate)?;
*current = candidate;
Ok(())
}
fn transition_terminal(
&self,
status: RecordingStatus,
finalized_at_utc: String,
active_duration_ns: u64,
terminal_metadata: Map<String, Value>,
) -> Result<(), StorageError> {
let mut current = lock_metadata(&self.metadata);
let mut candidate = current.clone();
candidate.status = status;
candidate.timing.finalized_at_utc = Some(finalized_at_utc);
candidate.timing.active_duration_ns = candidate
.timing
.active_duration_ns
.checked_add(active_duration_ns)
.ok_or(StorageError::OperationalDurationOverflow)?;
candidate.terminal_metadata = terminal_metadata;
commit_metadata(&self.root, &self.path, &candidate)?;
*current = candidate;
Ok(())
}
fn snapshot(&self) -> RecordingMetadata {
lock_metadata(&self.metadata).clone()
}
}
struct RecordingLease {
_directory: File,
}
impl RecordingLease {
fn acquire(root: &Path) -> Result<Self, StorageError> {
let directory = File::open(root).map_err(|source| StorageError::Io {
operation: "open output root for exclusive ownership",
path: root.to_path_buf(),
source,
})?;
match FileExt::try_lock_exclusive(&directory) {
Ok(()) => Ok(Self {
_directory: directory,
}),
Err(source) if source.kind() == std::io::ErrorKind::WouldBlock => {
Err(StorageError::RecordingDirectoryInUse {
path: root.to_path_buf(),
})
}
Err(source) => Err(StorageError::Io {
operation: "acquire exclusive output ownership",
path: root.to_path_buf(),
source,
}),
}
}
}
fn load_metadata(path: &Path) -> Result<RecordingMetadata, StorageError> {
let bytes = fs::read(path).map_err(|source| StorageError::Io {
operation: "read metadata for resume",
path: path.to_path_buf(),
source,
})?;
let metadata: RecordingMetadata =
serde_json::from_slice(&bytes).map_err(|source| StorageError::Json {
operation: "parse metadata for resume",
path: path.to_path_buf(),
source,
})?;
metadata.validate(path)?;
Ok(metadata)
}
fn ensure_resume_match(
path: &Path,
expected: &RecordingMetadata,
existing: &RecordingMetadata,
) -> Result<(), StorageError> {
let mut configuration = existing.clone();
for stream in &mut configuration.streams {
stream.chunks.clear();
}
configuration.status = RecordingStatus::Running;
configuration.timing = expected.timing.clone();
configuration.terminal_metadata.clear();
if &configuration != expected {
return Err(StorageError::RecordingConfigurationMismatch {
path: path.to_path_buf(),
reason: "builder time axis, user metadata, or stream declarations differ".to_owned(),
});
}
Ok(())
}
fn remove_stale_metadata_temp(root: &Path) -> Result<(), StorageError> {
let path = root.join(METADATA_TEMP_FILE);
match fs::remove_file(&path) {
Ok(()) => File::open(root)
.and_then(|directory| directory.sync_all())
.map_err(|source| StorageError::Io {
operation: "synchronize stale metadata cleanup",
path: root.to_path_buf(),
source,
}),
Err(source) if source.kind() == std::io::ErrorKind::NotFound => Ok(()),
Err(source) => Err(StorageError::Io {
operation: "remove stale temporary metadata",
path,
source,
}),
}
}
fn lock_metadata(metadata: &Mutex<RecordingMetadata>) -> MutexGuard<'_, RecordingMetadata> {
metadata
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
fn ensure_absent(root: &Path) -> Result<(), StorageError> {
match root.try_exists() {
Ok(false) => Ok(()),
Ok(true) => Err(StorageError::RecordingDirectoryExists {
path: root.to_path_buf(),
}),
Err(source) => Err(StorageError::Io {
operation: "inspect output root",
path: root.to_path_buf(),
source,
}),
}
}
fn create_root(root: &Path) -> Result<(), StorageError> {
if let Some(parent) = root.parent() {
fs::create_dir_all(parent).map_err(|source| StorageError::Io {
operation: "create recording parent directories",
path: parent.to_path_buf(),
source,
})?;
}
match fs::create_dir(root) {
Ok(()) => Ok(()),
Err(source) if source.kind() == std::io::ErrorKind::AlreadyExists => {
Err(StorageError::RecordingDirectoryExists {
path: root.to_path_buf(),
})
}
Err(source) => Err(StorageError::Io {
operation: "create output root",
path: root.to_path_buf(),
source,
}),
}
}
fn commit_metadata(
root: &Path,
metadata_path: &Path,
metadata: &RecordingMetadata,
) -> Result<(), StorageError> {
metadata.validate(metadata_path)?;
let mut bytes = serde_json::to_vec_pretty(metadata).map_err(|source| StorageError::Json {
operation: "serialize metadata",
path: metadata_path.to_path_buf(),
source,
})?;
bytes.push(b'\n');
let temporary_path = root.join(METADATA_TEMP_FILE);
let result = write_and_replace_metadata(root, metadata_path, &temporary_path, &bytes);
if result.is_err() {
let _ = fs::remove_file(&temporary_path);
}
result
}
fn write_and_replace_metadata(
root: &Path,
metadata_path: &Path,
temporary_path: &Path,
bytes: &[u8],
) -> Result<(), StorageError> {
let mut temporary = OpenOptions::new()
.write(true)
.create_new(true)
.open(temporary_path)
.map_err(|source| StorageError::Io {
operation: "create temporary metadata",
path: temporary_path.to_path_buf(),
source,
})?;
temporary
.write_all(bytes)
.map_err(|source| StorageError::Io {
operation: "write temporary metadata",
path: temporary_path.to_path_buf(),
source,
})?;
temporary.sync_all().map_err(|source| StorageError::Io {
operation: "sync temporary metadata",
path: temporary_path.to_path_buf(),
source,
})?;
drop(temporary);
fs::rename(temporary_path, metadata_path).map_err(|source| StorageError::Io {
operation: "publish metadata",
path: metadata_path.to_path_buf(),
source,
})?;
File::open(root)
.and_then(|directory| directory.sync_all())
.map_err(|source| StorageError::Io {
operation: "sync output root",
path: root.to_path_buf(),
source,
})
}