pw_build: Add --write-trace option to pw-wrap-ninja

This commit adds a new option to pw-wrap-ninja to write a trace file for
the build, in the Chromium trace_event format. This trace can be opened
with the Perfetto UI.

This trace only includes timing information from the build, not
dependency information from Ninja or additional compiler output. Thus,
it is simpler than a more comprehensive tool such as Fuchsia's
"ninjatrace". However, it's simple to use and still provides value.

Tested:
  Built Pigweed with pw-wrap-ninja and the --write-trace option.
  Confirmed the trace was viewable with the Perfetto UI.
Change-Id: I78e5717fb2257f74080def90fee253137defe3fd
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/126464
Commit-Queue: Eli Lipsitz <elipsitz@google.com>
Reviewed-by: Rob Mohr <mohrr@google.com>
Pigweed-Auto-Submit: Eli Lipsitz <elipsitz@google.com>
diff --git a/pw_build/docs.rst b/pw_build/docs.rst
index 368ea53..48325b2 100644
--- a/pw_build/docs.rst
+++ b/pw_build/docs.rst
@@ -828,8 +828,8 @@
 
 The templates for build time errors are defined in ``pw_build/error.gni``.
 
-Improved Ninja progress output
-------------------------------
+Improved Ninja interface
+------------------------
 Ninja includes a basic progress display, showing in a single line the number of
 targets finished, the total number of targets, and the name of the most recent
 target it has either started or finished.
@@ -864,6 +864,11 @@
 targets are bottlenecking the rest of the build, and which targets are taking
 an unusually long time to complete.
 
+``pw-wrap-ninja`` includes other useful functionality as well. The
+``--write-trace`` option writes a build trace to the specified path, which can
+be viewed in the `Perfetto UI <https://ui.perfetto.dev/>`_, or via Chrome's
+built-in ``chrome://tracing`` tool.
+
 CMake
 =====
 Pigweed's `CMake`_ support is provided primarily for projects that have an
diff --git a/pw_build/py/pw_build/wrap_ninja.py b/pw_build/py/pw_build/wrap_ninja.py
index 2a710eb..ec39947 100644
--- a/pw_build/py/pw_build/wrap_ninja.py
+++ b/pw_build/py/pw_build/wrap_ninja.py
@@ -17,6 +17,7 @@
 import argparse
 import dataclasses
 import enum
+import json
 import os
 import pty
 import re
@@ -182,23 +183,27 @@
         num_started: The number of started actions.
         num_finished: The number of finished actions.
         num_total: The (estimated) number of total actions in the build.
-        actions: The currently-running actions, keyed by the action name.
+        running_actions: The currently-running actions, keyed by name.
+        actions: A list of all actions, running and completed.
         last_action_completed: The last action that Ninja finished building.
         events: The events that have occured in the Ninja build.
         exited: Whether the Ninja subprocess has exited.
         lock: The lock guarding the rest of the attributes.
         process: The Python subprocess running Ninja.
