graphyte/utils/process_utils.py

# Copyright 2018 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

from __future__ import print_function

import contextlib
import copy
import getpass
import logging
import os
import pipes
import re
import signal
import subprocess
import sys
import threading
import time
import traceback

from six.moves import queue
from six.moves import StringIO


PIPE = subprocess.PIPE


# File descriptor for /dev/null.
dev_null = None


def GetLines(data, strip=False):
  """Returns a list of all lines in data.

  Args:
    strip: If True, each line is stripped.
  """
  ret = StringIO(data).readlines()
  if strip:
    ret = [x.strip() for x in ret]
  return ret


def OpenDevNull():
  """Opens and returns a readable/writable file pointing to /dev/null.

  The file object may be reused.
  """
  global dev_null  # pylint: disable=W0603
  if not dev_null:
    # There is a possible race condition here, but it is extremely
    # unlikely and won't hurt anyway (we'll just have multiple files
    # pointing at /dev/null).
    dev_null = open(os.devnull, 'r+')

  return dev_null


def IsProcessAlive(pid, ppid=None):
  """Returns true if the named process is alive and not a zombie.

  A PPID (parent PID) can be provided to be more specific to which process you
  are watching.  If there is a process with the same PID running but the PPID is
  not the same, then this is unlikely to be the same process, but a newly
  started one.  The function will return False in this case.

  Args:
    pid: process PID for checking
    ppid: specified the PID of the parent of given process.  If the PPID does
      not match, we assume that the named process is done, and we are looking at
      another process, the function returns False in this case.
  """
  try:
    with open('/proc/%d/stat' % pid) as f:
      stat = f.readline().split()
      if ppid is not None and int(stat[3]) != ppid:
        return False
      return stat[2] != 'Z'
  except IOError:
    return False


def CheckOutput(*args, **kwargs):
  """Runs command and returns its output.

  It is like subprocess.check_output but with the extra flexibility of Spawn.

  Args:
    Refer Spawn.

  Returns:
    stdout

  Raises:
    subprocess.CalledProcessError if returncode != 0.
  """
  kwargs['check_output'] = True
  return Spawn(*args, **kwargs).stdout_data


def SpawnOutput(*args, **kwargs):
  """Runs command and returns its output.

  Like CheckOutput. But it won't raise exception unless you set
  check_output=True.

  Args:
    Refer Spawn.

  Returns:
    stdout
  """
  kwargs['read_stdout'] = True
  return Spawn(*args, **kwargs).stdout_data


def LogAndCheckCall(*args, **kwargs):
  """Logs a command and invokes subprocess.check_call."""
  logging.info('Running: %s', ' '.join(pipes.quote(arg) for arg in args[0]))
  return subprocess.check_call(*args, **kwargs)


def LogAndCheckOutput(*args, **kwargs):
  """Logs a command and invokes subprocess.check_output."""
  logging.info('Running: %s', ' '.join(pipes.quote(arg) for arg in args[0]))
  return CheckOutput(*args, **kwargs)


class _ExtendedPopen(subprocess.Popen):
  """Popen subclass supported a few extra methods.

  Attributes:
    stdout_data, stderr_data: Data read by communicate().  These are set by
      the Spawn call if read_stdout/read_stderr are True.
  """
  stdout_data = None
  stderr_data = None

  def stdout_lines(self, strip=False):
    """Returns lines in stdout_data as a list.

    Args:
      strip: If True, each line is stripped.
    """
    return GetLines(self.stdout_data, strip)

  def stderr_lines(self, strip=False):
    """Returns lines in stderr_data as a list.

    Args:
      strip: If True, each line is stripped.
    """
    return GetLines(self.stderr_data, strip)

  def communicate(self, *args, **kwargs):
    if self.stdout_data is None and self.stderr_data is None:
      return super(_ExtendedPopen, self).communicate(*args, **kwargs)
    else:
      return self.stdout_data, self.stderr_data


