1# Copyright 2023 The Pigweed Authors 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); you may not 4# use this file except in compliance with the License. You may obtain a copy of 5# the License at 6# 7# https://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT 11# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the 12# License for the specific language governing permissions and limitations under 13# the License. 14"""Fetch active Project Builder Context.""" 15 16from __future__ import annotations 17 18import asyncio 19import concurrent.futures 20from contextvars import ContextVar 21from datetime import datetime 22from dataclasses import dataclass, field 23from enum import Enum 24import logging 25import os 26import subprocess 27import time 28from typing import Callable, NoReturn, TYPE_CHECKING 29 30from prompt_toolkit.formatted_text import ( 31 AnyFormattedText, 32 StyleAndTextTuples, 33) 34from prompt_toolkit.key_binding import KeyBindings 35from prompt_toolkit.layout.dimension import AnyDimension, D 36from prompt_toolkit.layout import ( 37 AnyContainer, 38 DynamicContainer, 39 FormattedTextControl, 40 Window, 41) 42from prompt_toolkit.shortcuts import ProgressBar, ProgressBarCounter 43from prompt_toolkit.shortcuts.progress_bar.formatters import ( 44 Formatter, 45 TimeElapsed, 46) 47from prompt_toolkit.shortcuts.progress_bar import formatters 48 49from pw_build.build_recipe import BuildRecipe 50 51if TYPE_CHECKING: 52 from pw_build.project_builder import ProjectBuilder 53 54_LOG = logging.getLogger('pw_build.watch') 55 56 57def _wait_for_terminate_then_kill( 58 proc: subprocess.Popen, timeout: int = 5 59) -> int: 60 """Wait for a process to end, then kill it if the timeout expires.""" 61 returncode = 1 62 try: 63 returncode = proc.wait(timeout=timeout) 64 except subprocess.TimeoutExpired: 65 proc_command = proc.args 66 if isinstance(proc.args, list): 67 proc_command = ' '.join(proc.args) 68 _LOG.debug('Killing %s', proc_command) 69 proc.kill() 70 return returncode 71 72 73class ProjectBuilderState(Enum): 74 IDLE = 'IDLE' 75 BUILDING = 'BUILDING' 76 ABORT = 'ABORT' 77 78 79# pylint: disable=unused-argument 80# Prompt Toolkit progress bar formatter classes follow: 81class BuildStatus(Formatter): 82 """Return OK/FAIL status and last output line for each build.""" 83 84 def __init__(self, ctx) -> None: 85 self.ctx = ctx 86 87 def format( 88 self, 89 progress_bar: ProgressBar, 90 progress: ProgressBarCounter, 91 width: int, 92 ) -> AnyFormattedText: 93 for cfg in self.ctx.recipes: 94 if cfg.display_name != progress.label: 95 continue 96 97 build_status: StyleAndTextTuples = [] 98 build_status.append( 99 cfg.status.status_slug(restarting=self.ctx.restart_flag) 100 ) 101 build_status.append(('', ' ')) 102 build_status.extend(cfg.status.current_step_formatted()) 103 104 return build_status 105 106 return [('', '')] 107 108 def get_width( # pylint: disable=no-self-use 109 self, progress_bar: ProgressBar 110 ) -> AnyDimension: 111 return D() 112 113 114class TimeElapsedIfStarted(TimeElapsed): 115 """Display the elapsed time if the build has started.""" 116 117 def __init__(self, ctx) -> None: 118 self.ctx = ctx 119 120 def format( 121 self, 122 progress_bar: ProgressBar, 123 progress: ProgressBarCounter, 124 width: int, 125 ) -> AnyFormattedText: 126 formatted_text: StyleAndTextTuples = [('', '')] 127 for cfg in self.ctx.recipes: 128 if cfg.display_name != progress.label: 129 continue 130 if cfg.status.started: 131 return super().format(progress_bar, progress, width) 132 return formatted_text 133 134 135# pylint: enable=unused-argument 136 137 138@dataclass 139class ProjectBuilderContext: # pylint: disable=too-many-instance-attributes,too-many-public-methods 140 """Maintains the state of running builds and active subproccesses.""" 141 142 current_state: ProjectBuilderState = ProjectBuilderState.IDLE 143 desired_state: ProjectBuilderState = ProjectBuilderState.BUILDING 144 procs: dict[BuildRecipe, subprocess.Popen] = field(default_factory=dict) 145 recipes: list[BuildRecipe] = field(default_factory=list) 146 147 def __post_init__(self) -> None: 148 self.project_builder: ProjectBuilder | None = None 149 150 self.progress_bar_formatters = [ 151 formatters.Text(' '), 152 formatters.Label(), 153 formatters.Text(' '), 154 BuildStatus(self), 155 formatters.Text(' '), 156 TimeElapsedIfStarted(self), 157 formatters.Text(' '), 158 ] 159 160 self._progress_bar_refresh_interval: float = 0.1 # 10 FPS 161 self._last_progress_bar_redraw_time: float = 0.0 162 163 self._enter_callback: Callable | None = None 164 165 key_bindings = KeyBindings() 166 167 @key_bindings.add('enter') 168 def _enter_pressed(_event): 169 """Run enter press function.""" 170 if self._enter_callback: 171 self._enter_callback() 172 173 self.key_bindings = key_bindings 174 175 self.progress_bar: ProgressBar | None = None 176 177 self._progress_bar_started: bool = False 178 179 self.bottom_toolbar: AnyFormattedText = None 180 self.horizontal_separator = '━' 181 self.title_bar_container: AnyContainer = Window( 182 char=self.horizontal_separator, height=1 183 ) 184 185 self.using_fullscreen: bool = False 186 self.restart_flag: bool = False 187 self.ctrl_c_pressed: bool = False 188 189 def using_progress_bars(self) -> bool: 190 return bool(self.progress_bar) or self.using_fullscreen 191 192 @property 193 def log_build_steps(self) -> bool: 194 if self.project_builder: 195 return self.project_builder.log_build_steps 196 return False 197 198 def interrupted(self) -> bool: 199 return self.ctrl_c_pressed or self.restart_flag 200 201 def set_bottom_toolbar(self, text: AnyFormattedText) -> None: 202 self.bottom_toolbar = text 203 204 def set_enter_callback(self, callback: Callable) -> None: 205 self._enter_callback = callback 206 207 def ctrl_c_interrupt(self) -> None: 208 """Abort function for when using ProgressBars.""" 209 self.ctrl_c_pressed = True 210 self.exit(1) 211 212 def startup_progress(self) -> None: 213 self.progress_bar = ProgressBar( 214 formatters=self.progress_bar_formatters, 215 key_bindings=self.key_bindings, 216 title=self.get_title_bar_text, 217 bottom_toolbar=self.bottom_toolbar, 218 cancel_callback=self.ctrl_c_interrupt, 219 ) 220 self.progress_bar.__enter__() # pylint: disable=unnecessary-dunder-call 221 222 self.create_title_bar_container() 223 self.progress_bar.app.layout.container.children[ # type: ignore 224 0 225 ] = DynamicContainer(lambda: self.title_bar_container) 226 self._progress_bar_started = True 227 228 def exit_progress(self) -> None: 229 if not self.progress_bar: 230 return 231 self.progress_bar.__exit__() # pylint: disable=unnecessary-dunder-call 232 233 def clear_progress_scrollback(self) -> None: 234 if not self.progress_bar: 235 return 236 if ( 237 self.progress_bar.app.is_running 238 and self.progress_bar.app.loop is not None 239 ): 240 self.progress_bar.app.loop.call_soon_threadsafe( 241 self.progress_bar.app.renderer.clear 242 ) 243 244 def redraw_progress(self) -> None: 245 if not self.progress_bar: 246 return 247 if hasattr(self.progress_bar, 'app'): 248 redraw_time = time.time() 249 # Has enough time passed since last redraw? 250 if redraw_time > ( 251 self._last_progress_bar_redraw_time 252 + self._progress_bar_refresh_interval 253 ): 254 # Update last redraw time 255 self._last_progress_bar_redraw_time = redraw_time 256 # Trigger Prompt Toolkit UI redraw. 257 self.progress_bar.invalidate() 258 259 def get_title_style(self) -> str: 260 if self.restart_flag: 261 return 'fg:ansiyellow' 262 263 # Assume passing 264 style = 'fg:ansigreen' 265 266 if self.current_state == ProjectBuilderState.BUILDING: 267 style = 'fg:ansiyellow' 268 269 for cfg in self.recipes: 270 if cfg.status.failed(): 271 style = 'fg:ansired' 272 273 return style 274 275 def exit_code(self) -> int: 276 """Returns a 0 for success, 1 for fail.""" 277 for cfg in self.recipes: 278 if cfg.status.failed(): 279 return 1 280 return 0 281 282 def get_title_bar_text( 283 self, include_separators: bool = True 284 ) -> StyleAndTextTuples: 285 title = '' 286 287 fail_count = 0 288 done_count = 0 289 for cfg in self.recipes: 290 if cfg.status.failed(): 291 fail_count += 1 292 if cfg.status.done: 293 done_count += 1 294 295 if self.restart_flag: 296 title = 'INTERRUPT' 297 elif fail_count > 0: 298 title = f'FAILED ({fail_count})' 299 elif self.current_state == ProjectBuilderState.IDLE and done_count > 0: 300 title = 'PASS' 301 else: 302 title = self.current_state.name 303 304 prefix = '' 305 if include_separators: 306 prefix += f'{self.horizontal_separator}{self.horizontal_separator} ' 307 308 return [(self.get_title_style(), f'{prefix}{title} ')] 309 310 def create_title_bar_container(self) -> None: 311 title_text = FormattedTextControl(self.get_title_bar_text) 312 self.title_bar_container = Window( 313 title_text, 314 char=self.horizontal_separator, 315 height=1, 316 # Expand width to max available space 317 dont_extend_width=False, 318 style=self.get_title_style, 319 ) 320 321 def add_progress_bars(self) -> None: 322 if not self._progress_bar_started: 323 self.startup_progress() 324 assert self.progress_bar 325 self.clear_progress_bars() 326 for cfg in self.recipes: 327 self.progress_bar(label=cfg.display_name, total=len(cfg.steps)) 328 329 def clear_progress_bars(self) -> None: 330 if not self.progress_bar: 331 return 332 self.progress_bar.counters = [] 333 334 def mark_progress_step_complete(self, recipe: BuildRecipe) -> None: 335 if not self.progress_bar: 336 return 337 for pbc in self.progress_bar.counters: 338 if pbc.label == recipe.display_name: 339 pbc.item_completed() 340 break 341 342 def mark_progress_done(self, recipe: BuildRecipe) -> None: 343 if not self.progress_bar: 344 return 345 346 for pbc in self.progress_bar.counters: 347 if pbc.label == recipe.display_name: 348 pbc.done = True 349 break 350 351 def mark_progress_started(self, recipe: BuildRecipe) -> None: 352 if not self.progress_bar: 353 return 354 355 for pbc in self.progress_bar.counters: 356 if pbc.label == recipe.display_name: 357 pbc.start_time = datetime.now() 358 break 359 360 def register_process( 361 self, recipe: BuildRecipe, proc: subprocess.Popen 362 ) -> None: 363 self.procs[recipe] = proc 364 365 def terminate_and_wait( 366 self, 367 exit_message: str | None = None, 368 ) -> None: 369 """End a subproces either cleanly or with a kill signal.""" 370 if self.is_idle() or self.should_abort(): 371 return 372 373 self._signal_abort() 374 375 with concurrent.futures.ThreadPoolExecutor( 376 max_workers=len(self.procs) 377 ) as executor: 378 futures = [] 379 for _recipe, proc in self.procs.items(): 380 if proc is None: 381 continue 382 383 proc_command = proc.args 384 if isinstance(proc.args, list): 385 proc_command = ' '.join(proc.args) 386 _LOG.debug('Stopping: %s', proc_command) 387 388 futures.append( 389 executor.submit(_wait_for_terminate_then_kill, proc) 390 ) 391 for future in concurrent.futures.as_completed(futures): 392 future.result() 393 394 if exit_message: 395 _LOG.info(exit_message) 396 self.set_idle() 397 398 def _signal_abort(self) -> None: 399 self.desired_state = ProjectBuilderState.ABORT 400 401 def build_stopping(self) -> bool: 402 """Return True if the build is restarting or quitting.""" 403 return self.should_abort() or self.interrupted() 404 405 def should_abort(self) -> bool: 406 """Return True if the build is restarting.""" 407 return self.desired_state == ProjectBuilderState.ABORT 408 409 def is_building(self) -> bool: 410 return self.current_state == ProjectBuilderState.BUILDING 411 412 def is_idle(self) -> bool: 413 return self.current_state == ProjectBuilderState.IDLE 414 415 def set_project_builder(self, project_builder) -> None: 416 self.project_builder = project_builder 417 self.recipes = project_builder.build_recipes 418 419 def set_idle(self) -> None: 420 self.current_state = ProjectBuilderState.IDLE 421 self.desired_state = ProjectBuilderState.IDLE 422 423 def set_building(self) -> None: 424 self.restart_flag = False 425 self.current_state = ProjectBuilderState.BUILDING 426 self.desired_state = ProjectBuilderState.BUILDING 427 428 def restore_stdout_logging(self) -> None: # pylint: disable=no-self-use 429 if not self.using_progress_bars(): 430 return 431 432 # Restore logging to STDOUT 433 stdout_handler = logging.StreamHandler() 434 if self.project_builder: 435 stdout_handler.setLevel(self.project_builder.default_log_level) 436 else: 437 stdout_handler.setLevel(logging.INFO) 438 root_logger = logging.getLogger() 439 if self.project_builder and self.project_builder.stdout_proxy: 440 self.project_builder.stdout_proxy.flush() 441 self.project_builder.stdout_proxy.close() 442 root_logger.addHandler(stdout_handler) 443 444 def restore_logging_and_shutdown( 445 self, 446 log_after_shutdown: Callable[[], None] | None = None, 447 ) -> None: 448 self.restore_stdout_logging() 449 _LOG.warning('Abort signal received, stopping processes...') 450 if log_after_shutdown: 451 log_after_shutdown() 452 self.terminate_and_wait() 453 # Flush all log handlers 454 # logging.shutdown() 455 456 def exit( 457 self, 458 exit_code: int = 1, 459 log_after_shutdown: Callable[[], None] | None = None, 460 ) -> None: 461 """Exit function called when the user presses ctrl-c.""" 462 463 # Note: The correct way to exit Python is via sys.exit() however this 464 # takes a number of seconds when running pw_watch with multiple parallel 465 # builds. Instead, this function calls os._exit() to shutdown 466 # immediately. This is similar to `pw_watch.watch._exit`: 467 # https://cs.opensource.google/pigweed/pigweed/+/main:pw_watch/py/pw_watch/watch.py?q=_exit.code 468 469 if not self.progress_bar: 470 self.restore_logging_and_shutdown(log_after_shutdown) 471 logging.shutdown() 472 os._exit(exit_code) # pylint: disable=protected-access 473 474 # Shut everything down after the progress_bar exits. 475 def _really_exit(future: asyncio.Future) -> NoReturn: 476 self.restore_logging_and_shutdown(log_after_shutdown) 477 logging.shutdown() 478 os._exit(future.result()) # pylint: disable=protected-access 479 480 if self.progress_bar.app.future: 481 self.progress_bar.app.future.add_done_callback(_really_exit) 482 self.progress_bar.app.exit(result=exit_code) # type: ignore 483 484 485PROJECT_BUILDER_CONTEXTVAR = ContextVar( 486 'pw_build_project_builder_state', default=ProjectBuilderContext() 487) 488 489 490def get_project_builder_context(): 491 return PROJECT_BUILDER_CONTEXTVAR.get() 492