py/utils/shell_utils.py

#!/usr/bin/env python
# Copyright (c) 2012 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

""" This module contains tools for running commands in a shell. """

import datetime
import os
import Queue
import select
import subprocess
import sys
import threading
import time

if 'nt' in os.name:
  import ctypes


DEFAULT_SECS_BETWEEN_ATTEMPTS = 10
POLL_MILLIS = 250
VERBOSE = True


class CommandFailedException(Exception):
  """Exception which gets raised when a command fails."""

  def __init__(self, output, *args):
    """Initialize the CommandFailedException.

    Args:
        output: string; output from the command.
    """
    Exception.__init__(self, *args)
    self._output = output

  @property
  def output(self):
    """Output from the command."""
    return self._output


class TimeoutException(CommandFailedException):
  """CommandFailedException which gets raised when a subprocess exceeds its
  timeout. """
  pass


def run_async(cmd, echo=None, shell=False):
  """ Run 'cmd' in a subprocess, returning a Popen class instance referring to
  that process.  (Non-blocking) """
  if echo is None:
    echo = VERBOSE
  if echo:
    print cmd
  if 'nt' in os.name:
    # Windows has a bad habit of opening a dialog when a console program
    # crashes, rather than just letting it crash.  Therefore, when a program
    # crashes on Windows, we don't find out until the build step times out.
    # This code prevents the dialog from appearing, so that we find out
    # immediately and don't waste time waiting around.
    SEM_NOGPFAULTERRORBOX = 0x0002
    ctypes.windll.kernel32.SetErrorMode(SEM_NOGPFAULTERRORBOX)
    flags = 0x8000000 # CREATE_NO_WINDOW
  else:
    flags = 0
  return subprocess.Popen(cmd, shell=shell, stderr=subprocess.STDOUT,
                          stdout=subprocess.PIPE, creationflags=flags,
                          bufsize=1)


class EnqueueThread(threading.Thread):
  """ Reads and enqueues lines from a file. """
  def __init__(self, file_obj, queue):
    threading.Thread.__init__(self)
    self._file = file_obj
    self._queue = queue
    self._stopped = False

  def run(self):
    if sys.platform.startswith('linux'):
      # Use a polling object to avoid the blocking call to readline().
      poll = select.poll()
      poll.register(self._file, select.POLLIN)
      while not self._stopped:
        has_output = poll.poll(POLL_MILLIS)
        if has_output:
          line = self._file.readline()
          if line == '':
            self._stopped = True
          self._queue.put(line)
    else:
      # Only Unix supports polling objects, so just read from the file,
      # Python-style.
      for line in iter(self._file.readline, ''):
        self._queue.put(line)
        if self._stopped:
          break

  def stop(self):
    self._stopped = True


def log_process_in_real_time(proc, echo=None, timeout=None, log_file=None,
                             halt_on_output=None, print_timestamps=True):
  """ Log the output of proc in real time until it completes. Return a tuple
  containing the exit code of proc and the contents of stdout.

  proc: an instance of Popen referring to a running subprocess.
  echo: boolean indicating whether to print the output received from proc.stdout
  timeout: number of seconds allotted for the process to run. Raises a
      TimeoutException if the run time exceeds the timeout.
  log_file: an open file for writing output
  halt_on_output: string; kill the process and return if this string is found
      in the output stream from the process.
  print_timestamps: boolean indicating whether a formatted timestamp should be
      prepended to each line of output.
  """
  if echo is None:
    echo = VERBOSE
  stdout_queue = Queue.Queue()
  log_thread = EnqueueThread(proc.stdout, stdout_queue)
  log_thread.start()
  try:
    all_output = []
    t_0 = time.time()
    while True:
      code = proc.poll()
      try:
        output = stdout_queue.get_nowait()
        all_output.append(output)
        if output and print_timestamps:
          timestamp = datetime.datetime.now().strftime('%H:%M:%S.%f')
          output = ''.join(['[%s] %s\n' % (timestamp, line)
                            for line in output.splitlines()])
        if echo:
          sys.stdout.write(output)
          sys.stdout.flush()
        if log_file:
          log_file.write(output)
          log_file.flush()
        if halt_on_output and halt_on_output in output:
          proc.terminate()
          break
      except Queue.Empty:
        if code != None: # proc has finished running
          break
        time.sleep(0.5)
      if timeout and time.time() - t_0 > timeout:
        proc.terminate()
        raise TimeoutException(
            ''.join(all_output),
            'Subprocess exceeded timeout of %ds' % timeout)
  finally:
    log_thread.stop()
    log_thread.join()
  return (code, ''.join(all_output))


