| import atexit | |
| import errno | |
| import os | |
| import selectors | |
| import signal | |
| import socket | |
| import struct | |
| import sys | |
| import threading | |
| import warnings | |
| from . import AuthenticationError | |
| from . import connection | |
| from . import process | |
| from .context import reduction | |
| from . import resource_tracker | |
| from . import spawn | |
| from . import util | |
| __all__ = ['ensure_running', 'get_inherited_fds', 'connect_to_new_process', | |
| 'set_forkserver_preload'] | |
| # | |
| # | |
| # | |
| MAXFDS_TO_SEND = 256 | |
| SIGNED_STRUCT = struct.Struct('q') # large enough for pid_t | |
| _AUTHKEY_LEN = 32 # <= PIPEBUF so it fits a single write to an empty pipe. | |
| # | |
| # Forkserver class | |
| # | |
| class ForkServer(object): | |
| def __init__(self): | |
| self._forkserver_authkey = None | |
| self._forkserver_address = None | |
| self._forkserver_alive_fd = None | |
| self._forkserver_pid = None | |
| self._inherited_fds = None | |
| self._lock = threading.Lock() | |
| self._preload_modules = ['__main__'] | |
| self._preload_on_error = 'ignore' | |
| def _stop(self): | |
| # Method used by unit tests to stop the server | |
| with self._lock: | |
| self._stop_unlocked() | |
| def _stop_unlocked(self): | |
| if self._forkserver_pid is None: | |
| return | |
| # close the "alive" file descriptor asks the server to stop | |
| os.close(self._forkserver_alive_fd) | |
| self._forkserver_alive_fd = None | |
| os.waitpid(self._forkserver_pid, 0) | |
| self._forkserver_pid = None | |
| if not util.is_abstract_socket_namespace(self._forkserver_address): | |
| os.unlink(self._forkserver_address) | |
| self._forkserver_address = None | |
| self._forkserver_authkey = None | |
| def set_forkserver_preload(self, modules_names, *, on_error='ignore'): | |
| '''Set list of module names to try to load in forkserver process. | |
| The on_error parameter controls how import failures are handled: | |
| "ignore" (default) silently ignores failures, "warn" emits warnings, | |
| and "fail" raises exceptions breaking the forkserver context. | |
| ''' | |
| if not all(type(mod) is str for mod in modules_names): | |
| raise TypeError('module_names must be a list of strings') | |
| if on_error not in ('ignore', 'warn', 'fail'): | |
| raise ValueError( | |
| f"on_error must be 'ignore', 'warn', or 'fail', " | |
| f"not {on_error!r}" | |
| ) | |
| self._preload_modules = modules_names | |
| self._preload_on_error = on_error | |
| def get_inherited_fds(self): | |
| '''Return list of fds inherited from parent process. | |
| This returns None if the current process was not started by fork | |
| server. | |
| ''' | |
| return self._inherited_fds | |
| def connect_to_new_process(self, fds): | |
| '''Request forkserver to create a child process. | |
| Returns a pair of fds (status_r, data_w). The calling process can read | |
| the child process's pid and (eventually) its returncode from status_r. | |
| The calling process should write to data_w the pickled preparation and | |
| process data. | |
| ''' | |
| self.ensure_running() | |
| assert self._forkserver_authkey | |
| if len(fds) + 4 >= MAXFDS_TO_SEND: | |
| raise ValueError('too many fds') | |
| with socket.socket(socket.AF_UNIX) as client: | |
| client.connect(self._forkserver_address) | |
| parent_r, child_w = os.pipe() | |
| child_r, parent_w = os.pipe() | |
| allfds = [child_r, child_w, self._forkserver_alive_fd, | |
| resource_tracker.getfd()] | |
| allfds += fds | |
| try: | |
| client.setblocking(True) | |
| wrapped_client = connection.Connection(client.fileno()) | |
| # The other side of this exchange happens in the child as | |
| # implemented in main(). | |
| try: | |
| connection.answer_challenge( | |
| wrapped_client, self._forkserver_authkey) | |
| connection.deliver_challenge( | |
| wrapped_client, self._forkserver_authkey) | |
| except (EOFError, ConnectionError, BrokenPipeError) as exc: | |
| if (self._preload_modules and | |
| self._preload_on_error == 'fail'): | |
| exc.add_note( | |
| "Forkserver process may have crashed during module " | |
| "preloading. Check stderr." | |
| ) | |
| raise | |
| finally: | |
| wrapped_client._detach() | |
| del wrapped_client | |
| reduction.sendfds(client, allfds) | |
| return parent_r, parent_w | |
| except: | |
| os.close(parent_r) | |
| os.close(parent_w) | |
| raise | |
| finally: | |
| os.close(child_r) | |
| os.close(child_w) | |
| def ensure_running(self): | |
| '''Make sure that a fork server is running. | |
| This can be called from any process. Note that usually a child | |
| process will just reuse the forkserver started by its parent, so | |
| ensure_running() will do nothing. | |
| ''' | |
| with self._lock: | |
| resource_tracker.ensure_running() | |
| if self._forkserver_pid is not None: | |
| # forkserver was launched before, is it still running? | |
| pid, status = os.waitpid(self._forkserver_pid, os.WNOHANG) | |
| if not pid: | |
| # still alive | |
| return | |
| # dead, launch it again | |
| os.close(self._forkserver_alive_fd) | |
| self._forkserver_authkey = None | |
| self._forkserver_address = None | |
| self._forkserver_alive_fd = None | |
| self._forkserver_pid = None | |
| # gh-144503: sys_argv is passed as real argv elements after the | |
| # ``-c cmd`` rather than repr'd into main_kws so that a large | |
| # parent sys.argv cannot push the single ``-c`` command string | |
| # over the OS per-argument length limit (MAX_ARG_STRLEN on Linux). | |
| # The child sees them as sys.argv[1:]. | |
| cmd = ('import sys; ' | |
| 'from multiprocessing.forkserver import main; ' | |
| 'main(%d, %d, %r, sys_argv=sys.argv[1:], **%r)') | |
| main_kws = {} | |
| sys_argv = None | |
| if self._preload_modules: | |
| data = spawn.get_preparation_data('ignore') | |
| if 'sys_path' in data: | |
| main_kws['sys_path'] = data['sys_path'] | |
| if 'init_main_from_path' in data: | |
| main_kws['main_path'] = data['init_main_from_path'] | |
| if 'sys_argv' in data: | |
| sys_argv = data['sys_argv'] | |
| if self._preload_on_error != 'ignore': | |
| main_kws['on_error'] = self._preload_on_error | |
| with socket.socket(socket.AF_UNIX) as listener: | |
| address = connection.arbitrary_address('AF_UNIX') | |
| listener.bind(address) | |
| if not util.is_abstract_socket_namespace(address): | |
| os.chmod(address, 0o600) | |
| listener.listen() | |
| # all client processes own the write end of the "alive" pipe; | |
| # when they all terminate the read end becomes ready. | |
| alive_r, alive_w = os.pipe() | |
| # A short lived pipe to initialize the forkserver authkey. | |
| authkey_r, authkey_w = os.pipe() | |
| try: | |
| fds_to_pass = [listener.fileno(), alive_r, authkey_r] | |
| main_kws['authkey_r'] = authkey_r | |
| cmd %= (listener.fileno(), alive_r, self._preload_modules, | |
| main_kws) | |
| exe = spawn.get_executable() | |
| args = [exe] + util._args_from_interpreter_flags() | |
| args += ['-c', cmd] | |
| if sys_argv is not None: | |
| args += sys_argv | |
| pid = util.spawnv_passfds(exe, args, fds_to_pass) | |
| except: | |
| os.close(alive_w) | |
| os.close(authkey_w) | |
| raise | |
| finally: | |
| os.close(alive_r) | |
| os.close(authkey_r) | |
| # Authenticate our control socket to prevent access from | |
| # processes we have not shared this key with. | |
| try: | |
| self._forkserver_authkey = os.urandom(_AUTHKEY_LEN) | |
| os.write(authkey_w, self._forkserver_authkey) | |
| finally: | |
| os.close(authkey_w) | |
| self._forkserver_address = address | |
| self._forkserver_alive_fd = alive_w | |
| self._forkserver_pid = pid | |
| # | |
| # | |
| # | |
| def _handle_import_error(on_error, modinfo, exc, *, warn_stacklevel): | |
| """Handle an import error according to the on_error policy.""" | |
| match on_error: | |
| case 'fail': | |
| raise | |
| case 'warn': | |
| warnings.warn( | |
| f"Failed to preload {modinfo}: {exc}", | |
| ImportWarning, | |
| stacklevel=warn_stacklevel + 1 | |
| ) | |
| case 'ignore': | |
| pass | |
| def _handle_preload(preload, main_path=None, sys_path=None, sys_argv=None, | |
| on_error='ignore'): | |
| """Handle module preloading with configurable error handling. | |
| Args: | |
| preload: List of module names to preload. | |
| main_path: Path to __main__ module if '__main__' is in preload. | |
| sys_path: sys.path to use for imports (None means use current). | |
| sys_argv: sys.argv to use (None means use current). | |
| on_error: How to handle import errors ("ignore", "warn", or "fail"). | |
| """ | |
| if not preload: | |
| return | |
| if sys_argv is not None: | |
| sys.argv[:] = sys_argv | |
| if sys_path is not None: | |
| sys.path[:] = sys_path | |
| if '__main__' in preload and main_path is not None: | |
| process.current_process()._inheriting = True | |
| try: | |
| spawn.import_main_path(main_path) | |
| except Exception as e: | |
| # Catch broad Exception because import_main_path() uses | |
| # runpy.run_path() which executes the script and can raise | |
| # any exception, not just ImportError | |
| _handle_import_error( | |
| on_error, f"__main__ from {main_path!r}", e, warn_stacklevel=2 | |
| ) | |
| finally: | |
| del process.current_process()._inheriting | |
| for modname in preload: | |
| try: | |
| __import__(modname) | |
| except ImportError as e: | |
| _handle_import_error( | |
| on_error, f"module {modname!r}", e, warn_stacklevel=2 | |
| ) | |
| # gh-135335: flush stdout/stderr in case any of the preloaded modules | |
| # wrote to them, otherwise children might inherit buffered data | |
| util._flush_std_streams() | |
| def main(listener_fd, alive_r, preload, main_path=None, sys_path=None, | |
| *, sys_argv=None, authkey_r=None, on_error='ignore'): | |
| """Run forkserver.""" | |
| if authkey_r is not None: | |
| try: | |
| authkey = os.read(authkey_r, _AUTHKEY_LEN) | |
| assert len(authkey) == _AUTHKEY_LEN, f'{len(authkey)} < {_AUTHKEY_LEN}' | |
| finally: | |
| os.close(authkey_r) | |
| else: | |
| authkey = b'' | |
| _handle_preload(preload, main_path, sys_path, sys_argv, on_error) | |
| util._close_stdin() | |
| sig_r, sig_w = os.pipe() | |
| os.set_blocking(sig_r, False) | |
| os.set_blocking(sig_w, False) | |
| def sigchld_handler(*_unused): | |
| # Dummy signal handler, doesn't do anything | |
| pass | |
| handlers = { | |
| # unblocking SIGCHLD allows the wakeup fd to notify our event loop | |
| signal.SIGCHLD: sigchld_handler, | |
| # protect the process from ^C | |
| signal.SIGINT: signal.SIG_IGN, | |
| } | |
| old_handlers = {sig: signal.signal(sig, val) | |
| for (sig, val) in handlers.items()} | |
| # calling os.write() in the Python signal handler is racy | |
| signal.set_wakeup_fd(sig_w) | |
| # map child pids to client fds | |
| pid_to_fd = {} | |
| with socket.socket(socket.AF_UNIX, fileno=listener_fd) as listener, \ | |
| selectors.DefaultSelector() as selector: | |
| _forkserver._forkserver_address = listener.getsockname() | |
| selector.register(listener, selectors.EVENT_READ) | |
| selector.register(alive_r, selectors.EVENT_READ) | |
| selector.register(sig_r, selectors.EVENT_READ) | |
| while True: | |
| try: | |
| while True: | |
| rfds = [key.fileobj for (key, events) in selector.select()] | |
| if rfds: | |
| break | |
| if alive_r in rfds: | |
| # EOF because no more client processes left | |
| assert os.read(alive_r, 1) == b'', "Not at EOF?" | |
| raise SystemExit | |
| if sig_r in rfds: | |
| # Got SIGCHLD | |
| os.read(sig_r, 65536) # exhaust | |
| while True: | |
| # Scan for child processes | |
| try: | |
| pid, sts = os.waitpid(-1, os.WNOHANG) | |
| except ChildProcessError: | |
| break | |
| if pid == 0: | |
| break | |
| child_w = pid_to_fd.pop(pid, None) | |
| if child_w is not None: | |
| returncode = os.waitstatus_to_exitcode(sts) | |
| # Send exit code to client process | |
| try: | |
| write_signed(child_w, returncode) | |
| except BrokenPipeError: | |
| # client vanished | |
| pass | |
| os.close(child_w) | |
| else: | |
| # This shouldn't happen really | |
| warnings.warn('forkserver: waitpid returned ' | |
| 'unexpected pid %d' % pid) | |
| if listener in rfds: | |
| # Incoming fork request | |
| with listener.accept()[0] as s: | |
| try: | |
| if authkey: | |
| wrapped_s = connection.Connection(s.fileno()) | |
| # The other side of this exchange happens in | |
| # in connect_to_new_process(). | |
| try: | |
| connection.deliver_challenge( | |
| wrapped_s, authkey) | |
| connection.answer_challenge( | |
| wrapped_s, authkey) | |
| finally: | |
| wrapped_s._detach() | |
| del wrapped_s | |
| # Receive fds from client | |
| fds = reduction.recvfds(s, MAXFDS_TO_SEND + 1) | |
| except (EOFError, BrokenPipeError, AuthenticationError): | |
| s.close() | |
| continue | |
| if len(fds) > MAXFDS_TO_SEND: | |
| raise RuntimeError( | |
| "Too many ({0:n}) fds to send".format( | |
| len(fds))) | |
| child_r, child_w, *fds = fds | |
| s.close() | |
| pid = os.fork() | |
| if pid == 0: | |
| # Child | |
| code = 1 | |
| try: | |
| listener.close() | |
| selector.close() | |
| unused_fds = [alive_r, child_w, sig_r, sig_w] | |
| unused_fds.extend(pid_to_fd.values()) | |
| atexit._clear() | |
| atexit.register(util._exit_function) | |
| code = _serve_one(child_r, fds, | |
| unused_fds, | |
| old_handlers) | |
| except Exception: | |
| sys.excepthook(*sys.exc_info()) | |
| sys.stderr.flush() | |
| finally: | |
| atexit._run_exitfuncs() | |
| os._exit(code) | |
| else: | |
| # Send pid to client process | |
| try: | |
| write_signed(child_w, pid) | |
| except BrokenPipeError: | |
| # client vanished | |
| pass | |
| pid_to_fd[pid] = child_w | |
| os.close(child_r) | |
| for fd in fds: | |
| os.close(fd) | |
| except OSError as e: | |
| if e.errno != errno.ECONNABORTED: | |
| raise | |
| def _serve_one(child_r, fds, unused_fds, handlers): | |
| # close unnecessary stuff and reset signal handlers | |
| signal.set_wakeup_fd(-1) | |
| for sig, val in handlers.items(): | |
| signal.signal(sig, val) | |
| for fd in unused_fds: | |
| os.close(fd) | |
| (_forkserver._forkserver_alive_fd, | |
| resource_tracker._resource_tracker._fd, | |
| *_forkserver._inherited_fds) = fds | |
| # Run process object received over pipe | |
| parent_sentinel = os.dup(child_r) | |
| code = spawn._main(child_r, parent_sentinel) | |
| return code | |
| # | |
| # Read and write signed numbers | |
| # | |
| def read_signed(fd): | |
| data = bytearray(SIGNED_STRUCT.size) | |
| unread = memoryview(data) | |
| while unread: | |
| count = os.readinto(fd, unread) | |
| if count == 0: | |
| raise EOFError('unexpected EOF') | |
| unread = unread[count:] | |
| return SIGNED_STRUCT.unpack(data)[0] | |
| def write_signed(fd, n): | |
| msg = SIGNED_STRUCT.pack(n) | |
| while msg: | |
| nbytes = os.write(fd, msg) | |
| if nbytes == 0: | |
| raise RuntimeError('should not get here') | |
| msg = msg[nbytes:] | |
| # | |
| # | |
| # | |
| _forkserver = ForkServer() | |
| ensure_running = _forkserver.ensure_running | |
| get_inherited_fds = _forkserver.get_inherited_fds | |
| connect_to_new_process = _forkserver.connect_to_new_process | |
| set_forkserver_preload = _forkserver.set_forkserver_preload |