def Spawn(args, **kwargs):
  """Popen wrapper with extra functionality:

    - Sets close_fds to True by default.  (You may still set
      close_fds=False to leave all fds open.)
    - Provides a consistent interface to functionality like the call,
      check_call, and check_output functions in subprocess.

  To get a command's output, logging stderr if the process fails:

    # Doesn't check retcode
    Spawn(['cmd'], read_stdout=True, log_stderr_on_error=True).stdout_data
    # Throws CalledProcessError on error
    Spawn(['cmd'], read_stdout=True, log_stderr_on_error=True,
          check_call=True).stdout_data

  To get a command's stdout and stderr, without checking the retcode:

    stdout, stderr = Spawn(
        ['cmd'], read_stdout=True, read_stderr=True).communicate()

  Args:
    log: Do a logging.info before running the command, or to any
      logging object to call its info method.
    stdout: Same as subprocess.Popen, but may be set to DEV_NULL to discard
      all stdout.
    stderr: Same as subprocess.Popen, but may be set to DEV_NULL to discard
      all stderr.
    call: Wait for the command to complete.
    check_call: Wait for the command to complete, throwing an
      exception if it fails.  This implies call=True.  This may be either
      True to signify that any non-zero exit status is failure, or a function
      that takes a returncode and returns True if that returncode is
      considered OK (e.g., lambda returncode: returncode in [0,1]).
    check_output: Wait for the command to complete, throwing an
      exception if it fails, and saves the contents to the return
      object's stdout_data attribute.  Implies check_call=True and
      read_stdout=True.
    log_stderr_on_error: Log stderr only if the command fails.
      Implies read_stderr=True and call=True.
    read_stdout: Wait for the command to complete, saving the contents
      to the return object's stdout_data attribute.  This implies
      call=True and stdout=PIPE.
    ignore_stdout: Ignore stdout.
    read_stderr: Wait for the command to complete, saving the contents
      to the return object's stderr_data attribute.  This implies
      call=True and stderr=PIPE.
    ignore_stderr: Ignore stderr.
    sudo: Prepend sudo to arguments if user is not root.
    env: Same as subprocess.Popen, set-up environment parameters if needed.

  Returns/Raises:
    Same as Popen.
  """
  kwargs.setdefault('close_fds', True)

  logger = logging
  log = kwargs.pop('log', False)
  if kwargs.get('shell'):
    args_to_log = args
  else:
    args_to_log = ' '.join(pipes.quote(arg) for arg in args)

  if log:
    if log != True:
      logger = log
    message = 'Running command: "%s"' % args_to_log
    if 'cwd' in kwargs:
      message += ' in %s' % kwargs['cwd']
    logger.info(message)

  call = kwargs.pop('call', False)
  check_call = kwargs.pop('check_call', False)
  check_output = kwargs.pop('check_output', False)
  read_stdout = kwargs.pop('read_stdout', False)
  ignore_stdin = kwargs.pop('ignore_stdin', False)
  ignore_stdout = kwargs.pop('ignore_stdout', False)
  read_stderr = kwargs.pop('read_stderr', False)
  ignore_stderr = kwargs.pop('ignore_stderr', False)
  log_stderr_on_error = kwargs.pop('log_stderr_on_error', False)
  sudo = kwargs.pop('sudo', False)

  if sudo and getpass.getuser() != 'root':
    if kwargs.pop('shell', False):
      args = ['sudo', 'sh', '-c', args]
    else:
      args = ['sudo'] + args

  if ignore_stdin:
    assert not kwargs.get('stdin')
    kwargs['stdin'] = OpenDevNull()
  if ignore_stdout:
    assert not read_stdout
    assert not kwargs.get('stdout')
    kwargs['stdout'] = OpenDevNull()
  if ignore_stderr:
    assert not read_stderr
    assert not log_stderr_on_error
    assert not kwargs.get('stderr')
    kwargs['stderr'] = OpenDevNull()

  if check_output:
    check_call = check_call or True
    read_stdout = True
  if check_call:
    call = True
  if log_stderr_on_error:
    read_stderr = True
  if read_stdout:
    call = True
    assert kwargs.get('stdout') in [None, PIPE]
    kwargs['stdout'] = PIPE
  if read_stderr:
    call = True
    assert kwargs.get('stderr') in [None, PIPE]
    kwargs['stderr'] = PIPE

  if call and (not read_stdout) and kwargs.get('stdout') == PIPE:
    raise ValueError('Cannot use call=True argument with stdout=PIPE, '
                     'since OS buffers may get filled up')
  if call and (not read_stderr) and kwargs.get('stderr') == PIPE:
    raise ValueError('Cannot use call=True argument with stderr=PIPE, '
                     'since OS buffers may get filled up')

  process = _ExtendedPopen(args, **kwargs)

  if call:
    if read_stdout or read_stderr:
      stdout, stderr = process.communicate()
      if read_stdout:
        process.stdout_data = stdout
      if read_stderr:
        process.stderr_data = stderr
    else:
      # No need to communicate; just wait
      process.wait()

    if callable(check_call):
      failed = not check_call(process.returncode)
    else:
      failed = process.returncode != 0
    if failed:
      if log or log_stderr_on_error:
        message = 'Exit code %d from command: "%s"' % (
            process.returncode, args_to_log)
        if log_stderr_on_error:
          message += '; stderr: """\n%s\n"""' % process.stderr_data
        logger.error(message)

      if check_call:
        raise subprocess.CalledProcessError(process.returncode, args)

  return process