def log_process_after_completion(proc, echo=None, timeout=None,
                                 log_file=None):
  """ Wait for proc to complete and return a tuple containing the exit code of
  proc and the contents of stdout. Unlike log_process_in_real_time, does not
  attempt to read stdout from proc in real time.

  proc: an instance of Popen referring to a running subprocess.
  echo: boolean indicating whether to print the output received from proc.stdout
  timeout: number of seconds allotted for the process to run. Raises a
      TimeoutException if the run time exceeds the timeout.
  log_file: an open file for writing outout
  """
  if echo is None:
    echo = VERBOSE
  t_0 = time.time()
  code = None
  while code is None:
    if timeout and time.time() - t_0 > timeout:
      raise TimeoutException(
          proc.communicate()[0],
          'Subprocess exceeded timeout of %ds' % timeout)
    time.sleep(0.5)
    code = proc.poll()
  output = proc.communicate()[0]
  if echo:
    print output
  if log_file:
    log_file.write(output)
    log_file.flush()
  return (code, output)


def run(cmd, echo=None, shell=False, timeout=None, print_timestamps=True,
        log_in_real_time=True):
  """ Run 'cmd' in a shell and return the combined contents of stdout and
  stderr (Blocking).  Throws an exception if the command exits non-zero.

  cmd: list of strings (or single string, iff shell==True) indicating the
      command to run
  echo: boolean indicating whether we should print the command and log output
  shell: boolean indicating whether we are using advanced shell features. Use
      only when absolutely necessary, since this allows a lot more freedom which
      could be exploited by malicious code. See the warning here:
      http://docs.python.org/library/subprocess.html#popen-constructor
  timeout: optional, integer indicating the maximum elapsed time in seconds
  print_timestamps: boolean indicating whether a formatted timestamp should be
      prepended to each line of output. Unused if echo or log_in_real_time is
      False.
  log_in_real_time: boolean indicating whether to read stdout from the
      subprocess in real time instead of when the process finishes. If echo is
      False, we never log in real time, even if log_in_real_time is True.
  """
  if echo is None:
    echo = VERBOSE
  proc = run_async(cmd, echo=echo, shell=shell)
  # If we're not printing the output, we don't care if the output shows up in
  # real time, so don't bother.
  if log_in_real_time and echo:
    (returncode, output) = log_process_in_real_time(proc, echo=echo,
        timeout=timeout, print_timestamps=print_timestamps)
  else:
    (returncode, output) = log_process_after_completion(proc, echo=echo,
                                                        timeout=timeout)
  if returncode != 0:
    raise CommandFailedException(
        output,
        'Command failed with code %d: %s' % (returncode, cmd))
  return output


def run_retry(cmd, echo=None, shell=False, attempts=1,
              secs_between_attempts=DEFAULT_SECS_BETWEEN_ATTEMPTS,
              timeout=None, print_timestamps=True):
  """ Wrapper for run() which makes multiple attempts until either the command
  succeeds or the maximum number of attempts is reached. """
  if echo is None:
    echo = VERBOSE
  attempt = 1
  while True:
    try:
      return run(cmd, echo=echo, shell=shell, timeout=timeout,
                 print_timestamps=print_timestamps)
    except CommandFailedException:
      if attempt >= attempts:
        raise
    print 'Command failed. Retrying in %d seconds...' % secs_between_attempts
    time.sleep(secs_between_attempts)
    attempt += 1