fast_local_dev_server.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336
  1. #!/usr/bin/env python3
  2. # Copyright 2021 The Chromium Authors. All rights reserved.
  3. # Use of this source code is governed by a BSD-style license that can be
  4. # found in the LICENSE file.
  5. """Creates an server to offload non-critical-path GN targets."""
  6. from __future__ import annotations
  7. import argparse
  8. import json
  9. import os
  10. import queue
  11. import shutil
  12. import socket
  13. import subprocess
  14. import sys
  15. import threading
  16. from typing import Callable, Dict, List, Optional, Tuple
  17. sys.path.append(os.path.join(os.path.dirname(__file__), 'gyp'))
  18. from util import server_utils
  19. def log(msg: str, *, end: str = ''):
  20. # Shrink the message (leaving a 2-char prefix and use the rest of the room
  21. # for the suffix) according to terminal size so it is always one line.
  22. width = shutil.get_terminal_size().columns
  23. prefix = f'[{TaskStats.prefix()}] '
  24. max_msg_width = width - len(prefix)
  25. if len(msg) > max_msg_width:
  26. length_to_show = max_msg_width - 5 # Account for ellipsis and header.
  27. msg = f'{msg[:2]}...{msg[-length_to_show:]}'
  28. # \r to return the carriage to the beginning of line.
  29. # \033[K to replace the normal \n to erase until the end of the line.
  30. # Avoid the default line ending so the next \r overwrites the same line just
  31. # like ninja's output.
  32. print(f'\r{prefix}{msg}\033[K', end=end, flush=True)
  33. class TaskStats:
  34. """Class to keep track of aggregate stats for all tasks across threads."""
  35. _num_processes = 0
  36. _completed_tasks = 0
  37. _total_tasks = 0
  38. _lock = threading.Lock()
  39. @classmethod
  40. def no_running_processes(cls):
  41. return cls._num_processes == 0
  42. @classmethod
  43. def add_task(cls):
  44. # Only the main thread calls this, so there is no need for locking.
  45. cls._total_tasks += 1
  46. @classmethod
  47. def add_process(cls):
  48. with cls._lock:
  49. cls._num_processes += 1
  50. @classmethod
  51. def remove_process(cls):
  52. with cls._lock:
  53. cls._num_processes -= 1
  54. @classmethod
  55. def complete_task(cls):
  56. with cls._lock:
  57. cls._completed_tasks += 1
  58. @classmethod
  59. def prefix(cls):
  60. # Ninja's prefix is: [205 processes, 6/734 @ 6.5/s : 0.922s ]
  61. # Time taken and task completion rate are not important for the build server
  62. # since it is always running in the background and uses idle priority for
  63. # its tasks.
  64. with cls._lock:
  65. word = 'process' if cls._num_processes == 1 else 'processes'
  66. return (f'{cls._num_processes} {word}, '
  67. f'{cls._completed_tasks}/{cls._total_tasks}')
  68. class TaskManager:
  69. """Class to encapsulate a threadsafe queue and handle deactivating it."""
  70. def __init__(self):
  71. self._queue: queue.SimpleQueue[Task] = queue.SimpleQueue()
  72. self._deactivated = False
  73. def add_task(self, task: Task):
  74. assert not self._deactivated
  75. TaskStats.add_task()
  76. self._queue.put(task)
  77. log(f'QUEUED {task.name}')
  78. self._maybe_start_tasks()
  79. def deactivate(self):
  80. self._deactivated = True
  81. while not self._queue.empty():
  82. try:
  83. task = self._queue.get_nowait()
  84. except queue.Empty:
  85. return
  86. task.terminate()
  87. @staticmethod
  88. def _num_running_processes():
  89. with open('/proc/stat') as f:
  90. for line in f:
  91. if line.startswith('procs_running'):
  92. return int(line.rstrip().split()[1])
  93. assert False, 'Could not read /proc/stat'
  94. return 0
  95. def _maybe_start_tasks(self):
  96. if self._deactivated:
  97. return
  98. # Include load avg so that a small dip in the number of currently running
  99. # processes will not cause new tasks to be started while the overall load is
  100. # heavy.
  101. cur_load = max(self._num_running_processes(), os.getloadavg()[0])
  102. num_started = 0
  103. # Always start a task if we don't have any running, so that all tasks are
  104. # eventually finished. Try starting up tasks when the overall load is light.
  105. # Limit to at most 2 new tasks to prevent ramping up too fast. There is a
  106. # chance where multiple threads call _maybe_start_tasks and each gets to
  107. # spawn up to 2 new tasks, but since the only downside is some build tasks
  108. # get worked on earlier rather than later, it is not worth mitigating.
  109. while num_started < 2 and (TaskStats.no_running_processes()
  110. or num_started + cur_load < os.cpu_count()):
  111. try:
  112. next_task = self._queue.get_nowait()
  113. except queue.Empty:
  114. return
  115. num_started += next_task.start(self._maybe_start_tasks)
  116. # TODO(wnwen): Break this into Request (encapsulating what ninja sends) and Task
  117. # when a Request starts to be run. This would eliminate ambiguity
  118. # about when and whether _proc/_thread are initialized.
  119. class Task:
  120. """Class to represent one task and operations on it."""
  121. def __init__(self, name: str, cwd: str, cmd: List[str], stamp_file: str):
  122. self.name = name
  123. self.cwd = cwd
  124. self.cmd = cmd
  125. self.stamp_file = stamp_file
  126. self._terminated = False
  127. self._lock = threading.Lock()
  128. self._proc: Optional[subprocess.Popen] = None
  129. self._thread: Optional[threading.Thread] = None
  130. self._return_code: Optional[int] = None
  131. @property
  132. def key(self):
  133. return (self.cwd, self.name)
  134. def start(self, on_complete_callback: Callable[[], None]) -> int:
  135. """Starts the task if it has not already been terminated.
  136. Returns the number of processes that have been started. This is called at
  137. most once when the task is popped off the task queue."""
  138. # The environment variable forces the script to actually run in order to
  139. # avoid infinite recursion.
  140. env = os.environ.copy()
  141. env[server_utils.BUILD_SERVER_ENV_VARIABLE] = '1'
  142. with self._lock:
  143. if self._terminated:
  144. return 0
  145. # Use os.nice(19) to ensure the lowest priority (idle) for these analysis
  146. # tasks since we want to avoid slowing down the actual build.
  147. # TODO(wnwen): Use ionice to reduce resource consumption.
  148. TaskStats.add_process()
  149. log(f'STARTING {self.name}')
  150. # This use of preexec_fn is sufficiently simple, just one os.nice call.
  151. # pylint: disable=subprocess-popen-preexec-fn
  152. self._proc = subprocess.Popen(
  153. self.cmd,
  154. stdout=subprocess.PIPE,
  155. stderr=subprocess.STDOUT,
  156. cwd=self.cwd,
  157. env=env,
  158. text=True,
  159. preexec_fn=lambda: os.nice(19),
  160. )
  161. self._thread = threading.Thread(
  162. target=self._complete_when_process_finishes,
  163. args=(on_complete_callback, ))
  164. self._thread.start()
  165. return 1
  166. def terminate(self):
  167. """Can be called multiple times to cancel and ignore the task's output."""
  168. with self._lock:
  169. if self._terminated:
  170. return
  171. self._terminated = True
  172. # It is safe to access _proc and _thread outside of _lock since they are
  173. # only changed by self.start holding _lock when self._terminate is false.
  174. # Since we have just set self._terminate to true inside of _lock, we know
  175. # that neither _proc nor _thread will be changed from this point onwards.
  176. if self._proc:
  177. self._proc.terminate()
  178. self._proc.wait()
  179. # Ensure that self._complete is called either by the thread or by us.
  180. if self._thread:
  181. self._thread.join()
  182. else:
  183. self._complete()
  184. def _complete_when_process_finishes(self,
  185. on_complete_callback: Callable[[], None]):
  186. assert self._proc
  187. # We know Popen.communicate will return a str and not a byte since it is
  188. # constructed with text=True.
  189. stdout: str = self._proc.communicate()[0]
  190. self._return_code = self._proc.returncode
  191. TaskStats.remove_process()
  192. self._complete(stdout)
  193. on_complete_callback()
  194. def _complete(self, stdout: str = ''):
  195. """Update the user and ninja after the task has run or been terminated.
  196. This method should only be run once per task. Avoid modifying the task so
  197. that this method does not need locking."""
  198. TaskStats.complete_task()
  199. failed = False
  200. if self._terminated:
  201. log(f'TERMINATED {self.name}')
  202. # Ignore stdout as it is now outdated.
  203. failed = True
  204. else:
  205. log(f'FINISHED {self.name}')
  206. if stdout or self._return_code != 0:
  207. failed = True
  208. # An extra new line is needed since we want to preserve the previous
  209. # _log line. Use a single print so that it is threadsafe.
  210. # TODO(wnwen): Improve stdout display by parsing over it and moving the
  211. # actual error to the bottom. Otherwise long command lines
  212. # in the Traceback section obscure the actual error(s).
  213. print('\n' + '\n'.join([
  214. f'FAILED: {self.name}',
  215. f'Return code: {self._return_code}',
  216. ' '.join(self.cmd),
  217. stdout,
  218. ]))
  219. if failed:
  220. # Force ninja to consider failed targets as dirty.
  221. try:
  222. os.unlink(os.path.join(self.cwd, self.stamp_file))
  223. except FileNotFoundError:
  224. pass
  225. else:
  226. # Ninja will rebuild targets when their inputs change even if their stamp
  227. # file has a later modified time. Thus we do not need to worry about the
  228. # script being run by the build server updating the mtime incorrectly.
  229. pass
  230. def _listen_for_request_data(sock: socket.socket):
  231. while True:
  232. conn = sock.accept()[0]
  233. received = []
  234. with conn:
  235. while True:
  236. data = conn.recv(4096)
  237. if not data:
  238. break
  239. received.append(data)
  240. if received:
  241. yield json.loads(b''.join(received))
  242. def _process_requests(sock: socket.socket):
  243. # Since dicts in python can contain anything, explicitly type tasks to help
  244. # make static type checking more useful.
  245. tasks: Dict[Tuple[str, str], Task] = {}
  246. task_manager = TaskManager()
  247. try:
  248. log('READY... Remember to set android_static_analysis="build_server" in '
  249. 'args.gn files')
  250. for data in _listen_for_request_data(sock):
  251. task = Task(name=data['name'],
  252. cwd=data['cwd'],
  253. cmd=data['cmd'],
  254. stamp_file=data['stamp_file'])
  255. existing_task = tasks.get(task.key)
  256. if existing_task:
  257. existing_task.terminate()
  258. tasks[task.key] = task
  259. task_manager.add_task(task)
  260. except KeyboardInterrupt:
  261. log('STOPPING SERVER...', end='\n')
  262. # Gracefully shut down the task manager, terminating all queued tasks.
  263. task_manager.deactivate()
  264. # Terminate all currently running tasks.
  265. for task in tasks.values():
  266. task.terminate()
  267. log('STOPPED', end='\n')
  268. def main():
  269. parser = argparse.ArgumentParser(description=__doc__)
  270. parser.add_argument(
  271. '--fail-if-not-running',
  272. action='store_true',
  273. help='Used by GN to fail fast if the build server is not running.')
  274. args = parser.parse_args()
  275. if args.fail_if_not_running:
  276. with socket.socket(socket.AF_UNIX) as sock:
  277. try:
  278. sock.connect(server_utils.SOCKET_ADDRESS)
  279. except socket.error:
  280. print('Build server is not running and '
  281. 'android_static_analysis="build_server" is set.\nPlease run '
  282. 'this command in a separate terminal:\n\n'
  283. '$ build/android/fast_local_dev_server.py\n')
  284. return 1
  285. else:
  286. return 0
  287. with socket.socket(socket.AF_UNIX) as sock:
  288. sock.bind(server_utils.SOCKET_ADDRESS)
  289. sock.listen()
  290. _process_requests(sock)
  291. return 0
  292. if __name__ == '__main__':
  293. sys.exit(main())