+        start_time: The time that the Ninja build started.
     """
 
     num_started: int
     num_finished: int
     num_total: int
-    actions: Dict[str, NinjaAction]
+    actions: List[NinjaAction]
+    running_actions: Dict[str, NinjaAction]
     last_action_completed: Optional[NinjaAction]
     events: List[NinjaEvent]
     exited: bool
     lock: threading.Lock
     process: subprocess.Popen
+    start_time: float
 
     def __init__(self, args: List[str]) -> None:
         if sys.platform == 'win32':
@@ -207,7 +212,8 @@
         self.num_started = 0
         self.num_finished = 0
         self.num_total = 0
-        self.actions = {}
+        self.actions = []
+        self.running_actions = {}
         self.last_action_completed = None
         self.events = []
         self.exited = False
@@ -227,6 +233,7 @@
             stderr=ptty_child,
         )
         os.close(ptty_child)
+        self.start_time = time.time()
 
         # Start the output monitor thread.
         thread = threading.Thread(
@@ -288,22 +295,30 @@
                 # Special case: first action in a new graph.
                 # This is important for GN's "Regenerating ninja files" action.
                 if actions_started == 1 and actions_finished == 0:
-                    self.actions = {}
+                    for action in self.running_actions.values():
+                        action.end_time = time.time()
+                        self.events.append(
+                            NinjaEvent(NinjaEventKind.ACTION_FINISHED, action)
+                        )
+                    self.running_actions = {}
                     self.last_action_completed = None
 
                 if did_start:
-                    action = self.actions.setdefault(name, NinjaAction(name))
+                    action = self.running_actions.setdefault(
+                        name, NinjaAction(name)
+                    )
                     if action.jobs == 0:
+                        self.actions.append(action)
                         self.events.append(
                             NinjaEvent(NinjaEventKind.ACTION_STARTED, action)
                         )
                     action.jobs += 1
 
-                if did_finish and name in self.actions:
-                    action = self.actions[name]
+                if did_finish and name in self.running_actions:
+                    action = self.running_actions[name]
                     action.jobs -= 1
                     if action.jobs <= 0:
-                        self.actions.pop(name)
+                        self.running_actions.pop(name)
                         self.last_action_completed = action
                         action.end_time = time.time()
                         self.events.append(
@@ -321,6 +336,43 @@
                     )
                 )
 
+    def write_trace(self, file: IO[str]) -> None:
+        """Write a Chromium trace_event-formatted trace to a file."""
+        now = time.time()
+        threads: List[float] = []
+        events = []
+        actions = sorted(
+            self.actions, key=lambda x: x.end_time or now, reverse=True
+        )
+        for action in actions:
+            # If this action hasn't completed, fake the end time.
+            end_time = action.end_time or now
+
+            # Allocate a "thread id" for the action. We look at threads in
+            # reverse end time and fill in threads from the end to get a
+            # better-looking trace.
+            for tid, busy_time in enumerate(threads):
+                if busy_time >= end_time:
+                    threads[tid] = action.start_time
+                    break
+            else:
+                tid = len(threads)
+                threads.append(action.start_time)
+
+            events.append(
+                {
+                    'name': action.name,
+                    'cat': 'actions',
+                    'ph': 'X',
+                    'ts': action.start_time * 1000000,
+                    'dur': (end_time - action.start_time) * 1000000,
+                    'pid': 0,
+                    'tid': tid,
+                    'args': {},
+                }
+            )
+        json.dump(events, file)
+
 
 class UI:
     """Class to handle UI state and rendering."""
@@ -328,17 +380,16 @@
     def __init__(self, ninja: Ninja, args: argparse.Namespace) -> None:
         self._ninja = ninja
         self._args = args
-        self._start_time = time.time()
         self._renderer = ConsoleRenderer()
         self._last_log_action: Optional[NinjaAction] = None
 
     def _get_status_display(self) -> List[str]:
         """Generates the status display. Must be called under the Ninja lock."""
         actions = sorted(
-            self._ninja.actions.values(), key=lambda x: x.start_time
+            self._ninja.running_actions.values(), key=lambda x: x.start_time
         )
         now = time.time()
-        total_elapsed = _format_duration(now - self._start_time)
+        total_elapsed = _format_duration(now - self._ninja.start_time)
         lines = [
             f'[{total_elapsed: >5}] '
             f'Building [{self._ninja.num_finished}/{self._ninja.num_total}] ...'
@@ -393,7 +444,7 @@
 
     def print_summary(self) -> None:
         """Prints the summary line at the end of the build."""
-        total_time = _format_duration(time.time() - self._start_time)
+        total_time = _format_duration(time.time() - self._ninja.start_time)
         num_finished = self._ninja.num_finished
         num_total = self._ninja.num_total
         self._renderer.print_line(
@@ -428,6 +479,11 @@
         help='whether to log when actions start and finish',
         action='store_true',
     )
+    parser.add_argument(
+        '--write-trace',
+        help='write a Chromium trace_event-style trace to the given file',
+        type=argparse.FileType('w'),
+    )
     return parser.parse_known_args()
 
 
@@ -451,10 +507,15 @@
         interface.update()
         time.sleep(args.ui_update_rate)
 
-    # Output the final summary build line.
+    # Finish up the build output.
     ninja.process.wait()
     interface.update()
     interface.print_summary()
+
+    if args.write_trace is not None:
+        ninja.write_trace(args.write_trace)
+        print(f'Wrote trace to {args.write_trace.name}')
+
     return ninja.process.returncode