Source code
Revision control
Copy as Markdown
Other Tools
#!/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())