234 lines
7.0 KiB
Python
234 lines
7.0 KiB
Python
#!/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
|