use crate::engine::message::{AuditTrail, Change, Message};
use crate::engine::utils::strip_hash_prefix;
use chrono::{DateTime, Utc};
use datavalue::OwnedDataValue;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::sync::Arc;
const NODE_SIZE: usize = std::mem::size_of::<usize>();
#[inline]
fn is_false(b: &bool) -> bool {
!*b
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum StepResult {
Executed,
Skipped,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AuditTrailScope {
#[default]
Full,
Own,
None,
}
#[derive(Clone, Debug)]
pub struct TraceOptions {
pub snapshots: bool,
pub mapping_contexts: bool,
pub changes: bool,
pub max_snapshot_bytes: usize,
pub redact_paths: Vec<String>,
pub snapshot_audit_trail: AuditTrailScope,
}
impl Default for TraceOptions {
fn default() -> Self {
Self {
snapshots: true,
mapping_contexts: true,
changes: false,
max_snapshot_bytes: 0,
redact_paths: Vec::new(),
snapshot_audit_trail: AuditTrailScope::Full,
}
}
}
impl TraceOptions {
pub fn timings_only() -> Self {
Self {
snapshots: false,
mapping_contexts: false,
changes: true,
..Default::default()
}
}
fn redact_segments(&self) -> Vec<Vec<String>> {
self.redact_paths
.iter()
.filter(|p| !p.is_empty())
.map(|p| p.split('.').map(str::to_string).collect())
.collect()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[non_exhaustive]
pub struct ExecutionStep {
pub workflow_id: String,
pub task_id: Option<String>,
pub result: StepResult,
#[serde(skip_serializing_if = "Option::is_none")]
pub message: Option<Message>,
#[serde(skip_serializing_if = "Option::is_none")]
pub mapping_contexts: Option<Vec<Value>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub started_at: Option<DateTime<Utc>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub duration_us: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub changes: Option<Vec<Change>>,
}
impl ExecutionStep {
pub fn executed(workflow_id: &str, task_id: &str, message: &Message) -> Self {
Self {
workflow_id: workflow_id.to_string(),
task_id: Some(task_id.to_string()),
result: StepResult::Executed,
message: Some(message.clone()),
mapping_contexts: None,
started_at: None,
duration_us: None,
changes: None,
}
}
pub fn task_skipped(workflow_id: &str, task_id: &str) -> Self {
Self {
workflow_id: workflow_id.to_string(),
task_id: Some(task_id.to_string()),
result: StepResult::Skipped,
message: None,
mapping_contexts: None,
started_at: None,
duration_us: None,
changes: None,
}
}
pub fn workflow_skipped(workflow_id: &str) -> Self {
Self {
workflow_id: workflow_id.to_string(),
task_id: None,
result: StepResult::Skipped,
message: None,
mapping_contexts: None,
started_at: None,
duration_us: None,
changes: None,
}
}
pub fn with_mapping_contexts(mut self, contexts: Vec<Value>) -> Self {
self.mapping_contexts = Some(contexts);
self
}
pub fn with_timing(mut self, started_at: DateTime<Utc>, duration_us: u64) -> Self {
self.started_at = Some(started_at);
self.duration_us = Some(duration_us);
self
}
pub fn with_changes(mut self, changes: Vec<Change>) -> Self {
self.changes = Some(changes);
self
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[non_exhaustive]
pub struct ExecutionTrace {
pub steps: Vec<ExecutionStep>,
#[serde(default, skip_serializing_if = "is_false")]
truncated: bool,
#[serde(skip)]
options: TraceOptions,
#[serde(skip)]
redact_segments: Vec<Vec<String>>,
#[serde(skip)]
snapshot_bytes: usize,
}
impl ExecutionTrace {
pub fn new() -> Self {
Self::with_options(TraceOptions::default())
}
pub fn with_options(options: TraceOptions) -> Self {
Self {
steps: Vec::new(),
truncated: false,
redact_segments: options.redact_segments(),
options,
snapshot_bytes: 0,
}
}
pub fn options(&self) -> &TraceOptions {
&self.options
}
pub fn truncated(&self) -> bool {
self.truncated
}
pub fn add_step(&mut self, step: ExecutionStep) {
self.steps.push(step);
}
pub(crate) fn add_executed_step(
&mut self,
workflow_id: &str,
task_id: &str,
message: &Message,
started_at: DateTime<Utc>,
duration_us: u64,
mapping_contexts: Option<Vec<Value>>,
) {
let mut step = ExecutionStep {
workflow_id: workflow_id.to_string(),
task_id: Some(task_id.to_string()),
result: StepResult::Executed,
message: None,
mapping_contexts: None,
started_at: Some(started_at),
duration_us: Some(duration_us),
changes: if self.options.changes {
Some(
own_audit_entry(message, workflow_id, task_id)
.map(|e| e.changes.clone())
.unwrap_or_default(),
)
} else {
None
},
};
if self.options.snapshots {
let projected = self.projected_snapshot_size(message, workflow_id, task_id);
if self.would_exceed(projected) {
self.truncated = true;
} else {
self.snapshot_bytes += projected;
step.message = Some(self.build_snapshot(message, workflow_id, task_id));
}
}
if self.options.mapping_contexts {
if let Some(mut contexts) = mapping_contexts {
for ctx in &mut contexts {
redact_json_in_place(ctx, &self.redact_segments);
}
let size: usize = contexts.iter().map(approx_json_size).sum();
if self.would_exceed(size) {
self.truncated = true;
} else {
self.snapshot_bytes += size;
step.mapping_contexts = Some(contexts);
}
}
}
self.steps.push(step);
}
#[inline]
fn would_exceed(&self, additional: usize) -> bool {
self.options.max_snapshot_bytes != 0
&& self.snapshot_bytes + additional > self.options.max_snapshot_bytes
}
fn projected_snapshot_size(
&self,
message: &Message,
workflow_id: &str,
task_id: &str,
) -> usize {
let mut size = redacted_size(&message.context, &self.redact_segments);
for entry in self.scoped_audit_trail(message, workflow_id, task_id) {
size += NODE_SIZE;
for change in &entry.changes {
size += change.path.len()
+ approx_owned_size(&change.old_value)
+ approx_owned_size(&change.new_value);
}
}
size
}
fn scoped_audit_trail<'m>(
&self,
message: &'m Message,
workflow_id: &str,
task_id: &str,
) -> Vec<&'m AuditTrail> {
match self.options.snapshot_audit_trail {
AuditTrailScope::Full => message.audit_trail.iter().collect(),
AuditTrailScope::Own => own_audit_entry(message, workflow_id, task_id)
.map(|e| vec![e])
.unwrap_or_default(),
AuditTrailScope::None => Vec::new(),
}
}
fn build_snapshot(&self, message: &Message, workflow_id: &str, task_id: &str) -> Message {
let (context, _) = redacting_clone(&message.context, &self.redact_segments);
let audit_trail: Vec<AuditTrail> = self
.scoped_audit_trail(message, workflow_id, task_id)
.into_iter()
.cloned()
.collect();
Message {
id: message.id.clone(),
payload: Arc::clone(&message.payload),
context,
audit_trail,
errors: message.errors.clone(),
capture_changes: message.capture_changes,
routing_bucket: message.routing_bucket,
}
}
pub fn final_message(&self) -> Option<&Message> {
self.steps
.iter()
.rev()
.find(|s| s.result == StepResult::Executed)
.and_then(|s| s.message.as_ref())
}
pub fn is_success(&self) -> bool {
self.final_message()
.map(|m| m.errors.is_empty())
.unwrap_or(true)
}
pub fn executed_count(&self) -> usize {
self.steps
.iter()
.filter(|s| s.result == StepResult::Executed)
.count()
}
pub fn skipped_count(&self) -> usize {
self.steps
.iter()
.filter(|s| s.result == StepResult::Skipped)
.count()
}
}
impl Default for ExecutionTrace {
fn default() -> Self {
Self::new()
}
}
#[inline]
fn own_audit_entry<'m>(
message: &'m Message,
workflow_id: &str,
task_id: &str,
) -> Option<&'m AuditTrail> {
match message.audit_trail.last() {
Some(entry)
if entry.task_id.as_ref() == task_id && entry.workflow_id.as_ref() == workflow_id =>
{
Some(entry)
}
_ => None,
}
}
#[inline]
pub(crate) fn duration_us_between(start: DateTime<Utc>, end: DateTime<Utc>) -> u64 {
(end - start)
.num_microseconds()
.unwrap_or(0)
.max(0)
.try_into()
.unwrap_or(0)
}
fn redacting_clone(value: &OwnedDataValue, paths: &[Vec<String>]) -> (OwnedDataValue, usize) {
let refs: Vec<&[String]> = paths.iter().map(|p| p.as_slice()).collect();
redacting_clone_inner(value, &refs)
}
fn narrow_for_object_key<'a>(paths: &[&'a [String]], key: &str) -> Vec<&'a [String]> {
paths
.iter()
.filter(|p| strip_hash_prefix(&p[0]) == key)
.map(|p| &p[1..])
.collect()
}
fn narrow_for_array_index<'a>(paths: &[&'a [String]], idx: usize) -> Vec<&'a [String]> {
paths
.iter()
.filter(|p| p[0].parse::<usize>() == Ok(idx))
.map(|p| &p[1..])
.collect()
}
fn redacting_clone_inner(value: &OwnedDataValue, paths: &[&[String]]) -> (OwnedDataValue, usize) {
if paths.iter().any(|p| p.is_empty()) {
return (OwnedDataValue::Null, NODE_SIZE);
}
match value {
OwnedDataValue::Object(pairs) => {
let mut out = Vec::with_capacity(pairs.len());
let mut size = NODE_SIZE;
for (key, child) in pairs {
let sub = narrow_for_object_key(paths, key);
let (cloned, child_size) = redacting_clone_inner(child, &sub);
size += key.len() + child_size;
out.push((key.clone(), cloned));
}
(OwnedDataValue::Object(out), size)
}
OwnedDataValue::Array(items) => {
let mut out = Vec::with_capacity(items.len());
let mut size = NODE_SIZE;
for (idx, child) in items.iter().enumerate() {
let sub = narrow_for_array_index(paths, idx);
let (cloned, child_size) = redacting_clone_inner(child, &sub);
size += child_size;
out.push(cloned);
}
(OwnedDataValue::Array(out), size)
}
OwnedDataValue::String(s) => (value.clone(), NODE_SIZE + s.len()),
other => (other.clone(), NODE_SIZE),
}
}
fn redacted_size(value: &OwnedDataValue, paths: &[Vec<String>]) -> usize {
let refs: Vec<&[String]> = paths.iter().map(|p| p.as_slice()).collect();
redacted_size_inner(value, &refs)
}
fn redacted_size_inner(value: &OwnedDataValue, paths: &[&[String]]) -> usize {
if paths.iter().any(|p| p.is_empty()) {
return NODE_SIZE;
}
match value {
OwnedDataValue::Object(pairs) => {
let mut size = NODE_SIZE;
for (key, child) in pairs {
let sub = narrow_for_object_key(paths, key);
size += key.len() + redacted_size_inner(child, &sub);
}
size
}
OwnedDataValue::Array(items) => {
let mut size = NODE_SIZE;
for (idx, child) in items.iter().enumerate() {
let sub = narrow_for_array_index(paths, idx);
size += redacted_size_inner(child, &sub);
}
size
}
OwnedDataValue::String(s) => NODE_SIZE + s.len(),
_ => NODE_SIZE,
}
}
fn approx_owned_size(value: &OwnedDataValue) -> usize {
match value {
OwnedDataValue::Object(pairs) => {
NODE_SIZE
+ pairs
.iter()
.map(|(k, v)| k.len() + approx_owned_size(v))
.sum::<usize>()
}
OwnedDataValue::Array(items) => {
NODE_SIZE + items.iter().map(approx_owned_size).sum::<usize>()
}
OwnedDataValue::String(s) => NODE_SIZE + s.len(),
_ => NODE_SIZE,
}
}
fn redact_json_in_place(value: &mut Value, paths: &[Vec<String>]) {
let refs: Vec<&[String]> = paths.iter().map(|p| p.as_slice()).collect();
redact_json_inner(value, &refs);
}
fn redact_json_inner(value: &mut Value, paths: &[&[String]]) {
if paths.is_empty() {
return;
}
if paths.iter().any(|p| p.is_empty()) {
*value = Value::Null;
return;
}
match value {
Value::Object(map) => {
for (key, child) in map.iter_mut() {
let sub = narrow_for_object_key(paths, key);
redact_json_inner(child, &sub);
}
}
Value::Array(items) => {
for (idx, child) in items.iter_mut().enumerate() {
let sub = narrow_for_array_index(paths, idx);
redact_json_inner(child, &sub);
}
}
_ => {}
}
}
fn approx_json_size(value: &Value) -> usize {
match value {
Value::Object(map) => {
NODE_SIZE
+ map
.iter()
.map(|(k, v)| k.len() + approx_json_size(v))
.sum::<usize>()
}
Value::Array(items) => NODE_SIZE + items.iter().map(approx_json_size).sum::<usize>(),
Value::String(s) => NODE_SIZE + s.len(),
_ => NODE_SIZE,
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn dv(v: serde_json::Value) -> OwnedDataValue {
OwnedDataValue::from(&v)
}
fn segments(paths: &[&str]) -> Vec<Vec<String>> {
TraceOptions {
redact_paths: paths.iter().map(|s| s.to_string()).collect(),
..Default::default()
}
.redact_segments()
}
#[test]
fn test_step_result_serialization() {
assert_eq!(
serde_json::to_string(&StepResult::Executed).unwrap(),
"\"executed\""
);
assert_eq!(
serde_json::to_string(&StepResult::Skipped).unwrap(),
"\"skipped\""
);
}
#[test]
fn test_execution_step_executed() {
let message = Message::from_value(&json!({"test": "data"}));
let step = ExecutionStep::executed("workflow1", "task1", &message);
assert_eq!(step.workflow_id, "workflow1");
assert_eq!(step.task_id, Some("task1".to_string()));
assert_eq!(step.result, StepResult::Executed);
assert!(step.message.is_some());
}
#[test]
fn test_execution_step_task_skipped() {
let step = ExecutionStep::task_skipped("workflow1", "task1");
assert_eq!(step.workflow_id, "workflow1");
assert_eq!(step.task_id, Some("task1".to_string()));
assert_eq!(step.result, StepResult::Skipped);
assert!(step.message.is_none());
}
#[test]
fn test_execution_step_workflow_skipped() {
let step = ExecutionStep::workflow_skipped("workflow1");
assert_eq!(step.workflow_id, "workflow1");
assert_eq!(step.task_id, None);
assert_eq!(step.result, StepResult::Skipped);
assert!(step.message.is_none());
}
#[test]
fn test_execution_step_with_mapping_contexts() {
let message = Message::from_value(&json!({"test": "data"}));
let contexts = vec![json!({"data": {"a": 1}}), json!({"data": {"a": 1, "b": 2}})];
let step = ExecutionStep::executed("workflow1", "task1", &message)
.with_mapping_contexts(contexts.clone());
assert_eq!(step.mapping_contexts, Some(contexts));
let serialized = serde_json::to_value(&step).unwrap();
assert!(serialized.get("mapping_contexts").is_some());
assert_eq!(serialized["mapping_contexts"].as_array().unwrap().len(), 2);
}
#[test]
fn test_execution_step_without_mapping_contexts_serialization() {
let message = Message::from_value(&json!({"test": "data"}));
let step = ExecutionStep::executed("workflow1", "task1", &message);
let serialized = serde_json::to_value(&step).unwrap();
assert!(serialized.get("mapping_contexts").is_none());
assert!(serialized.get("started_at").is_none());
assert!(serialized.get("duration_us").is_none());
assert!(serialized.get("changes").is_none());
}
#[test]
fn test_execution_trace() {
let mut trace = ExecutionTrace::new();
let message = Message::from_value(&json!({"test": "data"}));
trace.add_step(ExecutionStep::workflow_skipped("workflow0"));
trace.add_step(ExecutionStep::executed("workflow1", "task1", &message));
trace.add_step(ExecutionStep::task_skipped("workflow1", "task2"));
assert_eq!(trace.steps.len(), 3);
assert_eq!(trace.executed_count(), 1);
assert_eq!(trace.skipped_count(), 2);
assert!(trace.final_message().is_some());
assert!(trace.is_success());
}
#[test]
fn default_options_reproduce_historical_capture() {
let o = TraceOptions::default();
assert!(o.snapshots);
assert!(o.mapping_contexts);
assert!(!o.changes);
assert_eq!(o.max_snapshot_bytes, 0);
assert!(o.redact_paths.is_empty());
assert_eq!(o.snapshot_audit_trail, AuditTrailScope::Full);
}
#[test]
fn timings_only_drops_snapshots_and_keeps_the_diff() {
let o = TraceOptions::timings_only();
assert!(!o.snapshots);
assert!(!o.mapping_contexts);
assert!(o.changes);
}
#[test]
fn a_complete_trace_does_not_serialize_the_truncated_flag() {
let trace = ExecutionTrace::new();
let serialized = serde_json::to_value(&trace).unwrap();
assert!(
serialized.get("truncated").is_none(),
"a complete trace keeps the historical wire shape"
);
assert!(!trace.truncated());
}
#[test]
fn a_trace_deserializes_from_a_payload_without_the_truncated_flag() {
let trace: ExecutionTrace = serde_json::from_value(json!({ "steps": [] })).unwrap();
assert!(!trace.truncated());
}
#[test]
fn duration_clamps_a_backward_clock_step_to_zero() {
let start = Utc::now();
let earlier = start - chrono::Duration::seconds(5);
assert_eq!(duration_us_between(start, earlier), 0);
assert_eq!(duration_us_between(start, start), 0);
assert_eq!(
duration_us_between(start, start + chrono::Duration::microseconds(1500)),
1500
);
}
#[test]
fn redaction_nulls_only_the_named_subtree() {
let ctx = dv(json!({"data": {"secret": {"k": "v"}, "keep": 1}}));
let (out, _) = redacting_clone(&ctx, &segments(&["data.secret"]));
assert_eq!(
serde_json::Value::from(&out),
json!({"data": {"secret": null, "keep": 1}})
);
}
#[test]
fn redaction_of_an_unresolvable_path_creates_nothing() {
let ctx = dv(json!({"data": {"items": [1, 2, 3]}}));
let (out, _) = redacting_clone(&ctx, &segments(&["data.items.99"]));
assert_eq!(
serde_json::Value::from(&out),
json!({"data": {"items": [1, 2, 3]}})
);
}
#[test]
fn redaction_through_a_non_container_is_a_noop() {
let ctx = dv(json!({"data": {"name": "alice"}}));
let (out, _) = redacting_clone(&ctx, &segments(&["data.name.first"]));
assert_eq!(
serde_json::Value::from(&out),
json!({"data": {"name": "alice"}})
);
}
#[test]
fn an_empty_redact_path_is_ignored() {
let ctx = dv(json!({"data": {"a": 1}}));
let (out, _) = redacting_clone(&ctx, &segments(&[""]));
assert_eq!(serde_json::Value::from(&out), json!({"data": {"a": 1}}));
}
#[test]
fn redaction_honours_the_hash_escape() {
let obj = dv(json!({"data": {"20": "secret", "other": 1}}));
let (out, _) = redacting_clone(&obj, &segments(&["data.#20"]));
assert_eq!(
serde_json::Value::from(&out),
json!({"data": {"20": null, "other": 1}})
);
let arr = dv(json!({"data": [0, 1, 2]}));
let (out, _) = redacting_clone(&arr, &segments(&["data.1"]));
assert_eq!(serde_json::Value::from(&out), json!({"data": [0, null, 2]}));
let (out, _) = redacting_clone(&arr, &segments(&["data.#1"]));
assert_eq!(serde_json::Value::from(&out), json!({"data": [0, 1, 2]}));
}
#[test]
fn nested_and_duplicated_redact_paths_are_safe() {
let ctx = dv(json!({"data": {"a": {"b": 1, "c": 2}}}));
let (out, _) = redacting_clone(&ctx, &segments(&["data.a", "data.a.b"]));
assert_eq!(serde_json::Value::from(&out), json!({"data": {"a": null}}));
let (out, _) = redacting_clone(&ctx, &segments(&["data.a", "data.a"]));
assert_eq!(serde_json::Value::from(&out), json!({"data": {"a": null}}));
}
#[test]
fn redaction_matches_unicode_keys_and_sizes_strings_by_bytes() {
let ctx = dv(json!({"data": {"café": "secret", "keep": "née"}}));
let (out, size) = redacting_clone(&ctx, &segments(&["data.café"]));
assert_eq!(
serde_json::Value::from(&out),
json!({"data": {"café": null, "keep": "née"}})
);
let (_, unredacted) = redacting_clone(&ctx, &segments(&[]));
assert!(unredacted > size, "redacting must lower the counted size");
assert!(
approx_owned_size(&dv(json!("née"))) == NODE_SIZE + 4,
"str::len() bytes, not chars().count()"
);
}
#[test]
fn redacted_size_agrees_with_redacting_clone() {
let shapes = [
json!({}),
json!({"data": {"a": 1, "b": "hello"}}),
json!({"data": {"items": [1, "two", {"three": 3}], "nested": {"x": {"y": "z"}}}}),
json!({"data": {"secret": {"deep": [1, 2, 3]}, "keep": "café"}}),
];
let path_sets: [&[&str]; 4] = [&[], &["data.secret"], &["data.items.1"], &["data.nope"]];
for shape in &shapes {
for paths in path_sets {
let v = dv(shape.clone());
let segs = segments(paths);
let (_, cloned_size) = redacting_clone(&v, &segs);
assert_eq!(
redacted_size(&v, &segs),
cloned_size,
"probe and clone disagree for {shape:?} with {paths:?}"
);
}
}
}
#[test]
fn redaction_applies_to_mapping_contexts_too() {
let mut ctx = json!({"data": {"secret": {"k": "v"}, "keep": 1}});
redact_json_in_place(&mut ctx, &segments(&["data.secret"]));
assert_eq!(ctx, json!({"data": {"secret": null, "keep": 1}}));
}
#[test]
fn json_redaction_shares_the_owned_path_semantics() {
let mut arr = json!({"data": {"items": [1, 2, 3]}});
redact_json_in_place(&mut arr, &segments(&["data.items.99"]));
assert_eq!(arr, json!({"data": {"items": [1, 2, 3]}}));
let mut scalar = json!({"data": {"name": "alice"}});
redact_json_in_place(&mut scalar, &segments(&["data.name.first"]));
assert_eq!(scalar, json!({"data": {"name": "alice"}}));
let mut hash = json!({"data": {"20": "secret"}});
redact_json_in_place(&mut hash, &segments(&["data.#20"]));
assert_eq!(hash, json!({"data": {"20": null}}));
}
}