1
0
mirror of https://github.com/ARM-software/devlib.git synced 2025-03-04 17:27:51 +00:00

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.
This commit is contained in:
Douglas Raillard 2024-05-30 10:41:50 +01:00 committed by Marc Bonnici
parent eb9e0c9870
commit 9ec36e9040
6 changed files with 290 additions and 182 deletions

View File

@ -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 "$@"

View File

@ -17,6 +17,7 @@ from abc import ABC, abstractmethod
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 logging
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 @@ class BackgroundCommand(ABC):
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 @@ class BackgroundCommand(ABC):
: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 @@ class BackgroundCommand(ABC):
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 @@ class TransferHandleBase(ABC):
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 @@ class PopenTransferHandle(TransferHandleBase):
self.sample_size = lambda: sample_size(dest)
self.bg_cmd = bg_cmd
self.popen = popen
self.last_sample = 0
@staticmethod
@ -671,7 +738,7 @@ class PopenTransferHandle(TransferHandleBase):
return int(out.split()[0])
def cancel(self):
self.bg_cmd.cancel()
self.popen.terminate()
def isactive(self):
try:

View File

@ -147,6 +147,7 @@ class LocalConnection(ConnectionBase):
def preexec_fn():
os.setpgrp()
def make_init_kwargs(command):
popen = subprocess.Popen(
command,
stdout=stdout,
@ -155,8 +156,16 @@ class LocalConnection(ConnectionBase):
shell=True,
preexec_fn=preexec_fn,
)
bg_cmd = PopenBackgroundCommand(self, popen)
return bg_cmd
return dict(
popen=popen,
)
return PopenBackgroundCommand.from_factory(
conn=self,
cmd=command,
as_root=as_root,
make_init_kwargs=make_init_kwargs,
)
def _close(self):
pass

View File

@ -281,7 +281,7 @@ class Target(object):
@property
def shutils(self):
if self._shutils is None:
self._setup_shutils()
self._setup_scripts()
return self._shutils
def is_running(self, comm):
@ -591,7 +591,7 @@ class Target(object):
@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)
@ -1586,8 +1586,9 @@ fi
# 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:
@ -1598,6 +1599,8 @@ fi
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):

View File

@ -40,7 +40,7 @@ from shlex import quote
from devlib.exception import TargetTransientError, TargetStableError, HostError, TargetTransientCalledProcessError, TargetStableCalledProcessError, AdbRootError
from devlib.utils.misc import check_output, which, ABI_MAP, redirect_streams, get_subprocess
from devlib.connection import ConnectionBase, AdbBackgroundCommand, PopenBackgroundCommand, PopenTransferHandle
from devlib.connection import ConnectionBase, AdbBackgroundCommand, PopenTransferHandle
logger = logging.getLogger('android')
@ -341,7 +341,7 @@ class AdbConnection(ConnectionBase):
if timeout:
adb_command(self.device, command, timeout=timeout, adb_server=self.adb_server, adb_port=self.adb_port)
else:
bg_cmd = adb_command_background(
popen = adb_command_popen(
device=self.device,
conn=self,
command=command,
@ -351,12 +351,12 @@ class AdbConnection(ConnectionBase):
handle = PopenTransferHandle(
manager=self.transfer_manager,
bg_cmd=bg_cmd,
popen=popen,
dest=dest,
direction=action
)
with bg_cmd, self.transfer_manager.manage(sources, dest, action, handle):
bg_cmd.communicate()
with popen, self.transfer_manager.manage(sources, dest, action, handle):
popen.communicate()
# pylint: disable=unused-argument
def execute(self, command, timeout=None, check_exit_code=False,
@ -387,12 +387,18 @@ class AdbConnection(ConnectionBase):
return bg_cmd
def _background(self, command, stdout, stderr, as_root):
adb_shell, pid = adb_background_shell(self, command, stdout, stderr, as_root)
bg_cmd = AdbBackgroundCommand(
conn=self,
adb_popen=adb_shell,
def make_init_kwargs(command):
adb_popen, pid = adb_background_shell(self, command, stdout, stderr, as_root)
return dict(
adb_popen=adb_popen,
pid=pid,
as_root=as_root
)
bg_cmd = AdbBackgroundCommand.from_factory(
conn=self,
cmd=command,
as_root=as_root,
make_init_kwargs=make_init_kwargs,
)
return bg_cmd
@ -756,12 +762,11 @@ def adb_command(device, command, timeout=None, adb_server=None, adb_port=None):
return output
def adb_command_background(device, conn, command, adb_server=None, adb_port=None):
full_command = get_adb_command(device, command, adb_server, adb_port)
logger.debug(full_command)
popen = get_subprocess(full_command, shell=True)
cmd = PopenBackgroundCommand(conn=conn, popen=popen)
return cmd
def adb_command_popen(device, conn, command, adb_server=None, adb_port=None):
command = get_adb_command(device, command, adb_server, adb_port)
logger.debug(command)
popen = get_subprocess(command, shell=True)
return popen
def grant_app_permissions(target, package):

View File

@ -613,10 +613,10 @@ class SshConnection(SshConnectionBase):
return self._background(command, stdout, stderr, as_root)
def _background(self, command, stdout, stderr, as_root):
orig_command = command
stdout, stderr, command = redirect_streams(stdout, stderr, command)
def make_init_kwargs(command):
_stdout, _stderr, _command = redirect_streams(stdout, stderr, command)
command = "printf '%s\n' $$; exec sh -c {}".format(quote(command))
_command = "printf '%s\n' $$; exec sh -c {}".format(quote(_command))
channel = self._make_channel()
def executor(cmd, timeout):
@ -630,7 +630,7 @@ class SshConnection(SshConnectionBase):
)
stdin, stdout_in, stderr_in = self._execute_command(
command,
_command,
as_root=as_root,
log=False,
timeout=None,
@ -638,16 +638,16 @@ class SshConnection(SshConnectionBase):
)
pid = stdout_in.readline()
if not pid:
stderr = stderr_in.read()
_stderr = stderr_in.read()
if channel.exit_status_ready():
ret = channel.recv_exit_status()
else:
ret = 126
raise subprocess.CalledProcessError(
ret,
command,
_command,
b'',
stderr,
_stderr,
)
pid = int(pid)
@ -680,8 +680,8 @@ class SshConnection(SshConnectionBase):
out_streams = {
name: create_out_stream(stream_in, stream_out)
for stream_in, stream_out, name in (
(stdout_in, stdout, 'stdout'),
(stderr_in, stderr, 'stderr'),
(stdout_in, _stdout, 'stdout'),
(stderr_in, _stderr, 'stderr'),
)
}
@ -734,9 +734,7 @@ class SshConnection(SshConnectionBase):
)
redirect_thread.start()
return ParamikoBackgroundCommand(
conn=self,
as_root=as_root,
return dict(
chan=channel,
pid=pid,
stdin=stdin,
@ -744,7 +742,13 @@ class SshConnection(SshConnectionBase):
stdout=out_streams['stdout'][0],
stderr=out_streams['stderr'][0],
redirect_thread=redirect_thread,
cmd=orig_command,
)
return ParamikoBackgroundCommand.from_factory(
conn=self,
cmd=command,
as_root=as_root,
make_init_kwargs=make_init_kwargs,
)
def _close(self):