def TerminateOrKillProcess(process, wait_seconds=1, sudo=False):
  """Terminates a process and waits for it.

  The function sends SIGTERM to terminate the process, if it's not terminated
  in wait_seconds, then sends a SIGKILL.
  """
  pid = process.pid
  logging.info('Stopping process %d', pid)
  if sudo:
    return Spawn(['kill', str(pid)], sudo=True, check_call=True, log=True)
  else:
    process.terminate()

  reaped = threading.Event()

  def WaitAndKill():
    reaped.wait(wait_seconds)
    if not reaped.is_set():
      try:
        logging.info('Sending SIGKILL to process %d', pid)
        process.kill()
      except:  # pylint: disable=W0702
        pass
  thread = threading.Thread(target=WaitAndKill)
  thread.start()
  process.wait()
  reaped.set()
  thread.join()
  logging.info('Process %d stopped', pid)


def KillProcessTree(process, caption):
  """Kills a process and all its subprocesses.

  Args:
    process: The process to kill (opened with the subprocess module).
    caption: A caption describing the process.
  """
  # os.kill does not kill child processes. os.killpg kills all processes
  # sharing same group (and is usually used for killing process tree). But in
  # our case, to preserve PGID for autotest and upstart service, we need to
  # iterate through each level until leaf of the tree.

  def get_all_pids(root):
    ps_output = Spawn(['ps', '--no-headers', '-eo', 'pid,ppid'],
                      stdout=subprocess.PIPE)
    children = {}
    for line in ps_output.stdout:
      match = re.findall(r'\d+', line)
      children.setdefault(int(match[1]), []).append(int(match[0]))
    pids = []

    def add_children(pid):
      pids.append(pid)
      list(map(add_children, children.get(pid, [])))
    add_children(root)
    # Reverse the list to first kill children then parents.
    # Note reversed(pids) will return an iterator instead of real list, so
    # we must explicitly call pids.reverse() here.
    pids.reverse()
    return pids

  pids = get_all_pids(process.pid)
  for sig in [signal.SIGTERM, signal.SIGKILL]:
    logging.info('Stopping %s (pid=%s)...', caption, sorted(pids))

    for _ in range(25):  # Try 25 times (200 ms between tries)
      for pid in pids:
        try:
          logging.info('Sending signal %s to %d', sig, pid)
          os.kill(pid, sig)
        except OSError:
          pass
      pids = list(filter(IsProcessAlive, pids))
      if not pids:
        return
      time.sleep(0.2)  # Sleep 200 ms and try again

  logging.warn('Failed to stop %s process. Ignoring.', caption)


def WaitEvent(event):
  """Waits for an event without timeout, without blocking signals.

  event.wait() masks all signals until the event is set; this can be used
  instead to make sure that the signal is delivered within 100 ms.

  Returns:
    True if the event is set (i.e., always, since there is no timeout).  This
      return value is used so that this method behaves the same way as
      event.wait().
  """
  while not event.is_set():
    event.wait(0.1)
  return True


