| #!/usr/bin/env python3 |
| # Copyright 2021 The Chromium Authors. All rights reserved. |
| # 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 json |
| import os |
| import queue |
| import shutil |
| import socket |
| import subprocess |
| import sys |
| import threading |
| from typing import Callable, Dict, List, Optional, Tuple |
| |
| sys.path.append(os.path.join(os.path.dirname(__file__), 'gyp')) |
| from util import server_utils |
| |
| |
| def log(msg: str, *, end: str = ''): |
| # 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 |
| prefix = f'[{TaskStats.prefix()}] ' |
| 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=end, flush=True) |
| |
| |
| 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.Lock() |
| |
| @classmethod |
| def no_running_processes(cls): |
| return cls._num_processes == 0 |
| |
| @classmethod |
| def add_task(cls): |
| # Only the main thread calls this, so there is no need for locking. |
| 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 prefix(cls): |
| # 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: |
| word = 'process' if cls._num_processes == 1 else 'processes' |
| return (f'{cls._num_processes} {word}, ' |
| f'{cls._completed_tasks}/{cls._total_tasks}') |
| |
| |
| class TaskManager: |
| """Class to encapsulate a threadsafe queue and handle deactivating it.""" |
| |
| def __init__(self): |
| self._queue: queue.SimpleQueue[Task] = queue.SimpleQueue() |
| self._deactivated = False |
| |
| def add_task(self, task: Task): |
| assert not self._deactivated |
| TaskStats.add_task() |
| self._queue.put(task) |
| log(f'QUEUED {task.name}') |
| self._maybe_start_tasks() |
| |
| def deactivate(self): |
| self._deactivated = True |
| while not self._queue.empty(): |
| try: |
| task = self._queue.get_nowait() |
| except queue.Empty: |
| return |
| task.terminate() |
| |
| @staticmethod |
| 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' |
| |
| def _maybe_start_tasks(self): |
| if self._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(self._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()): |
| try: |
| next_task = self._queue.get_nowait() |
| except queue.Empty: |
| return |
| num_started += next_task.start(self._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, cwd: str, cmd: List[str], stamp_file: str): |
| self.name = name |
| self.cwd = cwd |
| self.cmd = cmd |
| self.stamp_file = stamp_file |
| self._terminated = False |
| self._lock = threading.Lock() |
| self._proc: Optional[subprocess.Popen] = None |
| self._thread: Optional[threading.Thread] = None |
| self._return_code: Optional[int] = None |
| |
| @property |
| def key(self): |
| return (self.cwd, self.name) |
| |
| 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.""" |
| |
| # The environment variable forces the script to actually run in order to |
| # avoid infinite recursion. |
| env = os.environ.copy() |
| env[server_utils.BUILD_SERVER_ENV_VARIABLE] = '1' |
| |
| 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. |
| TaskStats.add_process() |
| log(f'STARTING {self.name}') |
| self._proc = subprocess.Popen( |
| self.cmd, |
| stdout=subprocess.PIPE, |
| stderr=subprocess.STDOUT, |
| cwd=self.cwd, |
| env=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): |
| """Can be called multiple times to cancel and ignore the task's output.""" |
| |
| with self._lock: |
| if self._terminated: |
| return |
| self._terminated = True |
| # 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 |
| TaskStats.remove_process() |
| 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.""" |
| |
| TaskStats.complete_task() |
| failed = False |
| if self._terminated: |
| log(f'TERMINATED {self.name}') |
| # Ignore stdout as it is now outdated. |
| failed = True |
| else: |
| log(f'FINISHED {self.name}') |
| if stdout or self._return_code != 0: |
| failed = True |
| # An extra new line is needed since we want to preserve the previous |
| # _log line. Use a single print so that it is threadsafe. |
| # TODO(wnwen): Improve stdout display by parsing over it and moving the |
| # actual error to the bottom. Otherwise long command lines |
| # in the Traceback section obscure the actual error(s). |
| print('\n' + '\n'.join([ |
| f'FAILED: {self.name}', |
| f'Return code: {self._return_code}', |
| ' '.join(self.cmd), |
| stdout, |
| ])) |
| |
| if failed: |
| # Force ninja to consider failed targets as dirty. |
| try: |
| os.unlink(os.path.join(self.cwd, self.stamp_file)) |
| except FileNotFoundError: |
| pass |
| else: |
| # Ninja will rebuild targets when their inputs change even if their stamp |
| # file has a later modified time. Thus we do not need to worry about the |
| # script being run by the build server updating the mtime incorrectly. |
| pass |
| |
| |
| def _listen_for_request_data(sock: socket.socket): |
| while True: |
| conn = sock.accept()[0] |
| received = [] |
| with conn: |
| while True: |
| data = conn.recv(4096) |
| if not data: |
| break |
| received.append(data) |
| if received: |
| yield json.loads(b''.join(received)) |
| |
| |
| def _process_requests(sock: socket.socket): |
| # Since dicts in python can contain anything, explicitly type tasks to help |
| # make static type checking more useful. |
| tasks: Dict[Tuple[str, str], Task] = {} |
| task_manager = TaskManager() |
| try: |
| for data in _listen_for_request_data(sock): |
| task = Task(name=data['name'], |
| cwd=data['cwd'], |
| cmd=data['cmd'], |
| stamp_file=data['stamp_file']) |
| existing_task = tasks.get(task.key) |
| if existing_task: |
| existing_task.terminate() |
| tasks[task.key] = task |
| task_manager.add_task(task) |
| except KeyboardInterrupt: |
| log('STOPPING SERVER...', end='\n') |
| # Gracefully shut down the task manager, terminating all queued tasks. |
| task_manager.deactivate() |
| # Terminate all currently running tasks. |
| for task in tasks.values(): |
| task.terminate() |
| log('STOPPED', end='\n') |
| |
| |
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.parse_args() |
| with socket.socket(socket.AF_UNIX) as sock: |
| sock.bind(server_utils.SOCKET_ADDRESS) |
| sock.listen() |
| _process_requests(sock) |
| |
| |
| if __name__ == '__main__': |
| sys.exit(main()) |
