| #!/usr/bin/env python3 |
| # Copyright 2021 The Chromium Authors |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| """Creates an server to offload non-critical-path GN targets.""" |
| |
| from __future__ import annotations |
| |
| import argparse |
| import collections |
| import contextlib |
| import dataclasses |
| import datetime |
| import os |
| import pathlib |
| import re |
| import signal |
| import shlex |
| import shutil |
| import socket |
| import subprocess |
| import sys |
| import threading |
| import traceback |
| import time |
| from typing import Callable, Dict, List, Optional, Tuple, IO |
| |
| sys.path.append(os.path.join(os.path.dirname(__file__), 'gyp')) |
| from util import server_utils |
| |
| _SOCKET_TIMEOUT = 60 # seconds |
| |
| _LOGFILE_NAME = 'buildserver.log' |
| _MAX_LOGFILES = 6 |
| |
| FIRST_LOG_LINE = """\ |
| #### Start of log for build: {build_id} |
| #### CWD: {outdir} |
| """ |
| BUILD_ID_RE = re.compile(r'^#### Start of log for build: (?P<build_id>.+)') |
| |
| |
| def server_log(msg: str): |
| if OptionsManager.is_quiet(): |
| return |
| # Ensure we start our message on a new line. |
| print('\n' + msg) |
| |
| |
| def print_status(prefix: str, msg: str): |
| # No need to also output to the terminal if quiet. |
| if OptionsManager.is_quiet(): |
| return |
| # Shrink the message (leaving a 2-char prefix and use the rest of the room |
| # for the suffix) according to terminal size so it is always one line. |
| width = shutil.get_terminal_size().columns |
| max_msg_width = width - len(prefix) |
| if len(msg) > max_msg_width: |
| length_to_show = max_msg_width - 5 # Account for ellipsis and header. |
| msg = f'{msg[:2]}...{msg[-length_to_show:]}' |
| # \r to return the carriage to the beginning of line. |
| # \033[K to replace the normal \n to erase until the end of the line. |
| # Avoid the default line ending so the next \r overwrites the same line just |
| # like ninja's output. |
| print(f'\r{prefix}{msg}\033[K', end='', flush=True) |
| |
| |
| def _exception_hook(exctype: type, exc: Exception, tb): |
| # Let KeyboardInterrupt through. |
| if issubclass(exctype, KeyboardInterrupt): |
| sys.__excepthook__(exctype, exc, tb) |
| return |
| stacktrace = ''.join(traceback.format_exception(exctype, exc, tb)) |
| stacktrace_lines = [f'\n⛔{line}' for line in stacktrace.splitlines()] |
| # Output uncaught exceptions to all live terminals |
| # Extra newline since siso's output often erases the current line. |
| BuildManager.broadcast(''.join(stacktrace_lines) + '\n') |
| # Cancel all pending tasks cleanly (i.e. delete stamp files if necessary). |
| TaskManager.deactivate() |
| # Reset all remote terminal titles. |
| BuildManager.update_remote_titles('') |
| |
| |
| # Stores global options so as to not keep passing along and storing options |
| # everywhere. |
| class OptionsManager: |
| _quiet = None |
| _should_remote_print = None |
| |
| @classmethod |
| def set_options(cls, *, quiet, should_remote_print): |
| cls._quiet = quiet |
| cls._should_remote_print = should_remote_print |
| |
| @classmethod |
| def is_quiet(cls): |
| assert cls._quiet is not None |
| return cls._quiet |
| |
| @classmethod |
| def should_remote_print(cls): |
| assert cls._should_remote_print is not None |
| return cls._should_remote_print |
| |
| |
| class LogfileManager: |
| _logfiles: dict[str, IO[str]] = {} |
| _lock = threading.RLock() |
| |
| @classmethod |
| def create_logfile(cls, build_id, outdir): |
| with cls._lock: |
| if logfile := cls._logfiles.get(build_id, None): |
| return logfile |
| |
| outdir = pathlib.Path(outdir) |
| latest_logfile = outdir / f'{_LOGFILE_NAME}.0' |
| |
| if latest_logfile.exists(): |
| with latest_logfile.open('rt') as f: |
| first_line = f.readline() |
| if log_build_id := BUILD_ID_RE.search(first_line): |
| # If the newest logfile on disk is referencing the same build we are |
| # currently processing, we probably crashed previously and we should |
| # pick up where we left off in the same logfile. |
| if log_build_id.group('build_id') == build_id: |
| cls._logfiles[build_id] = latest_logfile.open('at') |
| return cls._logfiles[build_id] |
| |
| # Do the logfile name shift. |
| filenames = os.listdir(outdir) |
| logfiles = {f for f in filenames if f.startswith(_LOGFILE_NAME)} |
| for idx in reversed(range(_MAX_LOGFILES)): |
| current_name = f'{_LOGFILE_NAME}.{idx}' |
| next_name = f'{_LOGFILE_NAME}.{idx+1}' |
| if current_name in logfiles: |
| shutil.move(os.path.join(outdir, current_name), |
| os.path.join(outdir, next_name)) |
| |
| # Create a new 0th logfile. |
| logfile = latest_logfile.open('wt') |
| logfile.write(FIRST_LOG_LINE.format(build_id=build_id, outdir=outdir)) |
| logfile.flush() |
| cls._logfiles[build_id] = logfile |
| return logfile |
| |
| |
| class TaskStats: |
| """Class to keep track of aggregate stats for all tasks across threads.""" |
| _num_processes = 0 |
| _completed_tasks = 0 |
| _total_tasks = 0 |
| _lock = threading.RLock() |
| |
| @classmethod |
| def no_running_processes(cls): |
| with cls._lock: |
| return cls._num_processes == 0 |
| |
| @classmethod |
| def add_task(cls): |
| with cls._lock: |
| cls._total_tasks += 1 |
| |
| @classmethod |
| def add_process(cls): |
| with cls._lock: |
| cls._num_processes += 1 |
| |
| @classmethod |
| def remove_process(cls): |
| with cls._lock: |
| cls._num_processes -= 1 |
| |
| @classmethod |
| def complete_task(cls): |
| with cls._lock: |
| cls._completed_tasks += 1 |
| |
| @classmethod |
| def num_pending_tasks(cls): |
| with cls._lock: |
| return cls._total_tasks - cls._completed_tasks |
| |
| @classmethod |
| def num_completed_tasks(cls): |
| with cls._lock: |
| return cls._completed_tasks |
| |
| @classmethod |
| def total_tasks(cls): |
| with cls._lock: |
| return cls._total_tasks |
| |
| @classmethod |
| def get_title_message(cls): |
| with cls._lock: |
| return f'Analysis Steps: {cls._completed_tasks}/{cls._total_tasks}' |
| |
| @classmethod |
| def query_build(cls, query_build_id: str = None): |
| builds = [] |
| if query_build_id: |
| if build := BuildManager.get_build(query_build_id): |
| builds.append(build) |
| else: |
| builds = BuildManager.get_all_builds() |
| build_infos = [] |
| for build in builds: |
| build_infos.append(build.query_build_info()) |
| return { |
| 'pid': os.getpid(), |
| 'builds': build_infos, |
| } |
| |
| @classmethod |
| def prefix(cls, build_id: str = None): |
| # Ninja's prefix is: [205 processes, 6/734 @ 6.5/s : 0.922s ] |
| # Time taken and task completion rate are not important for the build server |
| # since it is always running in the background and uses idle priority for |
| # its tasks. |
| with cls._lock: |
| if build_id: |
| build = BuildManager.get_build(build_id) |
| _num_processes = build.process_count() |
| _completed_tasks = build.completed_task_count() |
| _total_tasks = build.total_task_count() |
| else: |
| _num_processes = cls._num_processes |
| _completed_tasks = cls._completed_tasks |
| _total_tasks = cls._total_tasks |
| word = 'process' if _num_processes == 1 else 'processes' |
| return (f'{_num_processes} {word}, ' |
| f'{_completed_tasks}/{_total_tasks}') |
| |
| |
| def check_pid_alive(pid: int): |
| try: |
| os.kill(pid, 0) |
| except OSError: |
| return False |
| return True |
| |
| |
| @dataclasses.dataclass |
| class Build: |
| id: str |
| pid: int |
| env: dict |
| stdout: IO[str] |
| cwd: Optional[str] = None |
| _logfile: Optional[IO[str]] = None |
| _is_ninja_alive: bool = True |
| _tasks: List[Task] = dataclasses.field(default_factory=list) |
| _completed_task_count = 0 |
| _active_process_count = 0 |
| _lock: threading.RLock = dataclasses.field(default_factory=threading.RLock, |
| repr=False, |
| init=False) |
| |
| def __hash__(self): |
| return hash((self.id, self.pid, self.cwd)) |
| |
| def add_task(self, task: Task): |
| self._status_update(f'QUEUED {task.name}') |
| with self._lock: |
| self._tasks.append(task) |
| TaskStats.add_task() |
| TaskManager.add_task(task) |
| |
| def add_process(self, task: Task): |
| self._status_update(f'STARTING {task.name}') |
| with self._lock: |
| self._active_process_count += 1 |
| TaskStats.add_process() |
| |
| def task_done(self, task: Task, status_string: str): |
| self._status_update(f'{status_string} {task.name}') |
| TaskStats.complete_task() |
| TaskManager.task_done(task) |
| with self._lock: |
| self._completed_task_count += 1 |
| |
| # We synchronize all terminal title info rather than having it per build |
| # since if two builds are happening in the same terminal concurrently, both |
| # builds will be overriding each other's titles continuously. Usually we |
| # only have the one build anyways so it should equivalent in most cases. |
| BuildManager.update_remote_titles() |
| with self._lock: |
| if not self.is_active(): |
| self._logfile.close() |
| # Reset in case its the last build. |
| BuildManager.update_remote_titles('') |
| |
| def process_complete(self): |
| with self._lock: |
| self._active_process_count -= 1 |
| TaskStats.remove_process() |
| |
| def ensure_logfile(self): |
| with self._lock: |
| if not self._logfile: |
| assert self.cwd is not None |
| self._logfile = LogfileManager.create_logfile(self.id, self.cwd) |
| |
| def log(self, message: str): |
| with self._lock: |
| self.ensure_logfile() |
| if self._logfile.closed: |
| # BuildManager#broadcast can call log after the build is done and the |
| # log is closed. Might make sense to separate out that flow so we can |
| # raise an exception here otherwise. |
| return |
| print(message, file=self._logfile, flush=True) |
| |
| def _status_update(self, status_message): |
| prefix = f'[{TaskStats.prefix(self.id)}] ' |
| self.log(f'{prefix}{status_message}') |
| print_status(prefix, status_message) |
| |
| def total_task_count(self): |
| with self._lock: |
| return len(self._tasks) |
| |
| def completed_task_count(self): |
| with self._lock: |
| return self._completed_task_count |
| |
| def pending_task_count(self): |
| with self._lock: |
| return self.total_task_count() - self.completed_task_count() |
| |
| def process_count(self): |
| with self._lock: |
| return self._active_process_count |
| |
| def is_active(self): |
| if self.pending_task_count() > 0: |
| return True |
| # Ninja is not coming back to life so only check on it if last we checked it |
| # was still alive. |
| if self._is_ninja_alive: |
| self._is_ninja_alive = check_pid_alive(self.pid) |
| return self._is_ninja_alive |
| |
| def query_build_info(self): |
| current_tasks = TaskManager.get_current_tasks(self.id) |
| return { |
| 'build_id': self.id, |
| 'is_active': self.is_active(), |
| 'completed_tasks': self.completed_task_count(), |
| 'pending_tasks': self.pending_task_count(), |
| 'active_tasks': [t.cmd for t in current_tasks], |
| 'outdir': self.cwd, |
| } |
| |
| |
| class BuildManager: |
| _builds_by_id: dict[str, Build] = dict() |
| _cached_ttys: dict[(int, int), tuple[IO[str], bool]] = dict() |
| _lock = threading.RLock() |
| |
| @classmethod |
| def register_builder(cls, env, pid, cwd): |
| build_id = env['AUTONINJA_BUILD_ID'] |
| stdout = cls.open_tty(env['AUTONINJA_STDOUT_NAME']) |
| # Tells the script not to re-delegate to build server. |
| env[server_utils.BUILD_SERVER_ENV_VARIABLE] = '1' |
| |
| with cls._lock: |
| build = Build(id=build_id, |
| pid=pid, |
| cwd=cwd, |
| env=env, |
| stdout=stdout) |
| cls.maybe_init_cwd(build, cwd) |
| cls._builds_by_id[build_id] = build |
| cls.update_remote_titles() |
| |
| @classmethod |
| def maybe_init_cwd(cls, build: Build, cwd: str): |
| if cwd is not None: |
| with cls._lock: |
| if build.cwd is None: |
| build.cwd = cwd |
| else: |
| assert pathlib.Path(cwd).samefile( |
| build.cwd), f'{repr(cwd)} != {repr(build.cwd)}' |
| build.ensure_logfile() |
| |
| @classmethod |
| def get_build(cls, build_id): |
| with cls._lock: |
| return cls._builds_by_id.get(build_id, None) |
| |
| @classmethod |
| def open_tty(cls, tty_path): |
| # Do not open the same tty multiple times. Use st_ino and st_dev to compare |
| # file descriptors. |
| tty = open(tty_path, 'at') |
| st = os.stat(tty.fileno()) |
| tty_key = (st.st_ino, st.st_dev) |
| with cls._lock: |
| # Dedupes ttys |
| if tty_key not in cls._cached_ttys: |
| # TTYs are kept open for the lifetime of the server so that broadcast |
| # messages (e.g. uncaught exceptions) can be sent to them even if they |
| # are not currently building anything. |
| cls._cached_ttys[tty_key] = (tty, tty.isatty()) |
| else: |
| tty.close() |
| return cls._cached_ttys[tty_key][0] |
| |
| @classmethod |
| def get_active_builds(cls) -> List[Build]: |
| builds = cls.get_all_builds() |
| return list(build for build in builds if build.is_active()) |
| |
| @classmethod |
| def get_all_builds(cls) -> List[Build]: |
| with cls._lock: |
| return list(cls._builds_by_id.values()) |
| |
| @classmethod |
| def broadcast(cls, msg: str): |
| with cls._lock: |
| ttys = list(cls._cached_ttys.values()) |
| builds = list(cls._builds_by_id.values()) |
| if OptionsManager.should_remote_print(): |
| for tty, _unused in ttys: |
| try: |
| tty.write(msg + '\n') |
| tty.flush() |
| except BrokenPipeError: |
| pass |
| for build in builds: |
| build.log(msg) |
| # Write to the current terminal if we have not written to it yet. |
| st = os.stat(sys.stderr.fileno()) |
| stderr_key = (st.st_ino, st.st_dev) |
| if stderr_key not in cls._cached_ttys: |
| print(msg, file=sys.stderr) |
| |
| @classmethod |
| def update_remote_titles(cls, new_title=None): |
| if new_title is None: |
| if not cls.has_active_builds() and TaskStats.num_pending_tasks() == 0: |
| # Setting an empty title causes most terminals to go back to the |
| # default title (and at least prevents the tab title from being |
| # "Analysis Steps: N/N" forevermore. |
| new_title = '' |
| else: |
| new_title = TaskStats.get_title_message() |
| |
| with cls._lock: |
| ttys = list(cls._cached_ttys.values()) |
| for tty, isatty in ttys: |
| if isatty: |
| try: |
| tty.write(f'\033]2;{new_title}\007') |
| tty.flush() |
| except BrokenPipeError: |
| pass |
| |
| @classmethod |
| def has_active_builds(cls): |
| return bool(cls.get_active_builds()) |
| |
| |
| class TaskManager: |
| """Class to encapsulate a threadsafe queue and handle deactivating it.""" |
| _queue: collections.deque[Task] = collections.deque() |
| _current_tasks: set[Task] = set() |
| _deactivated = False |
| _lock = threading.RLock() |
| |
| @classmethod |
| def add_task(cls, task: Task): |
| assert not cls._deactivated |
| with cls._lock: |
| cls._queue.appendleft(task) |
| cls._maybe_start_tasks() |
| |
| @classmethod |
| def task_done(cls, task: Task): |
| with cls._lock: |
| cls._current_tasks.discard(task) |
| |
| @classmethod |
| def get_current_tasks(cls, build_id): |
| with cls._lock: |
| return [t for t in cls._current_tasks if t.build.id == build_id] |
| |
| @classmethod |
| def deactivate(cls): |
| cls._deactivated = True |
| tasks_to_terminate: list[Task] = [] |
| with cls._lock: |
| while cls._queue: |
| task = cls._queue.pop() |
| tasks_to_terminate.append(task) |
| # Cancel possibly running tasks. |
| tasks_to_terminate.extend(cls._current_tasks) |
| # Terminate outside lock since task threads need the lock to finish |
| # terminating. |
| for task in tasks_to_terminate: |
| task.terminate() |
| |
| @classmethod |
| def cancel_build(cls, build_id): |
| terminated_pending_tasks: list[Task] = [] |
| terminated_current_tasks: list[Task] = [] |
| with cls._lock: |
| # Cancel pending tasks. |
| for task in cls._queue: |
| if task.build.id == build_id: |
| terminated_pending_tasks.append(task) |
| for task in terminated_pending_tasks: |
| cls._queue.remove(task) |
| # Cancel running tasks. |
| for task in cls._current_tasks: |
| if task.build.id == build_id: |
| terminated_current_tasks.append(task) |
| # Terminate tasks outside lock since task threads need the lock to finish |
| # terminating. |
| for task in terminated_pending_tasks: |
| task.terminate() |
| for task in terminated_current_tasks: |
| task.terminate() |
| |
| @staticmethod |
| # pylint: disable=inconsistent-return-statements |
| def _num_running_processes(): |
| with open('/proc/stat') as f: |
| for line in f: |
| if line.startswith('procs_running'): |
| return int(line.rstrip().split()[1]) |
| assert False, 'Could not read /proc/stat' |
| |
| @classmethod |
| def _maybe_start_tasks(cls): |
| if cls._deactivated: |
| return |
| # Include load avg so that a small dip in the number of currently running |
| # processes will not cause new tasks to be started while the overall load is |
| # heavy. |
| cur_load = max(cls._num_running_processes(), os.getloadavg()[0]) |
| num_started = 0 |
| # Always start a task if we don't have any running, so that all tasks are |
| # eventually finished. Try starting up tasks when the overall load is light. |
| # Limit to at most 2 new tasks to prevent ramping up too fast. There is a |
| # chance where multiple threads call _maybe_start_tasks and each gets to |
| # spawn up to 2 new tasks, but since the only downside is some build tasks |
| # get worked on earlier rather than later, it is not worth mitigating. |
| while num_started < 2 and (TaskStats.no_running_processes() |
| or num_started + cur_load < os.cpu_count()): |
| with cls._lock: |
| try: |
| next_task = cls._queue.pop() |
| cls._current_tasks.add(next_task) |
| except IndexError: |
| return |
| num_started += next_task.start(cls._maybe_start_tasks) |
| |
| |
| # TODO(wnwen): Break this into Request (encapsulating what ninja sends) and Task |
| # when a Request starts to be run. This would eliminate ambiguity |
| # about when and whether _proc/_thread are initialized. |
| class Task: |
| """Class to represent one task and operations on it.""" |
| |
| def __init__(self, name: str, build: Build, cmd: List[str], |
| stamp_file: Optional[str]): |
| self.name = name |
| self.build = build |
| self.cmd = cmd |
| self.stamp_file = stamp_file |
| self._terminated = False |
| self._replaced = False |
| self._lock = threading.RLock() |
| self._proc: Optional[subprocess.Popen] = None |
| self._thread: Optional[threading.Thread] = None |
| self._delete_stamp_thread: Optional[threading.Thread] = None |
| self._return_code: Optional[int] = None |
| |
| @property |
| def key(self): |
| return (self.build.cwd, self.name) |
| |
| def __hash__(self): |
| return hash((self.key, self.build.id)) |
| |
| def __eq__(self, other): |
| return self.key == other.key and self.build is other.build |
| |
| def start(self, on_complete_callback: Callable[[], None]) -> int: |
| """Starts the task if it has not already been terminated. |
| |
| Returns the number of processes that have been started. This is called at |
| most once when the task is popped off the task queue.""" |
| with self._lock: |
| if self._terminated: |
| return 0 |
| |
| # Use os.nice(19) to ensure the lowest priority (idle) for these analysis |
| # tasks since we want to avoid slowing down the actual build. |
| # TODO(wnwen): Use ionice to reduce resource consumption. |
| self.build.add_process(self) |
| # This use of preexec_fn is sufficiently simple, just one os.nice call. |
| # pylint: disable=subprocess-popen-preexec-fn |
| self._proc = subprocess.Popen( |
| self.cmd, |
| stdout=subprocess.PIPE, |
| stderr=subprocess.STDOUT, |
| cwd=self.build.cwd, |
| env=self.build.env, |
| text=True, |
| preexec_fn=lambda: os.nice(19), |
| ) |
| self._thread = threading.Thread( |
| target=self._complete_when_process_finishes, |
| args=(on_complete_callback, )) |
| self._thread.start() |
| return 1 |
| |
| def terminate(self, replaced=False): |
| """Can be called multiple times to cancel and ignore the task's output.""" |
| with self._lock: |
| if self._terminated: |
| return |
| self._terminated = True |
| self._replaced = replaced |
| |
| # It is safe to access _proc and _thread outside of _lock since they are |
| # only changed by self.start holding _lock when self._terminate is false. |
| # Since we have just set self._terminate to true inside of _lock, we know |
| # that neither _proc nor _thread will be changed from this point onwards. |
| if self._proc: |
| self._proc.terminate() |
| self._proc.wait() |
| # Ensure that self._complete is called either by the thread or by us. |
| if self._thread: |
| self._thread.join() |
| else: |
| self._complete() |
| |
| def _complete_when_process_finishes(self, |
| on_complete_callback: Callable[[], None]): |
| assert self._proc |
| # We know Popen.communicate will return a str and not a byte since it is |
| # constructed with text=True. |
| stdout: str = self._proc.communicate()[0] |
| self._return_code = self._proc.returncode |
| self.build.process_complete() |
| self._complete(stdout) |
| on_complete_callback() |
| |
| def _complete(self, stdout: str = ''): |
| """Update the user and ninja after the task has run or been terminated. |
| |
| This method should only be run once per task. Avoid modifying the task so |
| that this method does not need locking.""" |
| |
| delete_stamp = False |
| status_string = 'FINISHED' |
| if self._terminated: |
| status_string = 'TERMINATED' |
| # When tasks are replaced, avoid deleting the stamp file, context: |
| # https://issuetracker.google.com/301961827. |
| if not self._replaced: |
| delete_stamp = True |
| elif stdout or self._return_code != 0: |
| status_string = 'FAILED' |
| delete_stamp = True |
| preamble = [ |
| f'FAILED: {self.name}', |
| f'Return code: {self._return_code}', |
| 'CMD: ' + shlex.join(self.cmd), |
| 'STDOUT:', |
| ] |
| |
| message = '\n'.join(preamble + [stdout]) |
| self.build.log(message) |
| server_log(message) |
| |
| if OptionsManager.should_remote_print(): |
| # Add emoji to show that output is from the build server. |
| preamble = [f'⏩ {line}' for line in preamble] |
| remote_message = '\n'.join(preamble + [stdout]) |
| # Add a new line at start of message to clearly delineate from previous |
| # output/text already on the remote tty we are printing to. |
| self.build.stdout.write(f'\n{remote_message}') |
| self.build.stdout.flush() |
| if delete_stamp and self.stamp_file: |
| # Force siso to consider failed targets as dirty. |
| try: |
| os.unlink(os.path.join(self.build.cwd, self.stamp_file)) |
| except FileNotFoundError: |
| pass |
| self.build.task_done(self, status_string) |
| |
| |
| def _handle_add_task(data, current_tasks: Dict[Tuple[str, str], Task]): |
| """Handle messages of type ADD_TASK.""" |
| build_id = data['build_id'] |
| build = BuildManager.get_build(build_id) |
| BuildManager.maybe_init_cwd(build, data.get('cwd')) |
| |
| cmd = data['cmd'] |
| name = data.get('name') or shlex.join(cmd) |
| new_task = Task(name=name, |
| cmd=cmd, |
| build=build, |
| stamp_file=data['stamp_file']) |
| existing_task = current_tasks.get(new_task.key) |
| if existing_task: |
| existing_task.terminate(replaced=True) |
| current_tasks[new_task.key] = new_task |
| |
| build.add_task(new_task) |
| |
| |
| def _handle_query_build(data, connection: socket.socket): |
| """Handle messages of type QUERY_BUILD.""" |
| build_id = data['build_id'] |
| response = TaskStats.query_build(build_id) |
| try: |
| with connection: |
| server_utils.SendMessage(connection, response) |
| except BrokenPipeError: |
| # We should not die because the client died. |
| pass |
| |
| |
| def _handle_heartbeat(connection: socket.socket): |
| """Handle messages of type POLL_HEARTBEAT.""" |
| try: |
| with connection: |
| server_utils.SendMessage(connection, { |
| 'status': 'OK', |
| 'pid': os.getpid(), |
| }) |
| except BrokenPipeError: |
| # We should not die because the client died. |
| pass |
| |
| |
| def _handle_register_builder(data): |
| """Handle messages of type REGISTER_BUILDER.""" |
| env = data['env'] |
| pid = int(data['builder_pid']) |
| cwd = data['cwd'] |
| |
| BuildManager.register_builder(env, pid, cwd) |
| |
| |
| def _handle_cancel_build(data): |
| """Handle messages of type CANCEL_BUILD.""" |
| build_id = data['build_id'] |
| TaskManager.cancel_build(build_id) |
| BuildManager.update_remote_titles('') |
| |
| |
| def _handle_stop_server(): |
| """Handle messages of type STOP_SERVER.""" |
| server_log('STOPPING SERVER...') |
| TaskManager.deactivate() |
| server_log('STOPPED') |
| sys.exit(0) |
| |
| |
| def _listen_for_request_data(sock: socket.socket): |
| """Helper to encapsulate getting a new message.""" |
| while True: |
| conn = sock.accept()[0] |
| message = server_utils.ReceiveMessage(conn) |
| if message: |
| yield message, conn |
| |
| |
| def _register_cleanup_signal_handlers(): |
| original_sigint_handler = signal.getsignal(signal.SIGINT) |
| original_sigterm_handler = signal.getsignal(signal.SIGTERM) |
| |
| def _cleanup(signum, frame): |
| server_log('STOPPING SERVER...') |
| # Gracefully shut down the task manager, terminating all queued tasks. |
| TaskManager.deactivate() |
| server_log('STOPPED') |
| if signum == signal.SIGINT: |
| if callable(original_sigint_handler): |
| original_sigint_handler(signum, frame) |
| else: |
| raise KeyboardInterrupt() |
| if signum == signal.SIGTERM: |
| # Sometimes sigterm handler is not a callable. |
| if callable(original_sigterm_handler): |
| original_sigterm_handler(signum, frame) |
| else: |
| sys.exit(1) |
| |
| signal.signal(signal.SIGINT, _cleanup) |
| signal.signal(signal.SIGTERM, _cleanup) |
| |
| |
| def _process_requests(sock: socket.socket, exit_on_idle: bool): |
| """Main loop for build server receiving request messages.""" |
| # Since dicts in python can contain anything, explicitly type tasks to help |
| # make static type checking more useful. |
| tasks: Dict[Tuple[str, str], Task] = {} |
| server_log(f'Server started. PID={os.getpid()}') |
| _register_cleanup_signal_handlers() |
| # pylint: disable=too-many-nested-blocks |
| while True: |
| try: |
| for data, connection in _listen_for_request_data(sock): |
| message_type = data.get('message_type', server_utils.ADD_TASK) |
| if message_type == server_utils.POLL_HEARTBEAT: |
| _handle_heartbeat(connection) |
| elif message_type == server_utils.ADD_TASK: |
| connection.close() |
| _handle_add_task(data, tasks) |
| elif message_type == server_utils.QUERY_BUILD: |
| _handle_query_build(data, connection) |
| elif message_type == server_utils.REGISTER_BUILDER: |
| connection.close() |
| _handle_register_builder(data) |
| elif message_type == server_utils.CANCEL_BUILD: |
| connection.close() |
| _handle_cancel_build(data) |
| elif message_type == server_utils.STOP_SERVER: |
| connection.close() |
| _handle_stop_server() |
| else: |
| connection.close() |
| except TimeoutError: |
| # If we have not received a new task in a while and do not have any |
| # pending tasks or running builds, then exit. Otherwise keep waiting. |
| if (TaskStats.num_pending_tasks() == 0 |
| and not BuildManager.has_active_builds() and exit_on_idle): |
| break |
| except KeyboardInterrupt: |
| break |
| BuildManager.update_remote_titles('') |
| |
| |
| def query_build_info(build_id=None): |
| """Communicates with the main server to query build info.""" |
| return _send_message_with_response({ |
| 'message_type': server_utils.QUERY_BUILD, |
| 'build_id': build_id, |
| }) |
| |
| |
| def _wait_for_build(build_id): |
| """Comunicates with the main server waiting for a build to complete.""" |
| start_time = datetime.datetime.now() |
| while True: |
| try: |
| build_info = query_build_info(build_id)['builds'][0] |
| except ConnectionRefusedError: |
| print('No server running. It likely finished all tasks.') |
| print('You can check $OUTDIR/buildserver.log.0 to be sure.') |
| return 0 |
| |
| pending_tasks = build_info['pending_tasks'] |
| |
| if pending_tasks == 0: |
| print(f'\nAll tasks completed for build_id: {build_id}.') |
| return 0 |
| |
| current_time = datetime.datetime.now() |
| duration = current_time - start_time |
| print(f'\rWaiting for {pending_tasks} tasks [{str(duration)}]\033[K', |
| end='', |
| flush=True) |
| time.sleep(1) |
| |
| |
| def _wait_for_idle(): |
| """Communicates with the main server waiting for all builds to complete.""" |
| start_time = datetime.datetime.now() |
| while True: |
| try: |
| builds = query_build_info()['builds'] |
| except ConnectionRefusedError: |
| print('No server running. It likely finished all tasks.') |
| print('You can check $OUTDIR/buildserver.log.0 to be sure.') |
| return 0 |
| |
| all_pending_tasks = 0 |
| all_completed_tasks = 0 |
| for build_info in builds: |
| pending_tasks = build_info['pending_tasks'] |
| completed_tasks = build_info['completed_tasks'] |
| active = build_info['is_active'] |
| # Ignore completed builds. |
| if active or pending_tasks: |
| all_pending_tasks += pending_tasks |
| all_completed_tasks += completed_tasks |
| total_tasks = all_pending_tasks + all_completed_tasks |
| |
| if all_pending_tasks == 0: |
| print('\nServer Idle, All tasks complete.') |
| return 0 |
| |
| current_time = datetime.datetime.now() |
| duration = current_time - start_time |
| print( |
| f'\rWaiting for {all_pending_tasks} remaining tasks. ' |
| f'({all_completed_tasks}/{total_tasks} tasks complete) ' |
| f'[{str(duration)}]\033[K', |
| end='', |
| flush=True) |
| time.sleep(0.5) |
| |
| |
| def _send_message_and_close(message_dict): |
| with contextlib.closing(socket.socket(socket.AF_UNIX)) as sock: |
| sock.connect(server_utils.SOCKET_ADDRESS) |
| sock.settimeout(1) |
| server_utils.SendMessage(sock, message_dict) |
| |
| |
| def _send_message_with_response(message_dict): |
| with contextlib.closing(socket.socket(socket.AF_UNIX)) as sock: |
| sock.connect(server_utils.SOCKET_ADDRESS) |
| sock.settimeout(1) |
| server_utils.SendMessage(sock, message_dict) |
| return server_utils.ReceiveMessage(sock) |
| |
| |
| def _send_cancel_build(build_id): |
| _send_message_and_close({ |
| 'message_type': server_utils.CANCEL_BUILD, |
| 'build_id': build_id, |
| }) |
| return 0 |
| |
| |
| def _send_stop_server(): |
| try: |
| _send_message_and_close({ |
| 'message_type': server_utils.STOP_SERVER, |
| }) |
| except socket.error as e: |
| if e.errno == 111: |
| sys.stderr.write('No running build server found.\n') |
| return 1 |
| raise |
| return 0 |
| |
| |
| def _register_build(builder_pid, output_directory): |
| if output_directory is not None: |
| output_directory = str(pathlib.Path(output_directory).absolute()) |
| for _attempt in range(3): |
| try: |
| # Ensure environment variables that the server expects to be there are |
| # present. |
| server_utils.AssertEnvironmentVariables() |
| |
| _send_message_and_close({ |
| 'message_type': server_utils.REGISTER_BUILDER, |
| 'env': dict(os.environ), |
| 'builder_pid': builder_pid, |
| 'cwd': output_directory, |
| }) |
| return 0 |
| except OSError: |
| time.sleep(0.05) |
| print('Failed to register build. No server running?') |
| return 1 |
| |
| |
| def poll_server(retries=3): |
| """Communicates with the main server to query build info.""" |
| for _attempt in range(retries): |
| try: |
| response = _send_message_with_response( |
| {'message_type': server_utils.POLL_HEARTBEAT}) |
| if response: |
| break |
| except OSError: |
| time.sleep(0.05) |
| else: |
| return None |
| return response['pid'] |
| |
| |
| def _print_build_status_all(): |
| try: |
| query_data = query_build_info(None) |
| except ConnectionRefusedError: |
| print('No server running. Consult $OUTDIR/buildserver.log.0') |
| return 0 |
| builds = query_data['builds'] |
| pid = query_data['pid'] |
| all_active_tasks = [] |
| print(f'Build server (PID={pid}) has {len(builds)} registered builds') |
| for build_info in builds: |
| build_id = build_info['build_id'] |
| pending_tasks = build_info['pending_tasks'] |
| completed_tasks = build_info['completed_tasks'] |
| active_tasks = build_info['active_tasks'] |
| out_dir = build_info['outdir'] |
| active = build_info['is_active'] |
| total_tasks = pending_tasks + completed_tasks |
| all_active_tasks += active_tasks |
| if total_tasks == 0 and not active: |
| status = 'Finished without any jobs' |
| else: |
| if active: |
| status = 'Main build is still running' |
| else: |
| status = 'Main build completed' |
| if out_dir: |
| status += f' in {out_dir}' |
| status += f'. Tasks completed: {completed_tasks}/{total_tasks}' |
| if completed_tasks < total_tasks: |
| status += f' {len(active_tasks)} task(s) currently executing' |
| print(f'{build_id}: {status}') |
| if all_active_tasks: |
| total = len(all_active_tasks) |
| to_show = min(4, total) |
| print(f'Currently executing (showing {to_show} of {total}):') |
| for cmd in sorted(all_active_tasks)[:to_show]: |
| truncated = shlex.join(cmd) |
| if len(truncated) > 200: |
| truncated = truncated[:200] + '...' |
| print(truncated) |
| return 0 |
| |
| |
| def _print_build_status(build_id): |
| server_path = os.path.relpath(str(server_utils.SERVER_SCRIPT)) |
| try: |
| builds = query_build_info(build_id)['builds'] |
| if not builds: |
| print(f'⚠️ No build found with id ({build_id})') |
| print('⚠️ To see the status of all builds:', |
| shlex.join([server_path, '--print-status-all'])) |
| return 1 |
| build_info = builds[0] |
| except ConnectionRefusedError: |
| print('⚠️ No server running. Consult $OUTDIR/buildserver.log.0') |
| return 0 |
| pending_tasks = build_info['pending_tasks'] |
| |
| # Print nothing unless there are still pending tasks |
| if pending_tasks: |
| is_str = 'is' if pending_tasks == 1 else 'are' |
| job_str = 'job' if pending_tasks == 1 else 'jobs' |
| print(f'⏩ There {is_str} still {pending_tasks} static analysis {job_str}' |
| ' running in the background.') |
| print('⏩ To wait for them:', shlex.join([server_path, '--wait-for-idle'])) |
| return 0 |
| |
| |
| def _start_server(exit_on_idle): |
| sys.excepthook = _exception_hook |
| with socket.socket(socket.AF_UNIX) as sock: |
| sock.settimeout(_SOCKET_TIMEOUT) |
| try: |
| sock.bind(server_utils.SOCKET_ADDRESS) |
| except OSError as e: |
| # errno 98 is Address already in use |
| if e.errno == 98: |
| if not OptionsManager.is_quiet(): |
| pid = poll_server() |
| print(f'Another instance is already running (pid: {pid}).', |
| file=sys.stderr) |
| return 1 |
| raise |
| sock.listen() |
| _process_requests(sock, exit_on_idle) |
| return 0 |
| |
| |
| def _add_task(cmd): |
| build_id = f'default-{time.time()}' |
| os.environ['AUTONINJA_BUILD_ID'] = build_id |
| tty = os.readlink('/proc/self/fd/1') |
| if os.path.exists(tty): |
| os.environ['AUTONINJA_STDOUT_NAME'] = tty |
| else: |
| os.environ['AUTONINJA_STDOUT_NAME'] = '/dev/null' |
| |
| if code := _register_build(os.getpid(), os.getcwd()): |
| return code |
| |
| try: |
| _send_message_and_close({ |
| 'name': None, |
| 'message_type': server_utils.ADD_TASK, |
| 'cmd': cmd, |
| 'cwd': os.getcwd(), |
| 'build_id': build_id, |
| 'stamp_file': None, |
| }) |
| return 0 |
| except socket.error as e: |
| if e.errno == 111: |
| sys.stderr.write('No running build server found.\n') |
| return 1 |
| raise |
| |
| |
| def _main_old(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument( |
| '--exit-on-idle', |
| action='store_true', |
| help='Server started on demand. Exit when all tasks run out.') |
| parser.add_argument('--quiet', |
| action='store_true', |
| help='Do not output status updates.') |
| parser.add_argument('--no-remote-print', |
| action='store_true', |
| help='Do not output errors to remote terminals.') |
| parser.add_argument('--wait-for-build', |
| metavar='BUILD_ID', |
| help='Wait for build server to finish with all tasks ' |
| 'for BUILD_ID and output any pending messages.') |
| parser.add_argument('--wait-for-idle', |
| action='store_true', |
| help='Wait for build server to finish with all ' |
| 'pending tasks.') |
| parser.add_argument('--print-status', |
| metavar='BUILD_ID', |
| help='Print the current state of a build.') |
| parser.add_argument('--print-status-all', |
| action='store_true', |
| help='Print the current state of all active builds.') |
| parser.add_argument( |
| '--register-build-id', |
| metavar='BUILD_ID', |
| help='Inform the build server that a new build has started.') |
| parser.add_argument('--output-directory', |
| help='Build directory (use with --register-build-id)') |
| parser.add_argument('--builder-pid', |
| help='Builder process\'s pid for build BUILD_ID.') |
| parser.add_argument('--cancel-build', |
| metavar='BUILD_ID', |
| help='Cancel all pending and running tasks for BUILD_ID.') |
| args = parser.parse_args() |
| |
| OptionsManager.set_options(quiet=args.quiet, |
| should_remote_print=not args.no_remote_print) |
| |
| if args.wait_for_build: |
| return _wait_for_build(args.wait_for_build) |
| if args.wait_for_idle: |
| return _wait_for_idle() |
| if args.print_status: |
| return _print_build_status(args.print_status) |
| if args.print_status_all: |
| return _print_build_status_all() |
| if args.register_build_id: |
| return _register_build(args.builder_pid, args.output_directory) |
| if args.cancel_build: |
| return _send_cancel_build(args.cancel_build) |
| return _start_server(args.exit_on_idle) |
| |
| |
| def _main_new(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| sub_parsers = parser.add_subparsers(dest='command') |
| |
| sub_parser = sub_parsers.add_parser('start', help='Start the server') |
| sub_parser.add_argument('--quiet', |
| action='store_true', |
| help='Do not output status updates.') |
| sub_parser.add_argument('--no-remote-print', |
| action='store_true', |
| help='Do not output errors to remote terminals.') |
| sub_parser.add_argument( |
| '--exit-on-idle', |
| action='store_true', |
| help='Server started on demand. Exit when all tasks run out.') |
| |
| sub_parser = sub_parsers.add_parser('stop', |
| help='Stops the server if it is running') |
| |
| sub_parser = sub_parsers.add_parser( |
| 'register-build', help='Tell a running server about a new ninja session') |
| sub_parser.add_argument('--output-directory', |
| required=True, |
| help='CWD for the build') |
| sub_parser.add_argument('--builder-pid', |
| required=True, |
| help='Builder process\'s PID.') |
| |
| sub_parser = sub_parsers.add_parser( |
| 'unregister-build', |
| help='Tell a running server a ninja session has finished') |
| sub_parser.add_argument('--build-id', |
| required=True, |
| help='The AUTONINJA_BUILD_ID') |
| sub_parser.add_argument('--verbose', |
| action='store_true', |
| help='Print status if jobs exist.') |
| sub_parser.add_argument('--cancel-jobs', |
| action='store_true', |
| help='Cancel pending jobs') |
| |
| sub_parser = sub_parsers.add_parser('status', help='Print status and exit') |
| sub_parser.add_argument('--build-id', |
| help='The AUTONINJA_BUILD_ID of the session to query ' |
| '(otherwise prints all sessions).') |
| |
| sub_parser = sub_parsers.add_parser('wait', help='Wait for jobs to complete') |
| sub_parser.add_argument( |
| '--build-id', |
| help='The AUTONINJA_BUILD_ID of the session to wait for ' |
| '(otherwise waits for all sessions).') |
| |
| sub_parser = sub_parsers.add_parser('run', help='Adds a task.') |
| sub_parser.add_argument('cmd', nargs='+', help='The command to run') |
| |
| args = parser.parse_args() |
| |
| ret = 0 |
| if args.command == 'start': |
| OptionsManager.set_options(quiet=args.quiet, |
| should_remote_print=not args.no_remote_print) |
| ret = _start_server(args.exit_on_idle) |
| elif args.command == 'stop': |
| ret = _send_stop_server() |
| elif args.command == 'register-build': |
| ret = _register_build(args.builder_pid, args.output_directory) |
| elif args.command == 'unregister-build': |
| if args.verbose: |
| ret = _print_build_status(args.build_id) |
| if args.cancel_jobs: |
| ret = _send_cancel_build(args.cancel_build) |
| elif args.command == 'status': |
| if args.build_id: |
| ret = _print_build_status(args.build_id) |
| else: |
| ret = _print_build_status_all() |
| elif args.command == 'wait': |
| if args.build_id: |
| ret = _wait_for_build(args.build_id) |
| else: |
| ret = _wait_for_idle() |
| elif args.command == 'run': |
| ret = _add_task(args.cmd) |
| else: |
| parser.print_help() |
| return 1 |
| return ret |
| |
| |
| def main(): |
| if len(sys.argv) <= 1 or not sys.argv[1].startswith('-'): |
| return _main_new() |
| return _main_old() |
| |
| |
| if __name__ == '__main__': |
| sys.exit(main()) |
