use std::collections::HashMap;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::path::{Path, PathBuf};
use anyhow::{Context, Result, anyhow, bail, ensure};
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use rustc_hash::FxHashMap;
use serde::{Deserialize, Serialize};
use uuid::Uuid;
use super::driver::WorkloadDriver;
use super::types::{
AGENTIC_MOONCAKE_SCHEMA, AGENTIC_MOONCAKE_VERSION, AgenticDependency,
AgenticDependencyRelation, AgenticDependencyTrigger, AgenticHashIdScope, AgenticMooncakeHeader,
AgenticMooncakeRow, AgenticNode, AgenticPlay, AgenticSourceProvenance, AgenticTrace, DelaySpec,
LengthSpec, MooncakeRow, ReplayRequestHashes, SessionPartitionSpec, SessionTrace,
SyntheticTraceSpec, Trace, TraceFileFormat, TurnTrace, effective_replay_key,
};
use super::{SYNTHETIC_OUTPUT_SEED, planned_output_token_ids};
use crate::replay::protocol::DirectRequest;
#[derive(Debug, Deserialize)]
struct RawAppliedComputeAgenticRecord {
num_turns: usize,
input_prompt_length: usize,
assistant_response_length: Vec<usize>,
tool_call_output_length: Vec<usize>,
tool_call_latency: Vec<f64>,
final_assistant_response_length: usize,
}
#[derive(Debug, Deserialize)]
struct LegacyAgenticMooncakeRow {
request_id: String,
#[serde(default)]
session_id: Option<String>,
#[serde(default, alias = "input_tokens")]
input_length: Option<usize>,
#[serde(alias = "output_tokens")]
output_length: usize,
#[serde(default)]
output_token_ids: Option<Vec<u32>>,
hash_ids: Vec<u64>,
#[serde(default, alias = "created_time")]
timestamp: Option<f64>,
#[serde(default)]
delay: Option<f64>,
#[serde(default)]
delay_ms: Option<f64>,
#[serde(default)]
tool_wait_ms: f64,
#[serde(default)]
wait_for: Vec<String>,
#[serde(default)]
priority: Option<i32>,
#[serde(default)]
strict_priority: Option<u32>,
#[serde(default)]
policy_class: Option<String>,
}
pub fn load_agentic_mooncake(path: &Path, legacy_trace_block_size: usize) -> Result<AgenticTrace> {
let file = File::open(path)
.with_context(|| format!("failed to open trace file {}", path.display()))?;
let mut lines = BufReader::new(file).lines();
let first = loop {
let line = lines
.next()
.transpose()
.with_context(|| format!("failed to read trace file {}", path.display()))?
.context("agentic trace file is empty")?;
if !line.trim().is_empty() {
break line;
}
};
let first_json: serde_json::Value =
serde_json::from_str(&first).context("failed to parse first agentic trace row")?;
if first_json.get("schema").and_then(serde_json::Value::as_str) == Some(AGENTIC_MOONCAKE_SCHEMA)
{
return AgenticTrace::from_agentic_mooncake(path);
}
let mut raw_rows = vec![
serde_json::from_value::<LegacyAgenticMooncakeRow>(first_json)
.context("failed to parse first legacy agentic Mooncake row")?,
];
for (line_index, line) in lines.enumerate() {
let line = line.with_context(|| {
format!(
"failed to read legacy agentic trace line {}",
line_index + 2
)
})?;
if line.trim().is_empty() {
continue;
}
raw_rows.push(serde_json::from_str(&line).with_context(|| {
format!(
"failed to parse legacy agentic trace line {}",
line_index + 2
)
})?);
}
let mut rows = raw_rows
.into_iter()
.map(|raw| -> Result<AgenticMooncakeRow> {
if raw.request_id.trim().is_empty() {
bail!("request_id must be nonempty");
}
if raw.hash_ids.is_empty() {
bail!("hash_ids must be nonempty");
}
if !raw.tool_wait_ms.is_finite() || raw.tool_wait_ms < 0.0 {
bail!("tool_wait_ms must be finite and nonnegative");
}
if raw.delay.is_some() && raw.delay_ms.is_some() {
bail!("delay and delay_ms cannot both be set");
}
let delay = raw.delay.or(raw.delay_ms).unwrap_or(0.0) + raw.tool_wait_ms;
if !delay.is_finite() || delay < 0.0 {
bail!("dependency delay must be finite and nonnegative");
}
let relation = if raw.wait_for.len() > 1 {
AgenticDependencyRelation::Join
} else {
AgenticDependencyRelation::Sequence
};
let dependencies = raw
.wait_for
.into_iter()
.map(|request_id| AgenticDependency {
request_id,
trigger: AgenticDependencyTrigger::Completion,
delay_ms: delay,
relation,
})
.collect::<Vec<_>>();
Ok(AgenticMooncakeRow {
request_id: raw.request_id,
play_id: "agentic-play".to_string(),
session_id: raw
.session_id
.unwrap_or_else(|| "agentic-session".to_string()),
model: "unknown".to_string(),
input_length: raw.input_length,
output_length: Some(raw.output_length),
output_token_ids: raw.output_token_ids,
hash_ids: Some(raw.hash_ids),
not_before_ms: if dependencies.is_empty() {
raw.timestamp.unwrap_or(0.0)
} else {
0.0
},
priority: raw.priority,
strict_priority: raw.strict_priority,
policy_class: raw.policy_class,
dependencies,
})
})
.collect::<Result<Vec<_>>>()?;
assign_dependency_component_play_ids(&mut rows, "legacy-play");
AgenticTrace::from_agentic_mooncake_rows(
AgenticMooncakeHeader {
schema: AGENTIC_MOONCAKE_SCHEMA.to_string(),
version: AGENTIC_MOONCAKE_VERSION,
block_size: legacy_trace_block_size,
hash_id_scope: AgenticHashIdScope::Local,
source: AgenticSourceProvenance {
format: "legacy_agentic_mooncake".to_string(),
digest: format!("{}:{}", path.display(), rows.len()),
},
},
rows,
)
}
pub(super) fn assign_dependency_component_play_ids(rows: &mut [AgenticMooncakeRow], prefix: &str) {
fn find(parent: &mut [usize], value: usize) -> usize {
if parent[value] != value {
parent[value] = find(parent, parent[value]);
}
parent[value]
}
let by_id = rows
.iter()
.enumerate()
.map(|(index, row)| (row.request_id.clone(), index))
.collect::<HashMap<_, _>>();
let mut parent = (0..rows.len()).collect::<Vec<_>>();
for (index, row) in rows.iter().enumerate() {
for dependency in &row.dependencies {
let Some(&source) = by_id.get(&dependency.request_id) else {
continue;
};
let left = find(&mut parent, index);
let right = find(&mut parent, source);
if left != right {
parent[left] = right;
}
}
}
let mut roots_by_component: HashMap<usize, Vec<usize>> = HashMap::new();
for (index, row) in rows.iter().enumerate() {
if row.dependencies.is_empty() {
let component = find(&mut parent, index);
roots_by_component.entry(component).or_default().push(index);
}
}
let mut labels = HashMap::new();
for (component, roots) in roots_by_component {
let canonical = roots
.iter()
.copied()
.min_by(|left, right| {
rows[*left]
.not_before_ms
.total_cmp(&rows[*right].not_before_ms)
.then_with(|| rows[*left].request_id.cmp(&rows[*right].request_id))
})
.expect("a root component is nonempty");
labels.insert(component, rows[canonical].request_id.clone());
}
for (index, row) in rows.iter_mut().enumerate() {
let component = find(&mut parent, index);
let label = labels
.entry(component)
.or_insert_with(|| row.request_id.clone());
row.play_id = format!("{prefix}:{label}");
}
}
#[derive(Debug, Default)]
struct HashIdInterner {
canonical_ids: FxHashMap<u64, u32>,
}
impl HashIdInterner {
fn intern_all(&mut self, hash_ids: Vec<u64>) -> Result<Vec<u32>> {
hash_ids
.into_iter()
.map(|hash_id| self.intern(hash_id))
.collect()
}
fn intern(&mut self, hash_id: u64) -> Result<u32> {
let next_id = self.canonical_ids.len();
match self.canonical_ids.entry(hash_id) {
std::collections::hash_map::Entry::Occupied(entry) => Ok(*entry.get()),
std::collections::hash_map::Entry::Vacant(entry) => {
let canonical_id = u32::try_from(next_id)
.context("trace contains more unique hash IDs than u32 can represent")?;
entry.insert(canonical_id);
Ok(canonical_id)
}
}
}
}
pub fn validate_trace_files(format: TraceFileFormat, paths: &[PathBuf]) -> Result<()> {
if paths.is_empty() {
bail!("trace replay requires at least one trace file");
}
if format != TraceFileFormat::Dynamo && paths.len() != 1 {
bail!(
"trace_format='{}' requires exactly one trace file, got {}",
format.as_str(),
paths.len()
);
}
Ok(())
}
fn single_turn_request_uuid(_request_ordinal: usize) -> Uuid {
Uuid::new_v4()
}
pub(super) fn validate_synthesizable_prompt(
input_length: usize,
hash_ids: &[u32],
trace_block_size: usize,
) -> Result<()> {
if trace_block_size == 0 {
bail!("trace_block_size must be greater than 0");
}
let synthesizable_capacity = hash_ids
.len()
.checked_mul(trace_block_size)
.context("synthesized prompt capacity overflow")?;
let required_hash_ids = input_length.div_ceil(trace_block_size);
if hash_ids.len() < required_hash_ids {
bail!(
"input_length {} exceeds synthesized capacity {}",
input_length,
synthesizable_capacity
);
}
Ok(())
}
pub(super) fn synthesize_trace_tokens(
input_length: usize,
hash_ids: &[u32],
trace_block_size: usize,
) -> Result<Vec<u32>> {
validate_synthesizable_prompt(input_length, hash_ids, trace_block_size)?;
Ok(synthesize_validated_trace_tokens(
input_length,
hash_ids,
trace_block_size,
))
}
pub(super) fn synthesize_validated_trace_tokens(
input_length: usize,
hash_ids: &[u32],
trace_block_size: usize,
) -> Vec<u32> {
let mut tokens = Vec::with_capacity(input_length);
for &hash_id in hash_ids {
let remaining = input_length - tokens.len();
tokens.extend(std::iter::repeat_n(
hash_id,
remaining.min(trace_block_size),
));
if tokens.len() == input_length {
break;
}
}
debug_assert_eq!(tokens.len(), input_length);
tokens
}
fn trace_to_replay_hashes(
input_length: usize,
hash_ids: &[u32],
trace_block_size: usize,
engine_block_size: usize,
) -> Result<ReplayRequestHashes> {
if engine_block_size == 0 {
bail!("engine_block_size must be greater than 0");
}
let tokens = synthesize_trace_tokens(input_length, hash_ids, trace_block_size)?;
let engine_block_size =
u32::try_from(engine_block_size).context("engine_block_size does not fit in u32")?;
Ok(ReplayRequestHashes::from_tokens(&tokens, engine_block_size))
}
impl TurnTrace {
pub fn synthesize_tokens(&self, trace_block_size: usize) -> Result<Vec<u32>> {
synthesize_trace_tokens(self.input_length, &self.hash_ids, trace_block_size)
}
pub fn to_direct_request(
&self,
trace_block_size: usize,
request_uuid: Uuid,
arrival_timestamp_ms: Option<f64>,
) -> Result<DirectRequest> {
let tokens = self.synthesize_tokens(trace_block_size)?;
Ok(DirectRequest {
tokens,
max_output_tokens: self.max_output_tokens,
output_token_ids: self.output_token_ids.clone(),
uuid: Some(request_uuid),
dp_rank: 0,
preferred_dp_rank: None,
arrival_timestamp_ms,
priority: self.priority,
strict_priority: self.strict_priority,
policy_class: self.policy_class.clone(),
replay_context: None,
})
}
pub fn to_replay_hashes(
&self,
trace_block_size: usize,
engine_block_size: usize,
) -> Result<ReplayRequestHashes> {
trace_to_replay_hashes(
self.input_length,
&self.hash_ids,
trace_block_size,
engine_block_size,
)
}
}
struct MooncakeTraceBuilder {
trace_block_size: usize,
hash_id_interner: HashIdInterner,
sessions: Vec<SessionTrace>,
session_indices: HashMap<String, usize>,
last_timestamps: Vec<Option<f64>>,
}
impl MooncakeTraceBuilder {
fn new(trace_block_size: usize) -> Self {
Self {
trace_block_size,
hash_id_interner: HashIdInterner::default(),
sessions: Vec::new(),
session_indices: HashMap::new(),
last_timestamps: Vec::new(),
}
}
fn is_empty(&self) -> bool {
self.sessions.is_empty()
}
fn push(&mut self, line_idx: usize, raw: MooncakeRow) -> Result<()> {
let request_id = raw.request_id;
let raw_session_id = raw.session_id;
let session_id = raw_session_id
.clone()
.unwrap_or_else(|| format!("request_{}", line_idx + 1));
let hash_ids = raw
.hash_ids
.ok_or_else(|| anyhow!("trace line {} is missing hash_ids", line_idx + 1))?;
let synthesizable_capacity = hash_ids
.len()
.checked_mul(self.trace_block_size)
.ok_or_else(|| anyhow!("trace line {} synthesized capacity overflow", line_idx + 1))?;
let input_length = raw.input_length.unwrap_or(synthesizable_capacity);
ensure!(
input_length <= synthesizable_capacity,
"trace line {} input_length {} exceeds hash_ids capacity {}",
line_idx + 1,
input_length,
synthesizable_capacity
);
let output_length = raw
.output_length
.ok_or_else(|| anyhow!("trace line {} is missing output_length", line_idx + 1))?;
let output_token_ids = raw.output_token_ids;
if let Some(output_token_ids) = output_token_ids.as_ref()
&& output_token_ids.len() != output_length
{
bail!(
"trace line {} output_length {} does not match output_token_ids length {}",
line_idx + 1,
output_length,
output_token_ids.len()
);
}
let timestamp_ms = raw.timestamp;
let explicit_delay_ms = raw.delay;
let priority = raw.priority.unwrap_or(0);
let strict_priority = raw.strict_priority.unwrap_or(0);
let policy_class = raw.policy_class.clone();
let session_index = *self
.session_indices
.entry(session_id.clone())
.or_insert_with(|| {
let idx = self.sessions.len();
self.sessions.push(SessionTrace {
session_id: session_id.clone(),
first_arrival_timestamp_ms: timestamp_ms,
turns: Vec::new(),
});
self.last_timestamps.push(timestamp_ms);
idx
});
let session = self
.sessions
.get_mut(session_index)
.expect("newly inserted session must exist");
let turn_idx = session.turns.len();
let replay_key = output_token_ids.as_ref().map(|_| {
effective_replay_key(
request_id.as_deref(),
raw_session_id.as_deref(),
turn_idx,
line_idx,
)
});
let delay_after_previous_ms = if turn_idx == 0 {
let delay = explicit_delay_ms.unwrap_or(0.0);
if delay != 0.0 {
bail!(
"trace line {} sets delay on the first turn of session {}",
line_idx + 1,
session.session_id
);
}
0.0
} else if let Some(delay_ms) = explicit_delay_ms {
delay_ms
} else if let Some(timestamp_ms) = timestamp_ms {
let previous_timestamp_ms = self.last_timestamps[session_index].ok_or_else(|| {
anyhow!(
"trace line {} for session {} cannot infer delay without a previous timestamp",
line_idx + 1,
session.session_id
)
})?;
timestamp_ms - previous_timestamp_ms
} else {
0.0
};
if !delay_after_previous_ms.is_finite() || delay_after_previous_ms < 0.0 {
bail!(
"trace line {} has invalid delay {}",
line_idx + 1,
delay_after_previous_ms
);
}
if hash_ids.len() * self.trace_block_size < input_length {
bail!(
"trace line {} input_length {} exceeds synthesized capacity {}",
line_idx + 1,
input_length,
hash_ids.len() * self.trace_block_size
);
}
let hash_ids = self.hash_id_interner.intern_all(hash_ids)?;
session.turns.push(TurnTrace {
input_length,
max_output_tokens: output_length,
output_token_ids,
replay_key,
hash_ids,
delay_after_previous_ms,
priority,
strict_priority,
policy_class,
});
if let Some(timestamp_ms) = timestamp_ms {
self.last_timestamps[session_index] = Some(timestamp_ms);
}
Ok(())
}
fn finish(self) -> Trace {
Trace {
block_size: self.trace_block_size,
sessions: self.sessions,
}
}
}
impl Trace {
pub fn from_mooncake(path: &Path, trace_block_size: usize) -> Result<Self> {
if trace_block_size == 0 {
bail!("trace_block_size must be greater than 0");
}
let file = File::open(path)
.with_context(|| format!("failed to open trace file {}", path.display()))?;
let reader = BufReader::new(file);
let mut builder = MooncakeTraceBuilder::new(trace_block_size);
for (line_idx, line) in reader.lines().enumerate() {
let line = line.with_context(|| {
format!(
"failed to read line {} from {}",
line_idx + 1,
path.display()
)
})?;
if line.trim().is_empty() {
continue;
}
let row = serde_json::from_str(&line).with_context(|| {
format!(
"failed to parse line {} from {} as JSON",
line_idx + 1,
path.display()
)
})?;
builder.push(line_idx, row)?;
}
if builder.is_empty() {
bail!("trace file {} did not contain any requests", path.display());
}
Ok(builder.finish())
}
pub fn from_mooncake_rows(rows: Vec<MooncakeRow>, trace_block_size: usize) -> Result<Self> {
if trace_block_size == 0 {
bail!("trace_block_size must be greater than 0");
}
let mut builder = MooncakeTraceBuilder::new(trace_block_size);
for (line_idx, row) in rows.into_iter().enumerate() {
builder.push(line_idx, row)?;
}
if builder.is_empty() {
bail!("Mooncake rows did not contain any requests");
}
Ok(builder.finish())
}
pub fn from_applied_compute_agentic(
path: &Path,
trace_block_size: usize,
shared_prefix_ratio: f64,
num_prefix_groups: usize,
) -> Result<Self> {
if trace_block_size == 0 {
bail!("trace_block_size must be greater than 0");
}
if !(0.0..=1.0).contains(&shared_prefix_ratio) {
bail!(
"shared_prefix_ratio must be between 0.0 and 1.0, got {}",
shared_prefix_ratio
);
}
let file = File::open(path)
.with_context(|| format!("failed to open trace file {}", path.display()))?;
let reader = BufReader::new(file);
let mut sessions = Vec::new();
let mut hash_id_interner = HashIdInterner::default();
let mut next_unique_hash = 1_u64;
for (line_idx, line) in reader.lines().enumerate() {
let line = line.with_context(|| {
format!(
"failed to read line {} from {}",
line_idx + 1,
path.display()
)
})?;
if line.trim().is_empty() {
continue;
}
let raw: RawAppliedComputeAgenticRecord =
serde_json::from_str(&line).with_context(|| {
format!(
"failed to parse line {} from {} as JSON",
line_idx + 1,
path.display()
)
})?;
for (name, values) in [
(
"assistant_response_length",
raw.assistant_response_length.len(),
),
("tool_call_output_length", raw.tool_call_output_length.len()),
("tool_call_latency", raw.tool_call_latency.len()),
] {
if values != raw.num_turns {
bail!(
"trace line {} field {} length {} does not match num_turns {}",
line_idx + 1,
name,
values,
raw.num_turns
);
}
}
if raw.input_prompt_length == 0 {
bail!(
"trace line {} input_prompt_length must be positive",
line_idx + 1
);
}
let group_id = if shared_prefix_ratio > 0.0 && num_prefix_groups > 0 {
Some(line_idx % num_prefix_groups)
} else {
None
};
let mut current_input_length = raw.input_prompt_length;
let mut hash_ids = Vec::new();
let shared_initial_blocks = ((current_input_length.div_ceil(trace_block_size) as f64)
* shared_prefix_ratio)
.round() as usize;
extend_applied_compute_agentic_hash_ids(
&mut hash_ids,
current_input_length,
trace_block_size,
shared_initial_blocks,
group_id,
&mut next_unique_hash,
)?;
let mut turns = Vec::with_capacity(raw.num_turns + 1);
let mut next_turn_delay_ms = 0.0;
for turn_idx in 0..raw.num_turns {
let tool_call_latency = raw.tool_call_latency[turn_idx];
if !tool_call_latency.is_finite() || tool_call_latency < 0.0 {
bail!(
"trace line {} tool_call_latency[{}] must be a finite non-negative number",
line_idx + 1,
turn_idx
);
}
turns.push(TurnTrace {
input_length: current_input_length,
max_output_tokens: raw.assistant_response_length[turn_idx],
hash_ids: hash_id_interner.intern_all(hash_ids.clone())?,
delay_after_previous_ms: next_turn_delay_ms,
..Default::default()
});
current_input_length = current_input_length
.checked_add(raw.assistant_response_length[turn_idx])
.and_then(|value| value.checked_add(raw.tool_call_output_length[turn_idx]))
.ok_or_else(|| {
anyhow!(
"trace line {} cumulative input length overflow",
line_idx + 1
)
})?;
extend_applied_compute_agentic_hash_ids(
&mut hash_ids,
current_input_length,
trace_block_size,
shared_initial_blocks,
group_id,
&mut next_unique_hash,
)?;
next_turn_delay_ms = tool_call_latency * 1000.0;
}
turns.push(TurnTrace {
input_length: current_input_length,
max_output_tokens: raw.final_assistant_response_length,
hash_ids: hash_id_interner.intern_all(hash_ids)?,
delay_after_previous_ms: next_turn_delay_ms,
..Default::default()
});
sessions.push(SessionTrace {
session_id: format!("applied_compute_agentic_session_{}", line_idx + 1),
first_arrival_timestamp_ms: None,
turns,
});
}
if sessions.is_empty() {
bail!("trace file {} did not contain any requests", path.display());
}
Ok(Self {
block_size: trace_block_size,
sessions,
})
}
pub fn synthetic(spec: SyntheticTraceSpec) -> Result<Self> {
if spec.block_size == 0 {
bail!("block_size must be greater than 0");
}
if spec.num_sessions == 0 {
bail!("num_sessions must be greater than 0");
}
if spec.turns_per_session == 0 {
bail!("turns_per_session must be greater than 0");
}
if !(0.0..=1.0).contains(&spec.shared_prefix_ratio) {
bail!(
"shared_prefix_ratio must be between 0.0 and 1.0, got {}",
spec.shared_prefix_ratio
);
}
let mut rng = StdRng::seed_from_u64(spec.seed);
let mut sessions = Vec::with_capacity(spec.num_sessions);
let first_arrivals = spec
.first_turn_arrivals
.timestamps(spec.num_sessions, spec.arrival_seed)?;
let mut next_unique_hash = 1_u64;
let mut hash_id_interner = HashIdInterner::default();
for (session_idx, first_arrival_timestamp_ms) in first_arrivals.into_iter().enumerate() {
let group_id = if spec.num_prefix_groups > 0 && spec.shared_prefix_ratio > 0.0 {
Some(rng.random_range(0..spec.num_prefix_groups) as u64)
} else {
None
};
let mut turns = Vec::with_capacity(spec.turns_per_session);
for turn_idx in 0..spec.turns_per_session {
let input_length = sample_length(&spec.input_tokens, 1, &mut rng);
let max_output_tokens = sample_length(&spec.output_tokens, 1, &mut rng);
let num_blocks = input_length.div_ceil(spec.block_size);
let prefix_blocks =
((num_blocks as f64) * spec.shared_prefix_ratio).round() as usize;
let prefix_blocks = prefix_blocks.min(num_blocks);
let mut hash_ids = Vec::with_capacity(num_blocks);
for block_idx in 0..prefix_blocks {
if let Some(group_id) = group_id {
hash_ids.push(0xD00D_0000_0000_0000 | (group_id << 32) | block_idx as u64);
}
}
while hash_ids.len() < num_blocks {
hash_ids.push(next_unique_hash);
next_unique_hash = next_unique_hash
.checked_add(1)
.expect("synthetic hash id overflow");
}
turns.push(TurnTrace {
input_length,
max_output_tokens,
hash_ids: hash_id_interner.intern_all(hash_ids)?,
delay_after_previous_ms: if turn_idx == 0 {
0.0
} else {
sample_delay_ms(&spec.inter_turn_delays, &mut rng)?
},
..Default::default()
});
}
sessions.push(SessionTrace {
session_id: format!("session_{session_idx}"),
first_arrival_timestamp_ms: Some(first_arrival_timestamp_ms),
turns,
});
}
Ok(Self {
block_size: spec.block_size,
sessions,
})
}
pub fn validate_for_trace_mode(&self) -> Result<()> {
self.validate(false)
}
pub fn validate_for_concurrency_mode(&self) -> Result<()> {
self.validate(true)
}
pub fn normalize_session_starts(mut self) -> Result<Self> {
let Some(min_timestamp_ms) = self
.sessions
.iter()
.filter_map(|session| session.first_arrival_timestamp_ms)
.min_by(|left, right| left.total_cmp(right))
else {
return Ok(self);
};
for session in &mut self.sessions {
if let Some(timestamp_ms) = session.first_arrival_timestamp_ms.as_mut() {
*timestamp_ms -= min_timestamp_ms;
}
}
Ok(self)
}
pub fn speed_up_timing(mut self, ratio: f64) -> Result<Self> {
if !ratio.is_finite() || ratio <= 0.0 {
bail!("ratio must be a finite positive number, got {ratio}");
}
for session in &mut self.sessions {
if let Some(timestamp_ms) = session.first_arrival_timestamp_ms.as_mut() {
*timestamp_ms /= ratio;
}
for turn in &mut session.turns {
turn.delay_after_previous_ms /= ratio;
}
}
Ok(self)
}
pub fn rescale_session_start_span(mut self, duration_ms: u64) -> Result<Self> {
let Some(min_timestamp_ms) = self
.sessions
.iter()
.filter_map(|session| session.first_arrival_timestamp_ms)
.min_by(|left, right| left.total_cmp(right))
else {
return Ok(self);
};
let Some(max_timestamp_ms) = self
.sessions
.iter()
.filter_map(|session| session.first_arrival_timestamp_ms)
.max_by(|left, right| left.total_cmp(right))
else {
return Ok(self);
};
let target_span_ms = duration_ms as f64;
let source_span_ms = max_timestamp_ms - min_timestamp_ms;
for session in &mut self.sessions {
if let Some(timestamp_ms) = session.first_arrival_timestamp_ms.as_mut() {
*timestamp_ms = if source_span_ms == 0.0 {
0.0
} else {
(*timestamp_ms - min_timestamp_ms) * target_span_ms / source_span_ms
};
}
}
Ok(self)
}
pub fn rescale_ready_span(mut self, duration_ms: u64) -> Result<Self> {
let Some(min_start_ms) = self
.sessions
.iter()
.map(|session| session.first_arrival_timestamp_ms.unwrap_or(0.0))
.min_by(|left, right| left.total_cmp(right))
else {
return Ok(self);
};
let Some(max_ready_ms) = self
.sessions
.iter()
.map(|session| {
session.first_arrival_timestamp_ms.unwrap_or(0.0)
+ session
.turns
.iter()
.enumerate()
.filter(|(turn_idx, _)| *turn_idx > 0)
.map(|(_, turn)| turn.delay_after_previous_ms)
.sum::<f64>()
})
.max_by(|left, right| left.total_cmp(right))
else {
return Ok(self);
};
let ratio = duration_ms as f64 / (max_ready_ms - min_start_ms).max(1.0);
for session in &mut self.sessions {
if let Some(start_ms) = session.first_arrival_timestamp_ms.as_mut() {
*start_ms = (*start_ms - min_start_ms) * ratio;
}
for (turn_idx, turn) in session.turns.iter_mut().enumerate() {
if turn_idx > 0 {
turn.delay_after_previous_ms *= ratio;
}
}
}
Ok(self)
}
pub fn expand_hash_prefix_depth(mut self, factor: usize) -> Self {
if factor <= 1 {
return self;
}
let factor = u32::try_from(factor).expect("hash prefix expansion factor exceeds u32");
for session in &mut self.sessions {
for turn in &mut session.turns {
turn.input_length = turn
.input_length
.checked_mul(factor as usize)
.expect("input_length expansion overflow");
turn.hash_ids = turn
.hash_ids
.iter()
.flat_map(|&hash_id| {
let base = hash_id
.checked_mul(factor)
.expect("hash prefix expansion overflow");
(0..factor).map(move |offset| {
base.checked_add(offset)
.expect("hash prefix expansion overflow")
})
})
.collect();
}
}
self
}
pub fn duplicate_hash_space(mut self, copies: usize) -> Self {
if copies <= 1 {
return self;
}
let max_hash_id = self
.sessions
.iter()
.flat_map(|session| session.turns.iter())
.flat_map(|turn| turn.hash_ids.iter().copied())
.max()
.unwrap_or(0);
let offset_base = max_hash_id
.checked_add(1)
.expect("hash duplication offset overflow");
let original_sessions = self.sessions.clone();
self.sessions.clear();
for copy_idx in 0..copies {
let copy_idx = u32::try_from(copy_idx).expect("hash copy index exceeds u32");
let offset = offset_base
.checked_mul(copy_idx)
.expect("hash duplication offset overflow");
for session in &original_sessions {
let mut duplicated = session.clone();
duplicated.session_id = format!("{}:copy_{copy_idx}", session.session_id);
for turn in &mut duplicated.turns {
turn.hash_ids = turn
.hash_ids
.iter()
.map(|&hash_id| {
hash_id
.checked_add(offset)
.expect("hash duplication overflow")
})
.collect();
}
self.sessions.push(duplicated);
}
}
self
}
pub fn partition_by_session(&self, spec: SessionPartitionSpec) -> Vec<Self> {
let num_partitions = match spec {
SessionPartitionSpec::Random { num_partitions, .. } => num_partitions,
SessionPartitionSpec::RoundRobin { num_partitions } => num_partitions,
}
.max(1);
let mut partitions = vec![
Self {
block_size: self.block_size,
sessions: Vec::new(),
};
num_partitions
];
let mut rng = match spec {
SessionPartitionSpec::Random { seed, .. } => Some(StdRng::seed_from_u64(seed)),
SessionPartitionSpec::RoundRobin { .. } => None,
};
for (session_idx, session) in self.sessions.iter().cloned().enumerate() {
let partition_idx = match spec {
SessionPartitionSpec::Random { .. } => rng
.as_mut()
.expect("random partitioner must exist")
.random_range(0..num_partitions),
SessionPartitionSpec::RoundRobin { .. } => session_idx % num_partitions,
};
partitions[partition_idx].sessions.push(session);
}
partitions
}
pub fn to_single_turn_requests(&self) -> Result<Vec<DirectRequest>> {
let mut requests = Vec::with_capacity(self.sessions.len());
let mut output_rng = StdRng::seed_from_u64(SYNTHETIC_OUTPUT_SEED);
for (request_ordinal, session) in self.sessions.iter().enumerate() {
if session.turns.len() != 1 {
bail!(
"to_single_turn_requests requires exactly one turn per session, but session {} has {} turns",
session.session_id,
session.turns.len()
);
}
let request_uuid = single_turn_request_uuid(request_ordinal);
let mut request = session.turns[0].to_direct_request(
self.block_size,
request_uuid,
session.first_arrival_timestamp_ms,
)?;
request.output_token_ids = Some(planned_output_token_ids(
request.output_token_ids,
request.max_output_tokens,
&mut output_rng,
));
requests.push(request);
}
Ok(requests)
}
pub fn is_single_turn(&self) -> bool {
self.sessions.iter().all(|session| session.turns.len() == 1)
}
pub fn into_trace_driver(self) -> Result<WorkloadDriver> {
self.validate_for_trace_mode()?;
let engine_block_size = self.block_size;
WorkloadDriver::new_trace(self, engine_block_size)
}
pub fn into_concurrency_driver(self, max_in_flight: usize) -> Result<WorkloadDriver> {
self.validate_for_concurrency_mode()?;
let engine_block_size = self.block_size;
WorkloadDriver::new_concurrency(self, engine_block_size, max_in_flight)
}
pub fn into_trace_driver_with_block_size(
self,
engine_block_size: usize,
) -> Result<WorkloadDriver> {
self.validate_for_trace_mode()?;
WorkloadDriver::new_trace(self, engine_block_size)
}
pub fn into_delta_accumulating_trace_driver_with_block_size(
self,
engine_block_size: usize,
) -> Result<WorkloadDriver> {
self.validate_for_trace_mode()?;
WorkloadDriver::new_trace_accumulating_deltas(self, engine_block_size)
}
pub fn into_concurrency_driver_with_block_size(
self,
engine_block_size: usize,
max_in_flight: usize,
) -> Result<WorkloadDriver> {
self.validate_for_concurrency_mode()?;
WorkloadDriver::new_concurrency(self, engine_block_size, max_in_flight)
}
pub fn into_delta_accumulating_concurrency_driver_with_block_size(
self,
engine_block_size: usize,
max_in_flight: usize,
) -> Result<WorkloadDriver> {
self.validate_for_concurrency_mode()?;
WorkloadDriver::new_concurrency_accumulating_deltas(self, engine_block_size, max_in_flight)
}
fn validate(&self, allow_missing_first_timestamp: bool) -> Result<()> {
if self.block_size == 0 {
bail!("block_size must be greater than 0");
}
if self.sessions.is_empty() {
bail!("trace must contain at least one session");
}
for session in &self.sessions {
if session.turns.is_empty() {
bail!(
"session {} must contain at least one turn",
session.session_id
);
}
if !allow_missing_first_timestamp {
let timestamp_ms = session.first_arrival_timestamp_ms.ok_or_else(|| {
anyhow!(
"trace mode requires first_arrival_timestamp_ms for session {}",
session.session_id
)
})?;
if !timestamp_ms.is_finite() || timestamp_ms < 0.0 {
bail!(
"session {} has invalid first_arrival_timestamp_ms {}",
session.session_id,
timestamp_ms
);
}
} else if let Some(timestamp_ms) = session.first_arrival_timestamp_ms
&& (!timestamp_ms.is_finite() || timestamp_ms < 0.0)
{
bail!(
"session {} has invalid first_arrival_timestamp_ms {}",
session.session_id,
timestamp_ms
);
}
for (turn_idx, turn) in session.turns.iter().enumerate() {
if let Some(output_token_ids) = turn.output_token_ids.as_ref()
&& output_token_ids.len() != turn.max_output_tokens
{
bail!(
"session {} turn {} max_output_tokens {} does not match output_token_ids length {}",
session.session_id,
turn_idx,
turn.max_output_tokens,
output_token_ids.len()
);
}
if turn.input_length == 0 {
bail!(
"session {} turn {} must have a positive input_length",
session.session_id,
turn_idx
);
}
if turn.hash_ids.is_empty() {
bail!(
"session {} turn {} must contain at least one hash id",
session.session_id,
turn_idx
);
}
validate_synthesizable_prompt(turn.input_length, &turn.hash_ids, self.block_size)
.with_context(|| {
format!(
"session {} turn {} has invalid prompt",
session.session_id, turn_idx
)
})?;
if !turn.delay_after_previous_ms.is_finite() || turn.delay_after_previous_ms < 0.0 {
bail!(
"session {} turn {} has invalid delay {}",
session.session_id,
turn_idx,
turn.delay_after_previous_ms
);
}
if turn_idx == 0 && turn.delay_after_previous_ms != 0.0 {
bail!(
"session {} first turn must have delay_after_previous_ms == 0.0",
session.session_id
);
}
}
}
Ok(())
}
}
struct AgenticTraceBuilder {
header: AgenticMooncakeHeader,
nodes: Vec<AgenticNode>,
request_ids: std::collections::HashSet<String>,
}
pub struct AgenticGraphBuilder {
inner: AgenticTraceBuilder,
row_count: usize,
}
impl AgenticGraphBuilder {
pub fn new(header: AgenticMooncakeHeader) -> Result<Self> {
Ok(Self {
inner: AgenticTraceBuilder::new(header)?,
row_count: 0,
})
}
pub fn push(&mut self, row: AgenticMooncakeRow) -> Result<()> {
self.inner.push(self.row_count + 1, row)?;
self.row_count += 1;
Ok(())
}
pub fn finish(self) -> Result<AgenticTrace> {
if self.inner.is_empty() {
bail!("agentic Mooncake rows did not contain any requests");
}
self.inner.finish()
}
}
impl AgenticTraceBuilder {
fn new(header: AgenticMooncakeHeader) -> Result<Self> {
if header.schema != AGENTIC_MOONCAKE_SCHEMA {
bail!(
"unsupported agentic Mooncake schema {:?}; expected {:?}",
header.schema,
AGENTIC_MOONCAKE_SCHEMA
);
}
if header.version != AGENTIC_MOONCAKE_VERSION {
bail!(
"unsupported agentic Mooncake version {}; expected {}",
header.version,
AGENTIC_MOONCAKE_VERSION
);
}
if header.block_size == 0 {
bail!("agentic Mooncake block_size must be greater than 0");
}
if header.source.format.trim().is_empty() || header.source.digest.trim().is_empty() {
bail!("agentic Mooncake source provenance requires nonempty format and digest");
}
Ok(Self {
header,
nodes: Vec::new(),
request_ids: std::collections::HashSet::new(),
})
}
fn is_empty(&self) -> bool {
self.nodes.is_empty()
}
fn push(&mut self, line_idx: usize, raw: AgenticMooncakeRow) -> Result<()> {
if raw.request_id.trim().is_empty() {
bail!("trace line {} has empty request_id", line_idx + 1);
}
if raw.play_id.trim().is_empty() {
bail!("trace line {} has empty play_id", line_idx + 1);
}
if raw.session_id.trim().is_empty() {
bail!("trace line {} has empty session_id", line_idx + 1);
}
if !self.request_ids.insert(raw.request_id.clone()) {
bail!(
"trace line {} duplicates request_id {}",
line_idx + 1,
raw.request_id
);
}
let hash_ids = raw
.hash_ids
.ok_or_else(|| anyhow!("trace line {} is missing hash_ids", line_idx + 1))?;
if hash_ids.is_empty() {
bail!("trace line {} has empty hash_ids", line_idx + 1);
}
let input_length = raw
.input_length
.ok_or_else(|| anyhow!("trace line {} is missing input_length", line_idx + 1))?;
if input_length == 0 {
bail!("trace line {} has zero input_length", line_idx + 1);
}
let expected_hashes = input_length.div_ceil(self.header.block_size);
if hash_ids.len() != expected_hashes {
bail!(
"trace line {} has input_length {} at block_size {} and requires exactly {} hash_ids, got {}",
line_idx + 1,
input_length,
self.header.block_size,
expected_hashes,
hash_ids.len()
);
}
if raw.model.trim().is_empty() {
bail!("trace line {} has an empty model", line_idx + 1);
}
let output_length = raw
.output_length
.ok_or_else(|| anyhow!("trace line {} is missing output_length", line_idx + 1))?;
let output_token_ids = raw.output_token_ids;
if let Some(output_token_ids) = output_token_ids.as_ref()
&& output_token_ids.len() != output_length
{
bail!(
"trace line {} output_length {} does not match output_token_ids length {}",
line_idx + 1,
output_length,
output_token_ids.len()
);
}
if !raw.not_before_ms.is_finite() || raw.not_before_ms < 0.0 {
bail!(
"trace line {} has invalid not_before_ms {}",
line_idx + 1,
raw.not_before_ms
);
}
for dependency in &raw.dependencies {
if dependency.request_id.trim().is_empty() {
bail!(
"trace line {} has a dependency with an empty request_id",
line_idx + 1
);
}
if !dependency.delay_ms.is_finite() || dependency.delay_ms < 0.0 {
bail!(
"trace line {} has invalid dependency delay {}",
line_idx + 1,
dependency.delay_ms
);
}
match (dependency.relation, dependency.trigger) {
(AgenticDependencyRelation::Sequence, AgenticDependencyTrigger::Completion)
| (AgenticDependencyRelation::Spawn, _)
| (AgenticDependencyRelation::Join, AgenticDependencyTrigger::Completion)
| (
AgenticDependencyRelation::ReplayBarrier,
AgenticDependencyTrigger::Completion,
) => {}
(relation, trigger) => bail!(
"trace line {} has invalid {:?} dependency with {:?} trigger",
line_idx + 1,
relation,
trigger
),
}
}
let replay_key = output_token_ids.as_ref().map(|_| {
effective_replay_key(
Some(raw.request_id.as_str()),
Some(raw.session_id.as_str()),
0,
line_idx,
)
});
self.nodes.push(AgenticNode {
replay_key,
request_id: raw.request_id,
play_id: raw.play_id,
session_id: raw.session_id,
model: raw.model,
input_length,
max_output_tokens: output_length,
output_token_ids,
hash_ids,
not_before_ms: raw.not_before_ms,
priority: raw.priority.unwrap_or(0),
strict_priority: raw.strict_priority.unwrap_or(0),
policy_class: raw.policy_class,
dependencies: raw.dependencies,
});
Ok(())
}
fn finish(mut self) -> Result<AgenticTrace> {
self.nodes
.sort_by(|left, right| left.request_id.cmp(&right.request_id));
let index_by_id: HashMap<&str, usize> = self
.nodes
.iter()
.enumerate()
.map(|(index, node)| (node.request_id.as_str(), index))
.collect();
let mut nodes_by_play: HashMap<String, Vec<usize>> = HashMap::new();
for (node_index, node) in self.nodes.iter().enumerate() {
nodes_by_play
.entry(node.play_id.clone())
.or_default()
.push(node_index);
for dependency in &node.dependencies {
let Some(&dependency_index) = index_by_id.get(dependency.request_id.as_str())
else {
bail!(
"request {} references unknown request_id {}",
node.request_id,
dependency.request_id
);
};
let dependency_node = &self.nodes[dependency_index];
if dependency_index == node_index {
bail!("request {} cannot depend on itself", node.request_id);
}
if dependency_node.play_id != node.play_id {
bail!(
"request {} in play {} depends on request {} in play {}",
node.request_id,
node.play_id,
dependency.request_id,
dependency_node.play_id
);
}
}
}
validate_agentic_trace_is_acyclic(&self.nodes, &index_by_id)?;
let mut plays = Vec::with_capacity(nodes_by_play.len());
for (play_id, mut node_indices) in nodes_by_play {
node_indices.sort_unstable();
let roots: Vec<_> = node_indices
.iter()
.copied()
.filter(|node_index| self.nodes[*node_index].dependencies.is_empty())
.collect();
if roots.is_empty() {
bail!(
"play {} must have at least one root request, found {}",
play_id,
roots.len()
);
}
plays.push(AgenticPlay {
play_id,
root_nodes: roots,
nodes: node_indices,
});
}
plays.sort_by(|left, right| left.play_id.cmp(&right.play_id));
let graph_digest = canonical_agentic_graph_digest(self.header.block_size, &mut self.nodes)?;
Ok(AgenticTrace {
block_size: self.header.block_size,
source: self.header.source,
graph_digest,
nodes: self.nodes,
plays,
})
}
}
impl AgenticTrace {
pub fn from_agentic_mooncake(path: &Path) -> Result<Self> {
let file = File::open(path)
.with_context(|| format!("failed to open trace file {}", path.display()))?;
let reader = BufReader::new(file);
let mut builder = None;
for (line_idx, line) in reader.lines().enumerate() {
let line = line.with_context(|| {
format!(
"failed to read line {} from {}",
line_idx + 1,
path.display()
)
})?;
if line.trim().is_empty() {
continue;
}
if builder.is_none() {
let header = serde_json::from_str(&line).with_context(|| {
format!(
"failed to parse line {} from {} as the agentic Mooncake v2 header",
line_idx + 1,
path.display()
)
})?;
builder = Some(AgenticTraceBuilder::new(header)?);
continue;
}
let row = serde_json::from_str(&line).with_context(|| {
format!(
"failed to parse line {} from {} as an agentic Mooncake v2 request",
line_idx + 1,
path.display()
)
})?;
builder
.as_mut()
.expect("builder was initialized from the v2 header")
.push(line_idx, row)?;
}
let Some(builder) = builder else {
bail!(
"agentic trace file {} is missing its v2 header",
path.display()
);
};
if builder.is_empty() {
bail!(
"agentic trace file {} did not contain any requests",
path.display()
);
}
builder.finish()
}
pub fn from_agentic_mooncake_rows(
header: AgenticMooncakeHeader,
rows: Vec<AgenticMooncakeRow>,
) -> Result<Self> {
let mut builder = AgenticTraceBuilder::new(header)?;
for (line_idx, row) in rows.into_iter().enumerate() {
builder.push(line_idx + 1, row)?;
}
if builder.is_empty() {
bail!("agentic Mooncake rows did not contain any requests");
}
builder.finish()
}
pub fn normalize_starts(mut self) -> Self {
let min_timestamp_ms = self
.nodes
.iter()
.map(|node| node.not_before_ms)
.min_by(|left, right| left.total_cmp(right))
.unwrap_or(0.0);
for node in &mut self.nodes {
node.not_before_ms -= min_timestamp_ms;
}
self.graph_digest = canonical_agentic_graph_digest(self.block_size, &mut self.nodes)
.expect("validated agentic graph remains serializable after normalization");
self
}
pub fn speed_up_timing(mut self, ratio: f64) -> Result<Self> {
if !ratio.is_finite() || ratio <= 0.0 {
bail!("ratio must be a finite positive number, got {ratio}");
}
for node in &mut self.nodes {
node.not_before_ms /= ratio;
for dependency in &mut node.dependencies {
dependency.delay_ms /= ratio;
}
}
self.graph_digest = canonical_agentic_graph_digest(self.block_size, &mut self.nodes)?;
Ok(self)
}
pub fn into_trace_driver_with_block_size(
self,
engine_block_size: usize,
) -> Result<WorkloadDriver> {
WorkloadDriver::new_agentic_trace(self, engine_block_size)
}
pub fn into_trace_driver_with_options(
self,
engine_block_size: usize,
include_replay_hashes: bool,
agentic_lanes: Option<usize>,
) -> Result<WorkloadDriver> {
WorkloadDriver::new_agentic_trace_with_options(
self,
engine_block_size,
include_replay_hashes,
agentic_lanes,
)
}
}
fn validate_agentic_trace_is_acyclic(
nodes: &[AgenticNode],
index_by_id: &HashMap<&str, usize>,
) -> Result<()> {
let mut indegree = nodes
.iter()
.map(|node| node.dependencies.len())
.collect::<Vec<_>>();
let mut dependents = vec![Vec::new(); nodes.len()];
for (node_index, node) in nodes.iter().enumerate() {
for dependency in &node.dependencies {
let dependency_index = *index_by_id
.get(dependency.request_id.as_str())
.expect("dependencies were prevalidated");
dependents[dependency_index].push(node_index);
}
}
let mut ready = std::collections::VecDeque::from_iter(
indegree
.iter()
.enumerate()
.filter_map(|(index, degree)| (*degree == 0).then_some(index)),
);
let mut visited = 0;
while let Some(node_index) = ready.pop_front() {
visited += 1;
for dependent in &dependents[node_index] {
indegree[*dependent] -= 1;
if indegree[*dependent] == 0 {
ready.push_back(*dependent);
}
}
}
if visited == nodes.len() {
return Ok(());
}
let blocked = indegree
.iter()
.enumerate()
.filter_map(|(index, degree)| (*degree > 0).then_some(nodes[index].request_id.as_str()))
.take(8)
.collect::<Vec<_>>();
bail!(
"cycle detected among {} agentic requests; blocked request IDs: {}",
nodes.len() - visited,
blocked.join(", ")
)
}
fn canonical_agentic_graph_digest(block_size: usize, nodes: &mut [AgenticNode]) -> Result<String> {
struct Blake3Writer<'a>(&'a mut blake3::Hasher);
impl std::io::Write for Blake3Writer<'_> {
fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
self.0.update(bytes);
Ok(bytes.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
#[derive(Serialize)]
struct CanonicalGraph<'a> {
block_size: usize,
nodes: &'a [AgenticNode],
}
for node in nodes.iter_mut() {
node.dependencies.sort_by(|left, right| {
left.request_id
.cmp(&right.request_id)
.then_with(|| trigger_rank(left.trigger).cmp(&trigger_rank(right.trigger)))
.then_with(|| relation_rank(left.relation).cmp(&relation_rank(right.relation)))
.then_with(|| left.delay_ms.total_cmp(&right.delay_ms))
});
}
let mut hasher = blake3::Hasher::new();
serde_json::to_writer(
Blake3Writer(&mut hasher),
&CanonicalGraph { block_size, nodes },
)?;
Ok(hasher.finalize().to_hex().to_string())
}
fn trigger_rank(trigger: AgenticDependencyTrigger) -> u8 {
match trigger {
AgenticDependencyTrigger::Dispatch => 0,
AgenticDependencyTrigger::Completion => 1,
}
}
fn relation_rank(relation: AgenticDependencyRelation) -> u8 {
match relation {
AgenticDependencyRelation::Sequence => 0,
AgenticDependencyRelation::Spawn => 1,
AgenticDependencyRelation::Join => 2,
AgenticDependencyRelation::ReplayBarrier => 3,
}
}
fn extend_applied_compute_agentic_hash_ids(
hash_ids: &mut Vec<u64>,
input_length: usize,
trace_block_size: usize,
shared_initial_blocks: usize,
group_id: Option<usize>,
next_unique_hash: &mut u64,
) -> Result<()> {
let target_blocks = input_length.div_ceil(trace_block_size);
while hash_ids.len() < target_blocks {
let block_idx = hash_ids.len();
if block_idx < shared_initial_blocks
&& let Some(group_id) = group_id
{
hash_ids.push(0xA63E_0000_0000_0000 | ((group_id as u64) << 32) | block_idx as u64);
continue;
}
hash_ids.push(*next_unique_hash);
*next_unique_hash = next_unique_hash
.checked_add(1)
.ok_or_else(|| anyhow!("synthetic hash id overflow"))?;
}
Ok(())
}
fn sample_delay_ms(spec: &DelaySpec, rng: &mut StdRng) -> Result<f64> {
match spec {
DelaySpec::None => Ok(0.0),
DelaySpec::ConstantMs(delay_ms) => {
if !delay_ms.is_finite() || *delay_ms < 0.0 {
bail!("delay must be a finite non-negative number, got {delay_ms}");
}
Ok(*delay_ms)
}
DelaySpec::ExponentialMs { mean_ms } => {
if !mean_ms.is_finite() || *mean_ms < 0.0 {
bail!("mean_ms must be a finite non-negative number, got {mean_ms}");
}
Ok(sample_exponential_delay_ms(*mean_ms, rng))
}
}
}
fn sample_length(spec: &LengthSpec, min_value: usize, rng: &mut StdRng) -> usize {
if spec.stddev == 0.0 {
return spec.mean.max(min_value);
}
let stddev = spec.stddev.abs();
let u1 = (1.0 - rng.random::<f64>()).clamp(f64::MIN_POSITIVE, 1.0);
let u2 = rng.random::<f64>();
let z0 = (-2.0 * u1.ln()).sqrt() * (std::f64::consts::TAU * u2).cos();
let sample = spec.mean as f64 + z0 * stddev;
sample.round().max(min_value as f64) as usize
}
fn sample_exponential_delay_ms(mean_ms: f64, rng: &mut StdRng) -> f64 {
if mean_ms == 0.0 {
return 0.0;
}
let u = (1.0 - rng.random::<f64>()).clamp(f64::MIN_POSITIVE, 1.0);
-mean_ms * u.ln()
}