Merge pull request #74 from andjo403/collapse_threads

stack collaps events from multiple threads correct

Author: wesleywiser

flamegraph/src/main.rs

@@ -2,7 +2,6 @@ use std::error::Error;
 use std::fs::File;
 use std::io::BufWriter;
 use std::path::PathBuf;
-use std::time::Duration;
 
 use measureme::ProfilingData;
 
@@ -26,12 +25,7 @@ fn main() -> Result<(), Box<dyn Error>> {
 
     let profiling_data = ProfilingData::new(&opt.file_prefix)?;
 
-    let first_event_time = {
-        let current_time = profiling_data.iter().next().unwrap().timestamp;
-        current_time + Duration::from_millis(opt.interval)
-    };
-
-    let recorded_stacks = collapse_stacks(profiling_data.iter(), first_event_time, opt.interval)
+    let recorded_stacks = collapse_stacks(profiling_data.iter(), opt.interval)
         .iter()
         .map(|(unique_stack, count)| format!("{} {}", unique_stack, count))
         .collect::<Vec<_>>();

stack_collapse/src/main.rs

@@ -2,7 +2,6 @@ use std::error::Error;
 use std::fs::File;
 use std::io::{BufWriter, Write};
 use std::path::PathBuf;
-use std::time::Duration;
 
 use measureme::ProfilingData;
 
@@ -24,12 +23,7 @@ fn main() -> Result<(), Box<dyn Error>> {
 
     let profiling_data = ProfilingData::new(&opt.file_prefix)?;
 
-    let first_event_time = {
-        let current_time = profiling_data.iter().next().unwrap().timestamp;
-        current_time + Duration::from_millis(opt.interval)
-    };
-
-    let recorded_stacks = collapse_stacks(profiling_data.iter(), first_event_time, opt.interval);
+    let recorded_stacks = collapse_stacks(profiling_data.iter(), opt.interval);
 
     let mut file = BufWriter::new(File::create("out.stacks_folded")?);
 

tools_lib/src/stack_collapse.rs

@@ -5,38 +5,33 @@ use measureme::{Event, TimestampKind};
 
 pub fn collapse_stacks<'a>(
     events: impl Iterator<Item = Event<'a>>,
-    first_event_time: SystemTime,
     interval: u64,
 ) -> HashMap<String, usize> {
     let mut recorded_stacks = HashMap::<String, usize>::new();
-
-    let mut next_observation_time = first_event_time;
-
-    let mut thread_stacks: HashMap<u64, Vec<Event>> = HashMap::new();
+    let mut thread_stacks: HashMap<u64, (SystemTime, Vec<Event>)> = HashMap::new();
 
     for event in events {
+        let (next_observation_time, thread_stack) = thread_stacks
+            .entry(event.thread_id)
+            .or_insert((event.timestamp, Vec::new()));
         //if this event is after the next_observation_time then we need to record the current stacks
-        while event.timestamp > next_observation_time {
-            for (_tid, stack) in &thread_stacks {
-                let mut stack_string = String::new();
-                stack_string.push_str("rustc;");
+        if event.timestamp > *next_observation_time {
+            let mut stack_string = String::new();
+            stack_string.push_str("rustc");
 
-                for event in stack {
-                    stack_string.push_str(&event.label);
-                    stack_string.push(';');
-                }
-
-                //remove the trailing ';'
-                stack_string.remove(stack_string.len() - 1);
+            for event in thread_stack.iter() {
+                stack_string.push(';');
+                stack_string.push_str(&event.label);
+            }
 
-                *recorded_stacks.entry(stack_string).or_default() += 1;
+            let count = recorded_stacks.entry(stack_string).or_default();
 
-                next_observation_time += Duration::from_millis(interval);
+            while event.timestamp > *next_observation_time {
+                *count += 1;
+                *next_observation_time += Duration::from_millis(interval);
             }
         }
 
-        let thread_stack = thread_stacks.entry(event.thread_id).or_default();
-
         match event.timestamp_kind {
             TimestampKind::Start => {
                 thread_stack.push(event);
@@ -112,9 +107,7 @@ mod test {
             },
         ];
 
-        let first_event_time = events[0].timestamp;
-
-        let recorded_stacks = super::collapse_stacks(events.iter().cloned(), first_event_time, 1);
+        let recorded_stacks = super::collapse_stacks(events.iter().cloned(), 1);
 
         let mut expected_stacks = HashMap::<String, usize>::new();
         expected_stacks.insert("rustc;EventB;EventA".into(), 1000);
@@ -124,4 +117,67 @@ mod test {
 
         assert_eq!(expected_stacks, recorded_stacks);
     }
+
+    #[test]
+    fn multi_threaded_test() {
+        let events = [
+            Event {
+                event_kind: "Query".into(),
+                label: "EventA".into(),
+                additional_data: &[],
+                timestamp: SystemTime::UNIX_EPOCH + Duration::from_secs(1),
+                timestamp_kind: TimestampKind::Start,
+                thread_id: 1,
+            },
+            Event {
+                event_kind: "Query".into(),
+                label: "EventB".into(),
+                additional_data: &[],
+                timestamp: SystemTime::UNIX_EPOCH + Duration::from_secs(3),
+                timestamp_kind: TimestampKind::Start,
+                thread_id: 2,
+            },
+            Event {
+                event_kind: "Query".into(),
+                label: "EventA".into(),
+                additional_data: &[],
+                timestamp: SystemTime::UNIX_EPOCH + Duration::from_secs(2),
+                timestamp_kind: TimestampKind::End,
+                thread_id: 1,
+            },
+            Event {
+                event_kind: "Query".into(),
+                label: "EventA".into(),
+                additional_data: &[],
+                timestamp: SystemTime::UNIX_EPOCH + Duration::from_secs(4),
+                timestamp_kind: TimestampKind::Start,
+                thread_id: 2,
+            },
+            Event {
+                event_kind: "Query".into(),
+                label: "EventA".into(),
+                additional_data: &[],
+                timestamp: SystemTime::UNIX_EPOCH + Duration::from_secs(5),
+                timestamp_kind: TimestampKind::End,
+                thread_id: 2,
+            },
+            Event {
+                event_kind: "Query".into(),
+                label: "EventB".into(),
+                additional_data: &[],
+                timestamp: SystemTime::UNIX_EPOCH + Duration::from_secs(6),
+                timestamp_kind: TimestampKind::End,
+                thread_id: 2,
+            },
+        ];
+
+        let recorded_stacks = super::collapse_stacks(events.iter().cloned(), 1000);
+
+        let mut expected_stacks = HashMap::<String, usize>::new();
+        expected_stacks.insert("rustc;EventB;EventA".into(), 1);
+        expected_stacks.insert("rustc;EventB".into(), 2);
+        expected_stacks.insert("rustc;EventA".into(), 1);
+
+        assert_eq!(expected_stacks, recorded_stacks);
+    }
 }