connection: Support all signals in BackgroundCommand.send_signal()

Support sending any signal to background commands, instead of only
supporting properly SIGKILL/SIGTERM/SIGQUIT.

The main issue is figuring out what PID to send the signal to, as the
devlib API allows running a whole snippet of shell script that typically
is wrapped under many layers of sh -c and sudo calls. In order to lift
the ambiguity, the user has access to a "devlib-signal-target" command
that points devlib at what process should be the target of signals:

    # Run a "setup" command, then the main command that will receive the
    # signals
    cmd = 'echo setup; devlib-signal-target echo hello world'
    with target.background(cmd) as bg:
	bg.communicate()

The devlib-signal-target script can only be used once per background
command, so that it is never ambiguous what process is targeted, and so
that the Python code can cache the target PID.  Subsequent invocations
of devlib-signal-target will fail.

devlib/bin/scripts/devlib-signal-target

@@ -0,0 +1,20 @@
+(
+    # If there is no data dir, it means we are not running as a background
+    # command so we just do nothing
+    if [ -e "$_DEVLIB_BG_CMD_DATA_DIR" ]; then
+        pid_file="$_DEVLIB_BG_CMD_DATA_DIR/pid"
+        # Atomically check if the PID file already exist and make the write
+        # fail if it already does. This way we don't have any race condition
+        # with the Python API, as there is either no PID or the same PID for
+        # the duration of the command
+        set -o noclobber
+        if ! printf "%u\n" $$ > "$pid_file"; then
+            echo "$0 was already called for this command" >&2
+            exit 1
+        fi
+    fi
+) || exit $?
+
+# Use exec so that the PID of the command we run is the same as the current $$
+# PID that we just registered
+exec "$@"

devlib/connection.py

@@ -17,6 +17,7 @@
 from contextlib import contextmanager, nullcontext
 from shlex import quote
 import os
+from pathlib import Path
 import signal
 import subprocess
 import threading
@@ -25,14 +26,11 @@
 import select
 import fcntl
 
-from devlib.utils.misc import InitCheckpoint
+from devlib.utils.misc import InitCheckpoint, memoized
 
 _KILL_TIMEOUT = 3
 
 
-def _kill_pgid_cmd(pgid, sig, busybox):
-    return '{} kill -{} -{}'.format(busybox, sig.value, pgid)
-
 def _popen_communicate(bg, popen, input, timeout):
     try:
         stdout, stderr = popen.communicate(input=input, timeout=timeout)
@@ -130,8 +128,11 @@ class BackgroundCommand(ABC):
     semantic as :class:`subprocess.Popen`.
     """
 
-    def __init__(self, conn):
+    def __init__(self, conn, data_dir, cmd, as_root):
         self.conn = conn
+        self._data_dir = data_dir
+        self.as_root = as_root
+        self.cmd = cmd
 
         # Poll currently opened background commands on that connection to make
         # them deregister themselves if they are completed. This avoids
@@ -147,15 +148,65 @@ def __init__(self, conn):
 
         conn._current_bg_cmds.add(self)
 
+    @classmethod
+    def from_factory(cls, conn, cmd, as_root, make_init_kwargs):
+        cmd, data_dir = cls._with_data_dir(conn, cmd)
+        return cls(
+            conn=conn,
+            data_dir=data_dir,
+            cmd=cmd,
+            as_root=as_root,
+            **make_init_kwargs(cmd),
+        )
+
     def _deregister(self):
         try:
             self.conn._current_bg_cmds.remove(self)
         except KeyError:
             pass
 
-    @abstractmethod
-    def _send_signal(self, sig):
-        pass
+    @property
+    def _pid_file(self):
+        return str(Path(self._data_dir, 'pid'))
+
+    @property
+    @memoized
+    def _targeted_pid(self):
+        """
+        PID of the process pointed at by ``devlib-signal-target`` command.
+        """
+        path = quote(self._pid_file)
+        busybox = quote(self.conn.busybox)
+
+        def execute(cmd):
+            return self.conn.execute(cmd, as_root=self.as_root)
+
+        while self.poll() is None:
+            try:
+                pid = execute(f'{busybox} cat {path}')
+            except subprocess.CalledProcessError:
+                time.sleep(0.01)
+            else:
+                if pid.endswith('\n'):
+                    return int(pid.strip())
+                else:
+                    # We got a partial write in the PID file
+                    continue
+
+        raise ValueError(f'The background commmand did not use devlib-signal-target wrapper to designate which command should be the target of signals')
+
+    @classmethod
+    def _with_data_dir(cls, conn, cmd):
+        busybox = quote(conn.busybox)
+        data_dir = conn.execute(f'{busybox} mktemp -d').strip()
+        cmd = f'_DEVLIB_BG_CMD_DATA_DIR={data_dir} exec {busybox} sh -c {quote(cmd)}'
+        return cmd, data_dir
+
+    def _cleanup_data_dir(self):
+        path = quote(self._data_dir)
+        busybox = quote(self.conn.busybox)
+        cmd = f'{busybox} rm -r {path} || true'
+        self.conn.execute(cmd, as_root=self.as_root)
 
     def send_signal(self, sig):
         """
