meerschaum.utils.process

  1#! /usr/bin/env python
  2# -*- coding: utf-8 -*-
  3# vim:fenc=utf-8
  4
  5"""
  6Custom process-handling functions.
  7See `meerschaum.utils.pool` for multiprocessing and
  8`meerschaum.utils.threading` for threads.
  9"""
 10
 11from __future__ import annotations
 12
 13import os
 14import signal
 15import subprocess
 16import sys
 17import platform
 18
 19import meerschaum as mrsm
 20from meerschaum.utils.typing import Union, Optional, Any, Callable, Dict, Tuple
 21from meerschaum._internal.static import STATIC_CONFIG
 22
 23_child_processes = []
 24def signal_handler(sig, frame):
 25    for child in _child_processes:
 26        try:
 27            child.send_signal(sig)
 28            child.wait()
 29        except Exception:
 30            pass
 31
 32def run_process(
 33    *args,
 34    foreground: bool = False,
 35    as_proc: bool = False,
 36    line_callback: Optional[Callable[[bytes], Any]] = None,
 37    store_proc_dict: Optional[Dict[str, Any]] = None,
 38    store_proc_key: str = 'child_process',
 39    capture_output: bool = False,
 40    **kw: Any
 41) -> Union[int, subprocess.Popen]:
 42    """Original foreground solution found here:
 43    https://stackoverflow.com/questions/23826695/handling-keyboard-interrupt-when-using-subproccess
 44
 45    Parameters
 46    ----------
 47    *args:
 48        The sysargs to execute.
 49
 50    foreground: bool, default False
 51        If `True`, execute the process as a foreground process that passes Ctrl-C to children.
 52        From the original post:
 53        The "correct" way of spawning a new subprocess:
 54        signals like C-c must only go
 55        to the child process, and not to this python.
 56        
 57        Some side-info about "how ctrl-c works":
 58        https://unix.stackexchange.com/a/149756/1321
 59
 60    as_proc: bool, default False
 61        If `True`, return the `subprocess.Popen` object.
 62
 63    line_callback: Optional[Callable[[str], Any]], default None
 64        If provided, poll the process and execute the callback when `readline()` gets new text.
 65
 66    store_proc_dict: Optional[Dict[str, Any]], default None
 67        If provided, store the `subprocess.Popen` object under the key `store_proc_key`.
 68        Useful for accessing the process while it is polling in another thread.
 69
 70    store_proc_key: str, default 'child_process'
 71        If `store_proc_dict` is provided, store the process in the dictionary under this key.
 72
 73    kw: Any
 74        Additional keyword arguments to pass to `subprocess.Popen`.
 75
 76    Returns
 77    -------
 78    Either an int for the return code or a `subprocess.Popen` object.
 79    """
 80    try:
 81        import termios
 82    except ImportError:
 83        termios = None
 84
 85    if platform.system() == 'Windows':
 86        foreground = False
 87
 88    def print_line(line):
 89        sys.stdout.write(line.decode('utf-8'))
 90        sys.stdout.flush()
 91
 92  
 93    if capture_output or line_callback is not None:
 94        kw['stdout'] = subprocess.PIPE
 95        kw['stderr'] = subprocess.STDOUT
 96    elif os.environ.get(STATIC_CONFIG['environment']['daemon_id']):
 97        kw['stdout'] = subprocess.PIPE
 98        kw['stderr'] = subprocess.STDOUT
 99        if line_callback is None:
