client/cros/audio/cmd_utils.py

#!/usr/bin/python2
# Copyright (c) 2013 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.

import atexit
import itertools
import logging
import os
import pipes
import pwd
import select
import subprocess
import threading

from autotest_lib.client.common_lib.utils import TEE_TO_LOGS

_popen_lock = threading.Lock()
_logging_service = None
_command_serial_number = itertools.count(1)

_LOG_BUFSIZE = 4096
_PIPE_CLOSED = -1

class _LoggerProxy(object):

    def __init__(self, logger):
        self._logger = logger

    def fileno(self):
        """Returns the fileno of the logger pipe."""
        return self._logger._pipe[1]

    def __del__(self):
        self._logger.close()


class _PipeLogger(object):

    def __init__(self, level, prefix):
        self._pipe = list(os.pipe())
        self._level = level
        self._prefix = prefix

    def close(self):
        """Closes the logger."""
        if self._pipe[1] != _PIPE_CLOSED:
            os.close(self._pipe[1])
            self._pipe[1] = _PIPE_CLOSED


class _LoggingService(object):

    def __init__(self):
        # Python's list is thread safe
        self._loggers = []

        # Change tuple to list so that we can change the value when
        # closing the pipe.
        self._pipe = list(os.pipe())
        self._thread = threading.Thread(target=self._service_run)
        self._thread.daemon = True
        self._thread.start()


    def _service_run(self):
        terminate_loop = False
        while not terminate_loop:
            rlist = [l._pipe[0] for l in self._loggers]
            rlist.append(self._pipe[0])
            for r in select.select(rlist, [], [])[0]:
                data = os.read(r, _LOG_BUFSIZE)
                if r != self._pipe[0]:
                    self._output_logger_message(r, data)
                elif len(data) == 0:
                    terminate_loop = True
        # Release resources.
        os.close(self._pipe[0])
        for logger in self._loggers:
            os.close(logger._pipe[0])


    def _output_logger_message(self, r, data):
        logger = next(l for l in self._loggers if l._pipe[0] == r)

        if len(data) == 0:
            os.close(logger._pipe[0])
            self._loggers.remove(logger)
            return

        for line in data.split('\n'):
            logging.log(logger._level, '%s%s', logger._prefix, line)


    def create_logger(self, level=logging.DEBUG, prefix=''):
        """Creates a new logger.

        @param level: the desired logging level
        @param prefix: the prefix to add to each log entry
        """
        logger = _PipeLogger(level=level, prefix=prefix)
        self._loggers.append(logger)
        os.write(self._pipe[1], '\0')
        return _LoggerProxy(logger)


    def shutdown(self):
        """Shuts down the logger."""
        if self._pipe[1] != _PIPE_CLOSED:
            os.close(self._pipe[1])
            self._pipe[1] = _PIPE_CLOSED
            self._thread.join()


def create_logger(level=logging.DEBUG, prefix=''):
    """Creates a new logger.

    @param level: the desired logging level
    @param prefix: the prefix to add to each log entry
    """
    global _logging_service
    if _logging_service is None:
        _logging_service = _LoggingService()
        atexit.register(_logging_service.shutdown)
    return _logging_service.create_logger(level=level, prefix=prefix)


def kill_or_log_returncode(*popens):
    """Kills all the processes of the given Popens or logs the return code.

    @param popens: The Popens to be killed.
    """
    for p in popens:
        if p.poll() is None:
            try:
                p.kill()
            except Exception as e:
                logging.warning('failed to kill %d, %s', p.pid, e)
        else:
            logging.warning('command exit (pid=%d, rc=%d): %s',
                            p.pid, p.returncode, p.command)


def wait_and_check_returncode(*popens):
    """Wait for all the Popens and check the return code is 0.

    If the return code is not 0, it raises an RuntimeError.

    @param popens: The Popens to be checked.
    """
    error_message = None
    for p in popens:
        if p.wait() != 0:
            error_message = ('Command failed(%d, %d): %s' %
                             (p.pid, p.returncode, p.command))
            logging.error(error_message)
    if error_message:
        raise RuntimeError(error_message)


def execute(args, stdin=None, stdout=TEE_TO_LOGS, stderr=TEE_TO_LOGS,
            run_as=None):
    """Executes a child command and wait for it.

    Returns the output from standard output if 'stdout' is subprocess.PIPE.
    Raises RuntimeException if the return code of the child command is not 0.

    @param args: the command to be executed
    @param stdin: the executed program's standard input
    @param stdout: the executed program's standard output
    @param stderr: the executed program's standard error
    @param run_as: if not None, run the command as the given user
    """
    ps = popen(args, stdin=stdin, stdout=stdout, stderr=stderr,
               run_as=run_as)
    out = ps.communicate()[0] if stdout == subprocess.PIPE else None
    wait_and_check_returncode(ps)
    return out


def _run_as(user):
    """Changes the uid and gid of the running process to be that of the
    given user and configures its supplementary groups.

    Don't call this function directly, instead wrap it in a lambda and
    pass that to the preexec_fn argument of subprocess.Popen.

    Example usage:
    subprocess.Popen(..., preexec_fn=lambda: _run_as('chronos'))

    @param user: the user to run as
    """
    pw = pwd.getpwnam(user)
    os.setgid(pw.pw_gid)
    os.initgroups(user, pw.pw_gid)
    os.setuid(pw.pw_uid)


def popen(args, stdin=None, stdout=TEE_TO_LOGS, stderr=TEE_TO_LOGS, env=None,
          run_as=None):
    """Returns a Popen object just as subprocess.Popen does but with the
    executed command stored in Popen.command.

    @param args: the command to be executed
    @param stdin: the executed program's standard input
    @param stdout: the executed program's standard output
    @param stderr: the executed program's standard error
    @param env: the executed program's environment
    @param run_as: if not None, run the command as the given user
    """
    command_id = next(_command_serial_number)
    prefix = '[%04d] ' % command_id

    if stdout is TEE_TO_LOGS:
        stdout = create_logger(level=logging.DEBUG, prefix=prefix)
    if stderr is TEE_TO_LOGS:
        stderr = create_logger(level=logging.ERROR, prefix=prefix)

    command = ' '.join(pipes.quote(x) for x in args)
    logging.info('%sRunning: %s', prefix, command)

    preexec_fn = None
    if run_as is not None:
        preexec_fn = lambda: _run_as(run_as)

    # The lock is required for http://crbug.com/323843.
    with _popen_lock:
        ps = subprocess.Popen(args, stdin=stdin, stdout=stdout, stderr=stderr,
                              env=env, preexec_fn=preexec_fn)
    logging.info('%spid is %d', prefix, ps.pid)
    ps.command_id = command_id
    ps.command = command
    return ps