1# ===----------------------------------------------------------------------===## 2# 3# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4# See https://llvm.org/LICENSE.txt for license information. 5# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6# 7# ===----------------------------------------------------------------------===## 8 9import re 10import select 11import socket 12import subprocess 13import tempfile 14import threading 15from typing import List 16 17 18def _get_cpu_count() -> int: 19 # Determine the number of cores by listing a /sys directory. Older devices 20 # lack `nproc`. Even if a static toybox binary is pushed to the device, it may 21 # return an incorrect value. (e.g. On a Nexus 7 running Android 5.0, toybox 22 # nproc returns 1 even though the device has 4 CPUs.) 23 job = subprocess.run(["adb", "shell", "ls /sys/devices/system/cpu"], 24 encoding="utf8", check=False, 25 stdout=subprocess.PIPE, stderr=subprocess.PIPE) 26 if job.returncode == 1: 27 # Maybe adb is missing, maybe ANDROID_SERIAL needs to be defined, maybe the 28 # /sys subdir isn't there. Most errors will be handled later, just use one 29 # job. (N.B. The adb command still succeeds even if ls fails on older 30 # devices that lack the shell_v2 adb feature.) 31 return 1 32 # Make sure there are no CR characters in the output. Pre-shell_v2, the adb 33 # stdout comes from a master pty so newlines are CRLF-delimited. On Windows, 34 # LF might also get expanded to CRLF. 35 cpu_listing = job.stdout.replace('\r', '\n') 36 37 # Count lines that match "cpu${DIGITS}". 38 result = len([line for line in cpu_listing.splitlines() 39 if re.match(r'cpu(\d)+$', line)]) 40 41 # Restrict the result to something reasonable. 42 if result < 1: 43 result = 1 44 if result > 1024: 45 result = 1024 46 47 return result 48 49 50def _job_limit_socket_thread(temp_dir: tempfile.TemporaryDirectory, 51 server: socket.socket, job_count: int) -> None: 52 """Service the job limit server socket, accepting only as many connections 53 as there should be concurrent jobs. 54 """ 55 clients: List[socket.socket] = [] 56 while True: 57 rlist = list(clients) 58 if len(clients) < job_count: 59 rlist.append(server) 60 rlist, _, _ = select.select(rlist, [], []) 61 for sock in rlist: 62 if sock == server: 63 new_client, _ = server.accept() 64 new_client.send(b"x") 65 clients.append(new_client) 66 else: 67 sock.close() 68 clients.remove(sock) 69 70 71def adb_job_limit_socket() -> str: 72 """An Android device can frequently have many fewer cores than the host 73 (e.g. 4 versus 128). We want to exploit all the device cores without 74 overburdening it. 75 76 Create a Unix domain socket that only allows as many connections as CPUs on 77 the Android device. 78 """ 79 80 # Create the job limit server socket. 81 temp_dir = tempfile.TemporaryDirectory(prefix="libcxx_") 82 sock_addr = temp_dir.name + "/adb_job.sock" 83 server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) 84 server.bind(sock_addr) 85 server.listen(1) 86 87 # Spawn a thread to service the socket. As a daemon thread, its existence 88 # won't prevent interpreter shutdown. The temp dir will still be removed on 89 # shutdown. 90 cpu_count = _get_cpu_count() 91 threading.Thread(target=_job_limit_socket_thread, 92 args=(temp_dir, server, cpu_count), 93 daemon=True).start() 94 95 return sock_addr 96