@@ -165,8 +216,29 @@ def send_signal(self, sig):
         :param signal: Signal to send.
         :type signal: signal.Signals
         """
+
+        def execute(cmd):
+            return self.conn.execute(cmd, as_root=self.as_root)
+
+        def send(sig):
+            busybox = quote(self.conn.busybox)
+            # If the command has already completed, we don't want to send a
+            # signal to another process that might have gotten that PID in the
+            # meantime.
+            if self.poll() is None:
+                if sig in (signal.SIGTERM, signal.SIGQUIT, signal.SIGKILL):
+                    # Use -PGID to target a process group rather than just the
+                    # process itself. This will work in any condition and will
+                    # not require cooperation from the command.
+                    execute(f'{busybox} kill -{sig.value} -{self.pid}')
+                else:
+                    # Other signals require cooperation from the shell command
+                    # so that it points to a specific process using
+                    # devlib-signal-target
+                    pid = self._targeted_pid
+                    execute(f'{busybox} kill -{sig.value} {pid}')
         try:
-            return self._send_signal(sig)
+            return send(sig)
         finally:
             # Deregister if the command has finished
             self.poll()
@@ -287,6 +359,7 @@ def close(self):
             return self._close()
         finally:
             self._deregister()
+            self._cleanup_data_dir()
 
     def __enter__(self):
         return self
@@ -300,13 +373,15 @@ class PopenBackgroundCommand(BackgroundCommand):
     :class:`subprocess.Popen`-based background command.
     """
 
-    def __init__(self, conn, popen):
-        super().__init__(conn=conn)
+    def __init__(self, conn, data_dir, cmd, as_root, popen):
+        super().__init__(
+            conn=conn,
+            data_dir=data_dir,
+            cmd=cmd,
+            as_root=as_root,
+        )
         self.popen = popen
 
-    def _send_signal(self, sig):
-        return os.killpg(self.popen.pid, sig)
-
     @property
     def stdin(self):
         return self.popen.stdin
@@ -354,26 +429,20 @@ class ParamikoBackgroundCommand(BackgroundCommand):
     """
     :mod:`paramiko`-based background command.
     """
-    def __init__(self, conn, chan, pid, as_root, cmd, stdin, stdout, stderr, redirect_thread):
-        super().__init__(conn=conn)
+    def __init__(self, conn, data_dir, cmd, as_root, chan, pid, stdin, stdout, stderr, redirect_thread):
+        super().__init__(
+            conn=conn,
+            data_dir=data_dir,
+            cmd=cmd,
+            as_root=as_root,
+        )
+
         self.chan = chan
-        self.as_root = as_root
         self._pid = pid
         self._stdin = stdin
         self._stdout = stdout
         self._stderr = stderr
         self.redirect_thread = redirect_thread
-        self.cmd = cmd
-
-    def _send_signal(self, sig):
-        # If the command has already completed, we don't want to send a signal
-        # to another process that might have gotten that PID in the meantime.
-        if self.poll() is not None:
-            return
-        # Use -PGID to target a process group rather than just the process
-        # itself
-        cmd = _kill_pgid_cmd(self.pid, sig, self.conn.busybox)
-        self.conn.execute(cmd, as_root=self.as_root)
 
     @property
     def pid(self):