def SpawnTee(args, **kwargs):
  """Spawns a process and emulates tee.

  Starts a process with Spawn, redirect_streams stderr of the process to its
  stdout, and writes stdout of the process to both sys.stdout and the specified
  file.

  Args:
    args: Same as Spawn.
    output_file: The file to write to in addition to stdout.

  Returns:
    The created process object.
  """
  output_file = kwargs.get('output_file')
  if not output_file:
    raise ValueError('output_file must be specified')
  stdout = kwargs.get('stdout', sys.stdout)

  message_queue = queue.Queue()

  def Tee(out_fd):
    """Writes available data from message queue to stdout and output file."""
    try:
      while True:
        line = message_queue.get(timeout=0.1)
        stdout.write(line)
        stdout.flush()
        out_fd.write(line)
        out_fd.flush()
    except queue.Empty:
      return

  def EnqueueOutput(in_fd):
    """Keeps enqueue messages read."""
    for line in iter(in_fd.readline, ''):
      message_queue.put(line)

  tee_kwargs = copy.deepcopy(kwargs)
  del tee_kwargs['output_file']
  tee_kwargs['stdout'] = subprocess.PIPE
  tee_kwargs['stderr'] = subprocess.STDOUT
  # Do not block waiting for process to end in Spawn, as we want to have live
  # message output.
  tee_kwargs['check_call'] = False
  tee_kwargs['check_output'] = False
  tee_kwargs['call'] = False

  proc = Spawn(args, **tee_kwargs)
  enqueue_thread = threading.Thread(target=EnqueueOutput, args=(proc.stdout,))
  enqueue_thread.daemon = True
  enqueue_thread.start()

  with open(output_file, 'w') as f:
    while proc.poll() is None:
      Tee(f)
    proc.wait()
    Tee(f)
    if proc.returncode != 0 and kwargs.get('check_call'):
      raise subprocess.CalledProcessError(proc.returncode, args)
    return proc


def StartDaemonThread(*args, **kwargs):
  """Creates, starts, and returns a daemon thread.

  Args:
    interrupt_when_crash: If true, the thread sends interrupt signal when
        exception uncaught.
    For other parameters see threading.Thread().
  """
  if kwargs.pop('interrupt_on_crash', False):
    # 'target' is the second parameter of threading.Thread()
    target = args[1] if len(args) > 1 else kwargs.get('target')

    def _target(*_args, **_kwargs):
      try:
        target(*_args, **_kwargs)
      except Exception:
        logging.error(traceback.format_exc())
        os.kill(os.getpid(), signal.SIGINT)

    if len(args) > 1:
      args[1] = _target
    else:
      kwargs['target'] = _target

  thread = threading.Thread(*args, **kwargs)
  thread.daemon = True
  thread.start()
  return thread


class DummyFile(object):
  def write(self, x):  # pylint: disable=W0613
    pass

  def read(self, x):  # pylint: disable=W0613
    return ''


@contextlib.contextmanager
def RedirectStandardStreams(stdin=None, stdout=None, stderr=None):
  """Redirect standard stream.

  Args:
    stdin: A file object to override standard input.
    stdout: A file object to override standard output.
    stderr: A file object to override standard error.
    If stdin, stdout, stderr is None, then the stream is not redirected.

  Raises:
    IOError: raise the exception if the standard streams is redirected again
             within the context.
  """
  args = {'stdin': stdin, 'stdout': stdout, 'stderr': stderr}
  redirect_streams = dict((k, v) for k, v in args.iteritems() if v is not None)
  old_streams = dict((k, sys.__dict__[k]) for k in redirect_streams)

  for k, v in redirect_streams.iteritems():
    sys.__dict__[k] = v

  yield

  changed = dict((k, sys.__dict__[k]) for (k, v) in redirect_streams.iteritems()
                 if v is not sys.__dict__[k])
  if changed:
    raise IOError('Unexpected stadard stream redirections: %r' % changed)
  for k, v in old_streams.iteritems():
    sys.__dict__[k] = v