1#!/usr/bin/env python3 2# Copyright 2022 The Pigweed Authors 3# 4# Licensed under the Apache License, Version 2.0 (the "License"); you may not 5# use this file except in compliance with the License. You may obtain a copy of 6# the License at 7# 8# https://www.apache.org/licenses/LICENSE-2.0 9# 10# Unless required by applicable law or agreed to in writing, software 11# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT 12# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the 13# License for the specific language governing permissions and limitations under 14# the License. 15"""Wrapper for Ninja that adds a real-time status interface.""" 16 17import argparse 18import dataclasses 19import enum 20import json 21import os 22import re 23import signal 24import subprocess 25import sys 26import threading 27import time 28from typing import Any, IO 29 30if sys.platform != 'win32': 31 import pty 32 33# The status formatting string for Ninja to use. Includes a sentinel prefix. 34_NINJA_STATUS = '@@!!@@%s,%f,%t>' 35 36# "ANSI" terminal codes for controlling terminal output. This are more-or-less 37# standardized. See https://en.wikipedia.org/wiki/ANSI_escape_code for a more 38# comprehensive list. 39# Terminal code for clearing the entire current line. 40_TERM_CLEAR_LINE_FULL: str = '\x1B[2K' 41# Terminal code for clearing from the cursor to the end of the current line. 42_TERM_CLEAR_LINE_TO_END: str = '\x1B[0K' 43# Terminal code for moving the cursor to the previous line. 44_TERM_MOVE_PREVIOUS_LINE: str = '\x1B[1A' 45# Terminal code for stopping line wrapping. 46_TERM_DISABLE_WRAP: str = '\x1B[?7l' 47# Terminal code for enabling line wrapping. 48_TERM_ENABLE_WRAP: str = '\x1B[?7h' 49 50 51class SignalFlag: 52 """Helper class for intercepting signals and setting a flag to allow 53 gracefully handling them. 54 55 Attributes: 56 signaled: Initially False, becomes True when the process receives a 57 SIGTERM. 58 """ 59 60 def __init__(self) -> None: 61 """Initializes the signal handler.""" 62 self.signaled = False 63 signal.signal(signal.SIGTERM, self._signal_handler) 64 65 # pylint: disable=unused-argument 66 def _signal_handler(self, signal_number: int, stack_frame: Any) -> None: 67 """Sets the signaled flag, called when signaled.""" 68 self.signaled = True 69 70 71class ConsoleRenderer: 72 """Helper class for rendering the TUI using terminal control codes. 73 74 This class has two types of output. First, permanently printed lines 75 (with `print_line` method) behave like normal terminal output and stay in 76 the terminal once printed. Second, "temporary" printed lines follow the 77 most recently printed permanent line, and get rewritten on each call 78 to `render`. 79 80 Attributes: 81 smart_terminal: If true, will print a rich TUI using terminal control 82 codes. Otherwise, this class won't print temporary lines, and 83 permanent lines will be printed without any special control codes. 84 Defaults to true if stdout is connected to a TTY. 85 """ 86 87 def __init__(self) -> None: 88 """Initialize the console renderer.""" 89 self._queued_lines: list[str] = [] 90 self._temporary_lines: list[str] = [] 91 self._previous_line_count = 0 92 93 term = os.environ.get('TERM') 94 self.smart_terminal = term and (term != 'dumb') and sys.stdout.isatty() 95 96 def print_line(self, line: str) -> None: 97 """Queue a permanent line for printing during the next render.""" 98 self._queued_lines.append(line) 99 100 def set_temporary_lines(self, lines: list[str]) -> None: 101 """Set the temporary lines to be displayed during the next render.""" 102 self._temporary_lines = lines[:] 103 104 def render(self) -> None: 105 """Render the current state of the renderer.""" 106 107 # If we can't use terminal codes, print out permanent lines and exit. 108 if not self.smart_terminal: 109 for line in self._queued_lines: 110 print(line) 111 self._queued_lines.clear() 112 sys.stdout.flush() 113 return 114 115 # Go back to the end of the last permanent lines. 116 for _ in range(self._previous_line_count): 117 print(_TERM_MOVE_PREVIOUS_LINE, end='') 118 119 # Print any new permanent lines. 120 for line in self._queued_lines: 121 print(line, end='') 122 print(_TERM_CLEAR_LINE_TO_END) 123 124 # Print the new temporary lines. 125 print(_TERM_DISABLE_WRAP, end='') 126 for line in self._temporary_lines: 127 print(_TERM_CLEAR_LINE_FULL, end='') 128 print(line, end='') 129 print(_TERM_CLEAR_LINE_TO_END) 130 print(_TERM_ENABLE_WRAP, end='') 131 132 # Clear any leftover temporary lines from the previous render. 133 lines_to_clear = self._previous_line_count - len(self._temporary_lines) 134 for _ in range(lines_to_clear): 135 print(_TERM_CLEAR_LINE_FULL) 136 for _ in range(lines_to_clear): 137 print(_TERM_MOVE_PREVIOUS_LINE, end='') 138 139 # Flush and update state. 140 sys.stdout.flush() 141 self._previous_line_count = len(self._temporary_lines) 142 self._queued_lines.clear() 143 144 145def _format_duration(duration: float) -> str: 146 """Format a duration (in seconds) as a string.""" 147 if duration >= 60: 148 seconds = int(duration % 60) 149 minutes = int(duration / 60) 150 return f'{minutes}m{seconds:02}s' 151 return f'{duration:.1f}s' 152 153 154@dataclasses.dataclass 155class NinjaAction: 156 """A Ninja action: a task that needs to run and complete. 157 158 In Ninja, this is referred to as an "edge". 159 160 Attributes: 161 name: The name of the Ninja action. 162 jobs: The number of running jobs for the action. Some Ninja actions 163 have multiple sub-actions that have the same name. 164 start_time: The time that the Ninja action started. This is based on 165 time.time, not Ninja's own stopwatch. 166 end_time: The time that the Ninja action finished, or None if it is 167 still running. 168 """ 169 170 name: str 171 jobs: int = 0 172 start_time: float = dataclasses.field(default_factory=time.time) 173 end_time: float | None = None 174 175 176class NinjaEventKind(enum.Enum): 177 """The kind of Ninja event.""" 178 179 ACTION_STARTED = 1 180 ACTION_FINISHED = 2 181 ACTION_LOG = 3 182 183 184@dataclasses.dataclass 185class NinjaEvent: 186 """An event from the Ninja build. 187 188 Attributes: 189 kind: The kind of event. 190 action: The action this event relates to, if any. 191 message: The log message associated with this event, if it an 192 'ACTION_LOG' event. 193 num_started: The number of started actions when this event occurred. 194 num_finished: The number of finished actions when this event occurred. 195 num_total: The total number of actions when this event occurred. 196 """ 197 198 kind: NinjaEventKind 199 action: NinjaAction | None = None 200 log_message: str | None = None 201 num_started: int = 0 202 num_finished: int = 0 203 num_total: int = 0 204 205 206class Ninja: 207 """Wrapper around a Ninja subprocess. 208 209 Parses the Ninja output to maintain a representation of the build graph. 210 211 Attributes: 212 num_started: The number of started actions. 213 num_finished: The number of finished actions. 214 num_total: The (estimated) number of total actions in the build. 215 running_actions: The currently-running actions, keyed by name. 216 actions: A list of all actions, running and completed. 217 last_action_completed: The last action that Ninja finished building. 218 events: The events that have occured in the Ninja build. 219 exited: Whether the Ninja subprocess has exited. 220 lock: The lock guarding the rest of the attributes. 221 process: The Python subprocess running Ninja. 222 start_time: The time that the Ninja build started. 223 """ 224 225 num_started: int 226 num_finished: int 227 num_total: int 228 actions: list[NinjaAction] 229 running_actions: dict[str, NinjaAction] 230 last_action_completed: NinjaAction | None 231 events: list[NinjaEvent] 232 exited: bool 233 lock: threading.Lock 234 process: subprocess.Popen 235 start_time: float 236 237 def __init__(self, args: list[str]) -> None: 238 if sys.platform == 'win32': 239 raise RuntimeError('Ninja wrapper does not support Windows.') 240 241 self.num_started = 0 242 self.num_finished = 0 243 self.num_total = 0 244 self.actions = [] 245 self.running_actions = {} 246 self.last_action_completed = None 247 self.events = [] 248 self.exited = False 249 self.lock = threading.Lock() 250 251 # Launch ninja and configure pseudo-tty. 252 # pylint: disable-next=no-member,undefined-variable,used-before-assignment 253 ptty_parent, ptty_child = pty.openpty() # type: ignore 254 ptty_file = os.fdopen(ptty_parent, 'r') 255 env = dict(os.environ) 256 env['NINJA_STATUS'] = _NINJA_STATUS 257 if 'TERM' not in env: 258 env['TERM'] = 'vt100' 259 command = ['ninja'] + args 260 self.process = subprocess.Popen( 261 command, 262 env=env, 263 stdin=subprocess.DEVNULL, 264 stdout=ptty_child, 265 stderr=ptty_child, 266 ) 267 os.close(ptty_child) 268 self.start_time = time.time() 269 270 # Start the output monitor thread. 271 thread = threading.Thread( 272 target=self._monitor_thread, args=(ptty_file,) 273 ) 274 thread.daemon = True # Don't keep the process alive. 275 thread.start() 276 277 def _monitor_thread(self, file: IO[str]) -> None: 278 """Monitor the Ninja output. Run in a separate thread.""" 279 # A Ninja status line starts with "\r" and ends with "\x1B[K". This 280 # tracks whether we're currently in a status line. 281 ninja_status = False 282 buffer: list[str] = [] 283 try: 284 while True: 285 char = file.read(1) 286 if char == '\r': 287 ninja_status = True 288 continue 289 if char == '\n': 290 self._process_output(''.join(buffer)) 291 buffer = [] 292 continue 293 if char == '': 294 # End of file. 295 break 296 297 # Look for the end of a status line, ignoring partial matches. 298 if char == '\x1B' and ninja_status: 299 char = file.read(1) 300 if char == '[': 301 char = file.read(1) 302 if char == 'K': 303 self._process_output(''.join(buffer)) 304 buffer = [] 305 ninja_status = False 306 continue 307 buffer.append('\x1B[') 308 else: 309 buffer.append('\x1B') 310 311 buffer.append(char) 312 except OSError: 313 pass 314 315 self._process_output(''.join(buffer)) 316 self.exited = True 317 318 def _process_output(self, line: str) -> None: 319 """Process a line of output from Ninja, updating the internal state.""" 320 with self.lock: 321 if match := re.match(r'^@@!!@@(\d+),(\d+),(\d+)>(.*)', line): 322 actions_started = int(match.group(1)) 323 actions_finished = int(match.group(2)) 324 actions_total = int(match.group(3)) 325 name = match.group(4) 326 327 # Sometimes Ninja delimits lines without \r, which prevents 328 # _monitor_thread from stripping out the final control code, 329 # so just remove it here if it's present. 330 if name.endswith('\x1B[K'): 331 name = name[:-3] 332 333 did_start = actions_started > self.num_started 334 did_finish = actions_finished > self.num_finished 335 self.num_started = actions_started 336 self.num_finished = actions_finished 337 self.num_total = actions_total 338 339 # Special case: first action in a new graph. 340 # This is important for GN's "Regenerating ninja files" action. 341 if actions_started == 1 and actions_finished == 0: 342 for action in self.running_actions.values(): 343 action.end_time = time.time() 344 self._add_event( 345 NinjaEvent(NinjaEventKind.ACTION_FINISHED, action) 346 ) 347 self.running_actions = {} 348 self.last_action_completed = None 349 350 if did_start: 351 action = self.running_actions.setdefault( 352 name, NinjaAction(name) 353 ) 354 if action.jobs == 0: 355 self.actions.append(action) 356 self._add_event( 357 NinjaEvent(NinjaEventKind.ACTION_STARTED, action) 358 ) 359 action.jobs += 1 360 361 if did_finish and name in self.running_actions: 362 action = self.running_actions[name] 363 action.jobs -= 1 364 if action.jobs <= 0: 365 self.running_actions.pop(name) 366 self.last_action_completed = action 367 action.end_time = time.time() 368 self._add_event( 369 NinjaEvent(NinjaEventKind.ACTION_FINISHED, action) 370 ) 371 else: 372 context_action = None 373 if not line.startswith('ninja: '): 374 context_action = self.last_action_completed 375 self._add_event( 376 NinjaEvent( 377 NinjaEventKind.ACTION_LOG, 378 action=context_action, 379 log_message=line, 380 ) 381 ) 382 383 def _add_event(self, event: NinjaEvent) -> None: 384 """Add a new event to the event queue.""" 385 event.num_started = self.num_started 386 event.num_finished = self.num_finished 387 event.num_total = self.num_total 388 self.events.append(event) 389 390 def write_trace(self, file: IO[str]) -> None: 391 """Write a Chromium trace_event-formatted trace to a file.""" 392 now = time.time() 393 threads: list[float] = [] 394 events = [] 395 actions = sorted( 396 self.actions, key=lambda x: x.end_time or now, reverse=True 397 ) 398 for action in actions: 399 # If this action hasn't completed, fake the end time. 400 end_time = action.end_time or now 401 402 # Allocate a "thread id" for the action. We look at threads in 403 # reverse end time and fill in threads from the end to get a 404 # better-looking trace. 405 for tid, busy_time in enumerate(threads): 406 if busy_time >= end_time: 407 threads[tid] = action.start_time 408 break 409 else: 410 tid = len(threads) 411 threads.append(action.start_time) 412 413 events.append( 414 { 415 'name': action.name, 416 'cat': 'actions', 417 'ph': 'X', 418 'ts': action.start_time * 1000000, 419 'dur': (end_time - action.start_time) * 1000000, 420 'pid': 0, 421 'tid': tid, 422 'args': {}, 423 } 424 ) 425 json.dump(events, file) 426 427 428class UI: 429 """Class to handle UI state and rendering.""" 430 431 def __init__(self, ninja: Ninja, args: argparse.Namespace) -> None: 432 self._ninja = ninja 433 self._args = args 434 self._renderer = ConsoleRenderer() 435 self._last_log_action: NinjaAction | None = None 436 437 def _get_status_display(self) -> list[str]: 438 """Generates the status display. Must be called under the Ninja lock.""" 439 actions = sorted( 440 self._ninja.running_actions.values(), key=lambda x: x.start_time 441 ) 442 now = time.time() 443 total_elapsed = _format_duration(now - self._ninja.start_time) 444 lines = [ 445 f'[{total_elapsed: >5}] ' 446 f'Building [{self._ninja.num_finished}/{self._ninja.num_total}] ...' 447 ] 448 for action in actions[: self._args.ui_max_actions]: 449 elapsed = _format_duration(now - action.start_time) 450 lines.append(f' [{elapsed: >5}] {action.name}') 451 if len(actions) > self._args.ui_max_actions: 452 remaining = len(actions) - self._args.ui_max_actions 453 lines.append(f' ... and {remaining} more') 454 return lines 455 456 def _process_event(self, event: NinjaEvent) -> None: 457 """Processes a Ninja Event. Must be called under the Ninja lock.""" 458 show_started = self._args.log_actions 459 show_ended = self._args.log_actions or not self._renderer.smart_terminal 460 461 if event.kind == NinjaEventKind.ACTION_LOG: 462 if event.action and (event.action != self._last_log_action): 463 self._renderer.print_line(f'[{event.action.name}]') 464 self._last_log_action = event.action 465 assert event.log_message is not None 466 self._renderer.print_line(event.log_message) 467 468 if event.kind == NinjaEventKind.ACTION_STARTED and show_started: 469 assert event.action 470 self._renderer.print_line( 471 f'[{event.num_finished}/{event.num_total}] ' 472 f'Started [{event.action.name}]' 473 ) 474 475 if event.kind == NinjaEventKind.ACTION_FINISHED and show_ended: 476 assert event.action and event.action.end_time is not None 477 duration = _format_duration( 478 event.action.end_time - event.action.start_time 479 ) 480 self._renderer.print_line( 481 f'[{event.num_finished}/{event.num_total}] ' 482 f'Finished [{event.action.name}] ({duration})' 483 ) 484 485 def update(self) -> None: 486 """Updates and re-renders the UI.""" 487 with self._ninja.lock: 488 for event in self._ninja.events: 489 self._process_event(event) 490 self._ninja.events = [] 491 492 self._renderer.set_temporary_lines(self._get_status_display()) 493 494 self._renderer.render() 495 496 def print_summary(self) -> None: 497 """Prints the summary line at the end of the build.""" 498 total_time = _format_duration(time.time() - self._ninja.start_time) 499 num_finished = self._ninja.num_finished 500 num_total = self._ninja.num_total 501 self._renderer.print_line( 502 f'Built {num_finished}/{num_total} targets in {total_time}.' 503 ) 504 self._renderer.set_temporary_lines([]) 505 self._renderer.render() 506 507 def dump_running_actions(self) -> None: 508 """Prints a list of actions that are currently running.""" 509 actions = sorted( 510 self._ninja.running_actions.values(), key=lambda x: x.start_time 511 ) 512 now = time.time() 513 self._renderer.print_line(f'{len(actions)} running actions:') 514 for action in actions: 515 elapsed = _format_duration(now - action.start_time) 516 self._renderer.print_line(f' [{elapsed: >5}] {action.name}') 517 self._renderer.set_temporary_lines([]) 518 self._renderer.render() 519 520 def print_line(self, line: str) -> None: 521 """Print a line to the interface output.""" 522 self._renderer.print_line(line) 523 524 525def _parse_args() -> tuple[argparse.Namespace, list[str]]: 526 """Parses CLI arguments. 527 528 Returns: 529 The tuple containing the parsed arguments and the remaining unparsed 530 arguments to be passed to Ninja. 531 """ 532 parser = argparse.ArgumentParser(description=__doc__) 533 parser.add_argument( 534 '--ui-update-rate', 535 help='update rate of the UI (in seconds)', 536 type=float, 537 default=0.1, 538 ) 539 parser.add_argument( 540 '--ui-max-actions', 541 help='maximum build actions to display at once', 542 type=int, 543 default=8, 544 ) 545 parser.add_argument( 546 '--log-actions', 547 help='whether to log when actions start and finish', 548 action='store_true', 549 ) 550 parser.add_argument( 551 '--write-trace', 552 help='write a Chromium trace_event-style trace to the given file', 553 type=argparse.FileType('w'), 554 ) 555 return parser.parse_known_args() 556 557 558def main() -> int: 559 """Main entrypoint for script.""" 560 args, ninja_args = _parse_args() 561 562 if sys.platform == 'win32': 563 print( 564 'WARNING: pw-wrap-ninja does not support Windows. ' 565 'Running ninja directly.' 566 ) 567 process = subprocess.run(['ninja'] + ninja_args) 568 return process.returncode 569 570 signal_flag = SignalFlag() 571 ninja = Ninja(ninja_args) 572 interface = UI(ninja, args) 573 574 while True: 575 interface.update() 576 if signal_flag.signaled: 577 interface.print_line('Got SIGTERM, exiting...') 578 ninja.process.terminate() 579 break 580 if ninja.exited: 581 break 582 time.sleep(args.ui_update_rate) 583 584 # Finish up the build output. 585 ninja.process.wait() 586 interface.update() 587 interface.print_summary() 588 589 if args.write_trace is not None: 590 ninja.write_trace(args.write_trace) 591 print(f'Wrote trace to {args.write_trace.name}') 592 593 if signal_flag.signaled: 594 interface.dump_running_actions() 595 596 return ninja.process.returncode 597 598 599if __name__ == '__main__': 600 sys.exit(main()) 601