use crate::causal_graph::CausalGraph;
use anyhow::{Context, Result};
use std::collections::HashMap;
use trueno_graph::NodeId;
#[derive(Debug, Clone, PartialEq)]
pub struct CriticalPathResult {
pub path: Vec<NodeId>,
pub total_duration: u64,
pub node_durations: HashMap<NodeId, u64>,
pub span_names: Vec<String>,
}
impl CriticalPathResult {
pub fn critical_path_percentage(&self, total_trace_duration: u64) -> f64 {
if total_trace_duration == 0 {
return 0.0;
}
(self.total_duration as f64 / total_trace_duration as f64) * 100.0
}
pub fn longest_span(&self) -> Option<(NodeId, u64)> {
self.node_durations
.iter()
.max_by_key(|(_, &duration)| duration)
.map(|(&node, &duration)| (node, duration))
}
pub fn is_on_critical_path(&self, node: NodeId) -> bool {
self.path.contains(&node)
}
}
pub fn find_critical_path(graph: &CausalGraph) -> Result<CriticalPathResult> {
if graph.node_count() == 0 {
return Ok(CriticalPathResult {
path: Vec::new(),
total_duration: 0,
node_durations: HashMap::new(),
span_names: Vec::new(),
});
}
let topo_order = topological_sort(graph)?;
let mut dist: HashMap<NodeId, u64> = HashMap::new();
let mut parent: HashMap<NodeId, Option<NodeId>> = HashMap::new();
for &root in graph.roots() {
let span = graph.get_span(root).context("Root span not found in metadata")?;
dist.insert(root, span.duration_nanos);
parent.insert(root, None);
}
for &node in &topo_order {
let children = graph.children(node)?;
for (child, _edge_weight) in children {
let child_span = graph.get_span(child).context("Child span not found in metadata")?;
let new_dist = dist.get(&node).unwrap_or(&0) + child_span.duration_nanos;
if new_dist > *dist.get(&child).unwrap_or(&0) {
dist.insert(child, new_dist);
parent.insert(child, Some(node));
}
}
}
let (&critical_end, &total_duration) =
dist.iter().max_by_key(|(_, &d)| d).context("No paths found in graph")?;
let (path, node_durations, span_names) = reconstruct_path(graph, critical_end, &parent);
Ok(CriticalPathResult { path, total_duration, node_durations, span_names })
}
fn reconstruct_path(
graph: &CausalGraph,
start: NodeId,
parent: &HashMap<NodeId, Option<NodeId>>,
) -> (Vec<NodeId>, HashMap<NodeId, u64>, Vec<String>) {
let mut path = Vec::new();
let mut node_durations = HashMap::new();
let mut span_names = Vec::new();
let mut current = start;
loop {
path.push(current);
if let Some(span) = graph.get_span(current) {
node_durations.insert(current, span.duration_nanos);
span_names.push(span.span_name.clone());
}
match parent.get(¤t) {
Some(Some(p)) => current = *p,
_ => break,
}
}
path.reverse();
span_names.reverse();
(path, node_durations, span_names)
}
fn topological_sort(graph: &CausalGraph) -> Result<Vec<NodeId>> {
if !graph.is_dag()? {
anyhow::bail!("Graph contains cycles - cannot compute critical path");
}
let mut visited = std::collections::HashSet::new();
let mut stack = Vec::new();
for &root in graph.roots() {
if !visited.contains(&root) {
dfs_topo(graph, root, &mut visited, &mut stack)?;
}
}
stack.reverse();
Ok(stack)
}
fn dfs_topo(
graph: &CausalGraph,
node: NodeId,
visited: &mut std::collections::HashSet<NodeId>,
stack: &mut Vec<NodeId>,
) -> Result<()> {
visited.insert(node);
let children = graph.children(node)?;
for (child, _) in children {
if !visited.contains(&child) {
dfs_topo(graph, child, visited, stack)?;
}
}
stack.push(node);
Ok(())
}
pub fn find_all_critical_paths(
graph: &CausalGraph,
_tolerance_ns: u64,
) -> Result<Vec<CriticalPathResult>> {
Ok(vec![find_critical_path(graph)?])
}
static_assertions::assert_impl_all!(CriticalPathResult: Send, Sync);
#[cfg(test)]
mod tests {
use super::*;
use crate::span_record::{SpanKind, SpanRecord, StatusCode};
use std::collections::HashMap;
fn create_span(
span_id: u8,
parent_id: Option<u8>,
logical_clock: u64,
duration_nanos: u64,
) -> SpanRecord {
SpanRecord::new(
[1; 16],
[span_id; 8],
parent_id.map(|p| [p; 8]),
format!("span_{}", span_id),
SpanKind::Internal,
logical_clock * 1000,
logical_clock * 1000 + duration_nanos,
logical_clock,
StatusCode::Ok,
String::new(),
HashMap::new(),
HashMap::new(),
1234,
5678,
)
}
#[test]
fn test_empty_graph() {
let graph = CausalGraph::from_spans(&[]).expect("test");
let result = find_critical_path(&graph).expect("test");
assert_eq!(result.path.len(), 0);
assert_eq!(result.total_duration, 0);
}
#[test]
fn test_single_span() {
let root = create_span(1, None, 0, 1000);
let graph = CausalGraph::from_spans(&[root]).expect("test");
let result = find_critical_path(&graph).expect("test");
assert_eq!(result.path.len(), 1);
assert_eq!(result.total_duration, 1000);
assert_eq!(result.span_names, vec!["span_1"]);
}
#[test]
fn test_linear_path() {
let root = create_span(1, None, 0, 1000);
let child = create_span(2, Some(1), 1, 500);
let grandchild = create_span(3, Some(2), 2, 700);
let graph = CausalGraph::from_spans(&[root, child, grandchild]).expect("test");
let result = find_critical_path(&graph).expect("test");
assert_eq!(result.path.len(), 3);
assert_eq!(result.total_duration, 1000 + 500 + 700);
assert_eq!(result.span_names, vec!["span_1", "span_2", "span_3"]);
}
#[test]
fn test_branching_path() {
let root = create_span(1, None, 0, 1000);
let child_short = create_span(2, Some(1), 1, 500);
let child_long = create_span(3, Some(1), 2, 2000);
let graph = CausalGraph::from_spans(&[root, child_short, child_long]).expect("test");
let result = find_critical_path(&graph).expect("test");
assert_eq!(result.path.len(), 2);
assert_eq!(result.total_duration, 1000 + 2000);
assert!(result.span_names.contains(&"span_3".to_string()));
}
#[test]
fn test_complex_tree() {
let root = create_span(1, None, 0, 1000);
let child1 = create_span(2, Some(1), 1, 500);
let child2 = create_span(3, Some(1), 2, 800);
let grandchild1 = create_span(4, Some(2), 3, 300);
let grandchild2 = create_span(5, Some(3), 4, 1200);
let graph = CausalGraph::from_spans(&[root, child1, child2, grandchild1, grandchild2])
.expect("test");
let result = find_critical_path(&graph).expect("test");
assert_eq!(result.path.len(), 3);
assert_eq!(result.total_duration, 1000 + 800 + 1200);
assert_eq!(result.span_names, vec!["span_1", "span_3", "span_5"]);
}
#[test]
fn test_longest_span() {
let root = create_span(1, None, 0, 500);
let child = create_span(2, Some(1), 1, 2000); let grandchild = create_span(3, Some(2), 2, 300);
let graph = CausalGraph::from_spans(&[root, child, grandchild]).expect("test");
let result = find_critical_path(&graph).expect("test");
let (longest_node, duration) = result.longest_span().expect("test");
assert_eq!(duration, 2000);
assert_eq!(result.node_durations.get(&longest_node), Some(&2000));
}
#[test]
fn test_is_on_critical_path() {
let root = create_span(1, None, 0, 1000);
let child1 = create_span(2, Some(1), 1, 500);
let child2 = create_span(3, Some(1), 2, 2000);
let graph = CausalGraph::from_spans(&[root, child1, child2]).expect("test");
let result = find_critical_path(&graph).expect("test");
assert!(result.is_on_critical_path(NodeId(0))); assert!(result.is_on_critical_path(NodeId(2))); assert!(!result.is_on_critical_path(NodeId(1))); }
#[test]
fn test_critical_path_percentage() {
let root = create_span(1, None, 0, 1000);
let child = create_span(2, Some(1), 1, 2000);
let graph = CausalGraph::from_spans(&[root, child]).expect("test");
let result = find_critical_path(&graph).expect("test");
let percentage = result.critical_path_percentage(5000);
assert_eq!(percentage, 60.0);
}
}