# Copyright 2022 Google LLC. # SPDX-License-Identifier: Apache-2.0 """Module for job counters, limiting the amount of concurrent executions.""" import fcntl import functools import logging import multiprocessing import os import re import select import selectors import subprocess import sys logger = logging.getLogger('twister') logger.setLevel(logging.DEBUG) class JobHandle: """Small object to handle claim of a job.""" def __init__(self, release_func, *args, **kwargs): self.release_func = release_func self.args = args self.kwargs = kwargs def __enter__(self): pass def __exit__(self, exc_type, exc_value, traceback): if self.release_func: self.release_func(*self.args, **self.kwargs) class JobClient: """Abstract base class for all job clients.""" def get_job(self): """Claim a job.""" return JobHandle(None) @staticmethod def env(): """Get the environment variables necessary to share the job server.""" return {} @staticmethod def pass_fds(): """Returns the file descriptors that should be passed to subprocesses.""" return [] def popen(self, argv, **kwargs): """Start a process using subprocess.Popen All other arguments are passed to subprocess.Popen. Returns: A Popen object. """ kwargs.setdefault("env", os.environ) kwargs.setdefault("pass_fds", []) kwargs["env"].update(self.env()) kwargs["pass_fds"] += self.pass_fds() return subprocess.Popen( # pylint:disable=consider-using-with argv, **kwargs ) class GNUMakeJobClient(JobClient): """A job client for GNU make. A client of jobserver is allowed to run 1 job without contacting the jobserver, so maintain an optional self._internal_pipe to hold that job. """ def __init__(self, inheritable_pipe, jobs, internal_jobs=0, makeflags=None): self._makeflags = makeflags self._inheritable_pipe = inheritable_pipe self.jobs = jobs self._selector = selectors.DefaultSelector() if internal_jobs: self._internal_pipe = os.pipe() os.write(self._internal_pipe[1], b"+" * internal_jobs) os.set_blocking(self._internal_pipe[0], False) self._selector.register( self._internal_pipe[0], selectors.EVENT_READ, self._internal_pipe[1], ) else: self._internal_pipe = None if self._inheritable_pipe is not None: os.set_blocking(self._inheritable_pipe[0], False) self._selector.register( self._inheritable_pipe[0], selectors.EVENT_READ, self._inheritable_pipe[1], ) def __del__(self): if self._inheritable_pipe: os.close(self._inheritable_pipe[0]) os.close(self._inheritable_pipe[1]) if self._internal_pipe: os.close(self._internal_pipe[0]) os.close(self._internal_pipe[1]) @classmethod def from_environ(cls, env=None, jobs=0): """Create a job client from an environment with the MAKEFLAGS variable. If we are started under a GNU Make Job Server, we can search the environment for a string "--jobserver-auth=R,W", where R and W will be the read and write file descriptors to the pipe respectively. If we don't find this environment variable (or the string inside of it), this will raise an OSError. The specification for MAKEFLAGS is: * If the first char is "n", this is a dry run, just exit. * If the flags contains -j1, go to sequential mode. * If the flags contains --jobserver-auth=R,W AND those file descriptors are valid, use the jobserver. Otherwise output a warning. Args: env: Optionally, the environment to search. jobs: The number of jobs set by the user on the command line. Returns: A GNUMakeJobClient configured appropriately or None if there is no MAKEFLAGS environment variable. """ if env is None: env = os.environ makeflags = env.get("MAKEFLAGS") if not makeflags: return None match = re.search(r"--jobserver-auth=(\d+),(\d+)", makeflags) if match: pipe = [int(x) for x in match.groups()] if jobs: pipe = None logger.warning( "-jN forced on command line; ignoring GNU make jobserver" ) else: try: # Use F_GETFL to see if file descriptors are valid if pipe: rc = fcntl.fcntl(pipe[0], fcntl.F_GETFL) if not rc & os.O_ACCMODE == os.O_RDONLY: logger.warning( "FD %s is not readable (flags=%x); " "ignoring GNU make jobserver", pipe[0], rc) pipe = None if pipe: rc = fcntl.fcntl(pipe[1], fcntl.F_GETFL) if not rc & os.O_ACCMODE == os.O_WRONLY: logger.warning( "FD %s is not writable (flags=%x); " "ignoring GNU make jobserver", pipe[1], rc) pipe = None if pipe: logger.info("using GNU make jobserver") except OSError: pipe = None logger.warning( "No file descriptors; ignoring GNU make jobserver" ) else: pipe = None if not jobs: match = re.search(r"-j(\d+)", makeflags) if match: jobs = int(match.group(1)) if jobs == 1: logger.info("Running in sequential mode (-j1)") if makeflags[0] == "n": logger.info("MAKEFLAGS contained dry-run flag") sys.exit(0) return cls(pipe, jobs, internal_jobs=1, makeflags=makeflags) def get_job(self): """Claim a job. Returns: A JobHandle object. """ while True: ready_items = self._selector.select() if len(ready_items) > 0: read_fd = ready_items[0][0].fd write_fd = ready_items[0][0].data try: byte = os.read(read_fd, 1) return JobHandle( functools.partial(os.write, write_fd, byte) ) except BlockingIOError: pass def env(self): """Get the environment variables necessary to share the job server.""" if self._makeflags: return {"MAKEFLAGS": self._makeflags} flag = "" if self.jobs: flag += f" -j{self.jobs}" if self.jobs != 1 and self._inheritable_pipe is not None: flag += f" --jobserver-auth={self._inheritable_pipe[0]},{self._inheritable_pipe[1]}" return {"MAKEFLAGS": flag} def pass_fds(self): """Returns the file descriptors that should be passed to subprocesses.""" if self.jobs != 1 and self._inheritable_pipe is not None: return self._inheritable_pipe return [] class GNUMakeJobServer(GNUMakeJobClient): """Implements a GNU Make POSIX Job Server. See https://www.gnu.org/software/make/manual/html_node/POSIX-Jobserver.html for specification. """ def __init__(self, jobs=0): if not jobs: jobs = multiprocessing.cpu_count() elif jobs > select.PIPE_BUF: jobs = select.PIPE_BUF super().__init__(os.pipe(), jobs) os.write(self._inheritable_pipe[1], b"+" * jobs)