100            line_callback = print_line
101
102    if 'env' not in kw:
103        kw['env'] = os.environ
104
105    user_preexec_fn = kw.get("preexec_fn", None)
106
107    if foreground:
108        try:
109            old_pgrp = os.tcgetpgrp(sys.stdin.fileno())
110        except Exception as e:
111            termios = None
112        if termios:
113            try:
114                old_attr = termios.tcgetattr(sys.stdin.fileno())
115            except Exception as e:
116                termios = None
117
118    def new_pgid():
119        if user_preexec_fn:
120            user_preexec_fn()
121
122        # set a new process group id
123        os.setpgid(os.getpid(), os.getpid())
124
125        # generally, the child process should stop itself
126        # before exec so the parent can set its new pgid.
127        # (setting pgid has to be done before the child execs).
128        # however, Python 'guarantee' that `preexec_fn`
129        # is run before `Popen` returns.
130        # this is because `Popen` waits for the closure of
131        # the error relay pipe '`errpipe_write`',
132        # which happens at child's exec.
133        # this is also the reason the child can't stop itself
134        # in Python's `Popen`, since the `Popen` call would never
135        # terminate then.
136        # `os.kill(os.getpid(), signal.SIGSTOP)`
137
138    if foreground:
139        kw['preexec_fn'] = new_pgid
140
141    try:
142        child = subprocess.Popen(*args, **kw)
143        _child_processes.append(child)
144
145        # we can't set the process group id from the parent since the child
146        # will already have exec'd. and we can't SIGSTOP it before exec,
147        # see above.
148        # `os.setpgid(child.pid, child.pid)`
149
150        if foreground:
151            # set the child's process group as new foreground
152            try:
153                os.tcsetpgrp(sys.stdin.fileno(), child.pid)
154            except Exception as e:
155                pass
156            # revive the child,
157            # because it may have been stopped due to SIGTTOU or
158            # SIGTTIN when it tried using stdout/stdin
159            # after setpgid was called, and before we made it
160            # forward process by tcsetpgrp.
161            os.kill(child.pid, signal.SIGCONT)
162
163        # wait for the child to terminate
164        if store_proc_dict is not None:
165            store_proc_dict[store_proc_key] = child
166        _ret = poll_process(child, line_callback) if line_callback is not None else child.wait()
167        ret = _ret if not as_proc else child
168    except KeyboardInterrupt:
169        child.send_signal(signal.SIGINT)
170        ret = child.wait() if not as_proc else child
171    finally:
172        if foreground:
173            # we have to mask SIGTTOU because tcsetpgrp
174            # raises SIGTTOU to all current background
175            # process group members (i.e. us) when switching tty's pgrp
176            # it we didn't do that, we'd get SIGSTOP'd
177            hdlr = signal.signal(signal.SIGTTOU, signal.SIG_IGN)
178            # make us tty's foreground again
179            try:
180                os.tcsetpgrp(sys.stdin.fileno(), old_pgrp)
181            except Exception:
182                pass
183            # now restore the handler
184            signal.signal(signal.SIGTTOU, hdlr)
185            # restore terminal attributes
186            if termios:
187                termios.tcsetattr(sys.stdin.fileno(), termios.TCSADRAIN, old_attr)
188
189    return ret
190
191
192def poll_process(
193    proc: subprocess.Popen,
194    line_callback: Callable[[bytes], Any],
195    timeout_seconds: Union[int, float, None] = None,
196    timeout_callback: Optional[Callable[[Any], Any]] = None,
197    timeout_callback_args: Optional[Tuple[Any]] = None,
198    timeout_callback_kwargs: Optional[Dict[str, Any]] = None,
199) -> int:
200    """
201    Poll a process and execute a callback function for each line printed to the process's `stdout`.
202    """
203    from meerschaum.utils.threading import Timer
204    from meerschaum.utils.warnings import warn
205
206    def timeout_handler():
207        nonlocal timeout_callback_args, timeout_callback_kwargs
208        try:
209            if platform.system() != 'Windows':
210                ### The process being killed may have children.
211                os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
212            else:
213                proc.send_signal(signal.CTRL_BREAK_EVENT)
214            proc.terminate()
215        except Exception as e:
216            warn(f"Failed to kill process:\n{e}")
217        if timeout_callback_args is None:
218            timeout_callback_args = []
219        if timeout_callback_kwargs is None:
220            timeout_callback_kwargs = {}
221        timeout_callback(*timeout_callback_args, **timeout_callback_kwargs)
222
223    if timeout_seconds is not None:
224        watchdog_thread = Timer(timeout_seconds, timeout_handler)
225        watchdog_thread.daemon = True
226        watchdog_thread.start()
227
228    while proc.poll() is None:
229        line = proc.stdout.readline()
230        line_callback(line)
231
232    if timeout_seconds is not None:
233        watchdog_thread.cancel()
234
235    return proc.poll()
236
237
238def _stop_process(
239    proc: subprocess.Popen,
240    timeout_seconds: int = 8,
241):
242    """
243    Stop a `subproccess.Popen` object.
244    """
245    proc.terminate()
246    try:
247        proc.wait(timeout=timeout_seconds)
248    except subprocess.TimeoutExpired:
249        proc.kill()