@@ -517,18 +586,16 @@ class AdbBackgroundCommand(BackgroundCommand):
     ``adb``-based background command.
     """
 
-    def __init__(self, conn, adb_popen, pid, as_root):
-        super().__init__(conn=conn)
-        self.as_root = as_root
+    def __init__(self, conn, data_dir, cmd, as_root, adb_popen, pid):
+        super().__init__(
+            conn=conn,
+            data_dir=data_dir,
+            cmd=cmd,
+            as_root=as_root,
+        )
         self.adb_popen = adb_popen
         self._pid = pid
 
-    def _send_signal(self, sig):
-        self.conn.execute(
-            _kill_pgid_cmd(self.pid, sig, self.conn.busybox),
-            as_root=self.as_root,
-        )
-
     @property
     def stdin(self):
         return self.adb_popen.stdin
@@ -638,7 +705,7 @@ def cancel(self):
 
 
 class PopenTransferHandle(TransferHandleBase):
-    def __init__(self, bg_cmd, dest, direction, *args, **kwargs):
+    def __init__(self, popen, dest, direction, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
         if direction == 'push':
@@ -650,7 +717,7 @@ def __init__(self, bg_cmd, dest, direction, *args, **kwargs):
 
         self.sample_size = lambda: sample_size(dest)
 
-        self.bg_cmd = bg_cmd
+        self.popen = popen
         self.last_sample = 0
 
     @staticmethod
@@ -671,7 +738,7 @@ def _push_dest_size(self, dest):
         return int(out.split()[0])
 
     def cancel(self):
-        self.bg_cmd.cancel()
+        self.popen.terminate()
 
     def isactive(self):
         try:

devlib/host.py

@@ -145,16 +145,25 @@ def background(self, command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, as
         def preexec_fn():
             os.setpgrp()
 
-        popen = subprocess.Popen(
-            command,
-            stdout=stdout,
-            stderr=stderr,
-            stdin=subprocess.PIPE,
-            shell=True,
-            preexec_fn=preexec_fn,
+        def make_init_kwargs(command):
+            popen = subprocess.Popen(
+                command,
+                stdout=stdout,
+                stderr=stderr,
+                stdin=subprocess.PIPE,
+                shell=True,
+                preexec_fn=preexec_fn,
+            )
+            return dict(
+                popen=popen,
+            )
+
+        return PopenBackgroundCommand.from_factory(
+            conn=self,
+            cmd=command,
+            as_root=as_root,
+            make_init_kwargs=make_init_kwargs,
         )
-        bg_cmd = PopenBackgroundCommand(self, popen)
-        return bg_cmd
 
     def _close(self):
         pass

devlib/target.py

@@ -272,7 +272,7 @@ def page_size_kb(self):
     @property
     def shutils(self):
         if self._shutils is None:
-            self._setup_shutils()
+            self._setup_scripts()
         return self._shutils
 
     def is_running(self, comm):
@@ -567,7 +567,7 @@ def wait_boot_complete(self, timeout=10):
 
     @asyn.asyncf
     async def setup(self, executables=None):
-        await self._setup_shutils.asyn()
+        await self._setup_scripts.asyn()
 
         for host_exe in (executables or []):  # pylint: disable=superfluous-parens
             await self.install.asyn(host_exe)
@@ -1538,8 +1538,9 @@ def install_module(self, mod, **params):
     # internal methods
 
     @asyn.asyncf
-    async def _setup_shutils(self):
-        shutils_ifile = os.path.join(PACKAGE_BIN_DIRECTORY, 'scripts', 'shutils.in')
+    async def _setup_scripts(self):
+        scripts = os.path.join(PACKAGE_BIN_DIRECTORY, 'scripts')
+        shutils_ifile = os.path.join(scripts, 'shutils.in')
         with open(shutils_ifile) as fh:
             lines = fh.readlines()
         with tempfile.TemporaryDirectory() as folder:
@@ -1550,6 +1551,8 @@ async def _setup_shutils(self):
                     ofile.write(line)
             self._shutils = await self.install.asyn(shutils_ofile)
 
+        await self.install.asyn(os.path.join(scripts, 'devlib-signal-target'))
+
     @asyn.asyncf
     @call_conn
     async def _execute_util(self, command, timeout=None, check_exit_code=True, as_root=False):