Merge branch 'cobrien7-event-callbacks'

2 files changed, 74 insertions, 0 deletions

diff --git a/tests/test_ftrace.py b/tests/test_ftrace.py
index 7d7874a..e6f6319 100644
--- a/tests/test_ftrace.py
+++ b/tests/test_ftrace.py
@@ -230,6 +230,31 @@ class TestFTrace(BaseTestThermal):
         # Make sure there are no NaNs in the middle of the array
         self.assertTrue(allfreqs[0][1]["A57_freq_in"].notnull().all())
 
+    def test_apply_callbacks(self):
+        """Test apply_callbacks()"""
+
+        counts = {
+            "cpu_in_power": 0,
+            "cpu_out_power": 0
+        }
+
+        def cpu_in_power_fn(data):
+            counts["cpu_in_power"] += 1
+
+        def cpu_out_power_fn(data):
+            counts["cpu_out_power"] += 1
+
+        fn_map = {
+            "cpu_in_power": cpu_in_power_fn,
+            "cpu_out_power": cpu_out_power_fn
+        }
+        trace = trappy.FTrace()
+
+        trace.apply_callbacks(fn_map)
+
+        self.assertEqual(counts["cpu_in_power"], 134)
+        self.assertEqual(counts["cpu_out_power"], 134)
+
     def test_plot_freq_hists(self):
         """Test that plot_freq_hists() doesn't bomb"""
 
diff --git a/trappy/ftrace.py b/trappy/ftrace.py
index dae1e43..06e2a93 100644
--- a/trappy/ftrace.py
+++ b/trappy/ftrace.py
@@ -401,6 +401,55 @@ is part of the trace.
 
         return ret
 
+    def apply_callbacks(self, fn_map):
+        """
+        Apply callback functions to trace events in chronological order.
+
+        This method iterates over a user-specified subset of the available trace
+        event dataframes, calling different user-specified functions for each
+        event type. These functions are passed a dictionary mapping 'Index' and
+        the column names to their values for that row.
+
+        For example, to iterate over trace t, applying your functions callback_fn1
+        and callback_fn2 to each sched_switch and sched_wakeup event respectively:
+
+        t.apply_callbacks({
+            "sched_switch": callback_fn1,
+            "sched_wakeup": callback_fn2
+        })
+        """
+        dfs = {event: getattr(self, event).data_frame for event in fn_map.keys()}
+        events = [event for event in fn_map.keys() if not dfs[event].empty]
+        iters = {event: dfs[event].itertuples() for event in events}
+        next_rows = {event: iterator.next() for event,iterator in iters.iteritems()}
+
+        # Column names beginning with underscore will not be preserved in tuples
+        # due to constraints on namedtuple field names, so store mappings from
+        # column name to column number for each trace event.
+        col_idxs = {event: {
+            name: idx for idx, name in enumerate(
+                ['Index'] + dfs[event].columns.tolist()
+            )
+        } for event in events}
+
+        def getLine(event):
+            line_col_idx = col_idxs[event]['__line']
+            return next_rows[event][line_col_idx]
+
+        while events:
+            event_name = min(events, key=getLine)
+            event_tuple = next_rows[event_name]
+
+            event_dict = {
+                col: event_tuple[idx] for col, idx in col_idxs[event_name].iteritems()
+            }
+            fn_map[event_name](event_dict)
+            event_row = next(iters[event_name], None)
+            if event_row:
+                next_rows[event_name] = event_row
+            else:
+                events.remove(event_name)
+
     def plot_freq_hists(self, map_label, ax):
         """Plot histograms for each actor input and output frequency