xref: /aosp_15_r20/external/autotest/client/cros/audio/cmd_utils.py (revision 9c5db1993ded3edbeafc8092d69fe5de2ee02df7)
1*9c5db199SXin Li# Lint as: python2, python3
2*9c5db199SXin Li# Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
3*9c5db199SXin Li# Use of this source code is governed by a BSD-style license that can be
4*9c5db199SXin Li# found in the LICENSE file.
5*9c5db199SXin Li
6*9c5db199SXin Liimport atexit
7*9c5db199SXin Liimport itertools
8*9c5db199SXin Liimport logging
9*9c5db199SXin Liimport os
10*9c5db199SXin Liimport pipes
11*9c5db199SXin Liimport pwd
12*9c5db199SXin Liimport select
13*9c5db199SXin Liimport subprocess
14*9c5db199SXin Liimport threading
15*9c5db199SXin Li
16*9c5db199SXin Lifrom autotest_lib.client.common_lib.utils import TEE_TO_LOGS
17*9c5db199SXin Li
18*9c5db199SXin Li_popen_lock = threading.Lock()
19*9c5db199SXin Li_logging_service = None
20*9c5db199SXin Li_command_serial_number = itertools.count(1)
21*9c5db199SXin Li
22*9c5db199SXin Li_LOG_BUFSIZE = 4096
23*9c5db199SXin Li_PIPE_CLOSED = -1
24*9c5db199SXin Li
25*9c5db199SXin Liclass _LoggerProxy(object):
26*9c5db199SXin Li
27*9c5db199SXin Li    def __init__(self, logger):
28*9c5db199SXin Li        self._logger = logger
29*9c5db199SXin Li
30*9c5db199SXin Li    def fileno(self):
31*9c5db199SXin Li        """Returns the fileno of the logger pipe."""
32*9c5db199SXin Li        return self._logger._pipe[1]
33*9c5db199SXin Li
34*9c5db199SXin Li    def __del__(self):
35*9c5db199SXin Li        self._logger.close()
36*9c5db199SXin Li
37*9c5db199SXin Li
38*9c5db199SXin Liclass _PipeLogger(object):
39*9c5db199SXin Li
40*9c5db199SXin Li    def __init__(self, level, prefix):
41*9c5db199SXin Li        self._pipe = list(os.pipe())
42*9c5db199SXin Li        self._level = level
43*9c5db199SXin Li        self._prefix = prefix
44*9c5db199SXin Li
45*9c5db199SXin Li    def close(self):
46*9c5db199SXin Li        """Closes the logger."""
47*9c5db199SXin Li        if self._pipe[1] != _PIPE_CLOSED:
48*9c5db199SXin Li            os.close(self._pipe[1])
49*9c5db199SXin Li            self._pipe[1] = _PIPE_CLOSED
50*9c5db199SXin Li
51*9c5db199SXin Li
52*9c5db199SXin Liclass _LoggingService(object):
53*9c5db199SXin Li
54*9c5db199SXin Li    def __init__(self):
55*9c5db199SXin Li        # Python's list is thread safe
56*9c5db199SXin Li        self._loggers = []
57*9c5db199SXin Li
58*9c5db199SXin Li        # Change tuple to list so that we can change the value when
59*9c5db199SXin Li        # closing the pipe.
60*9c5db199SXin Li        self._pipe = list(os.pipe())
61*9c5db199SXin Li        self._thread = threading.Thread(target=self._service_run)
62*9c5db199SXin Li        self._thread.daemon = True
63*9c5db199SXin Li        self._thread.start()
64*9c5db199SXin Li
65*9c5db199SXin Li
66*9c5db199SXin Li    def _service_run(self):
67*9c5db199SXin Li        terminate_loop = False
68*9c5db199SXin Li        while not terminate_loop:
69*9c5db199SXin Li            rlist = [l._pipe[0] for l in self._loggers]
70*9c5db199SXin Li            rlist.append(self._pipe[0])
71*9c5db199SXin Li            for r in select.select(rlist, [], [])[0]:
72*9c5db199SXin Li                data = os.read(r, _LOG_BUFSIZE)
73*9c5db199SXin Li                if r != self._pipe[0]:
74*9c5db199SXin Li                    self._output_logger_message(r, data)
75*9c5db199SXin Li                elif len(data) == 0:
76*9c5db199SXin Li                    terminate_loop = True
77*9c5db199SXin Li        # Release resources.
78*9c5db199SXin Li        os.close(self._pipe[0])
79*9c5db199SXin Li        for logger in self._loggers:
80*9c5db199SXin Li            os.close(logger._pipe[0])
81*9c5db199SXin Li
82*9c5db199SXin Li
83*9c5db199SXin Li    def _output_logger_message(self, r, data):
84*9c5db199SXin Li        logger = next(l for l in self._loggers if l._pipe[0] == r)
85*9c5db199SXin Li
86*9c5db199SXin Li        if len(data) == 0:
87*9c5db199SXin Li            os.close(logger._pipe[0])
88*9c5db199SXin Li            self._loggers.remove(logger)
89*9c5db199SXin Li            return
90*9c5db199SXin Li
91*9c5db199SXin Li        for line in data.split(b'\n'):
92*9c5db199SXin Li            logging.log(logger._level, '%s%s', logger._prefix, line)
93*9c5db199SXin Li
94*9c5db199SXin Li
95*9c5db199SXin Li    def create_logger(self, level=logging.DEBUG, prefix=''):
96*9c5db199SXin Li        """Creates a new logger.
97*9c5db199SXin Li
98*9c5db199SXin Li        @param level: the desired logging level
99*9c5db199SXin Li        @param prefix: the prefix to add to each log entry
100*9c5db199SXin Li        """
101*9c5db199SXin Li        logger = _PipeLogger(level=level, prefix=prefix)
102*9c5db199SXin Li        self._loggers.append(logger)
103*9c5db199SXin Li        os.write(self._pipe[1], b'\0')
104*9c5db199SXin Li        return _LoggerProxy(logger)
105*9c5db199SXin Li
106*9c5db199SXin Li
107*9c5db199SXin Li    def shutdown(self):
108*9c5db199SXin Li        """Shuts down the logger."""
109*9c5db199SXin Li        if self._pipe[1] != _PIPE_CLOSED:
110*9c5db199SXin Li            os.close(self._pipe[1])
111*9c5db199SXin Li            self._pipe[1] = _PIPE_CLOSED
112*9c5db199SXin Li            self._thread.join()
113*9c5db199SXin Li
114*9c5db199SXin Li
115*9c5db199SXin Lidef create_logger(level=logging.DEBUG, prefix=''):
116*9c5db199SXin Li    """Creates a new logger.
117*9c5db199SXin Li
118*9c5db199SXin Li    @param level: the desired logging level
119*9c5db199SXin Li    @param prefix: the prefix to add to each log entry
120*9c5db199SXin Li    """
121*9c5db199SXin Li    global _logging_service
122*9c5db199SXin Li    if _logging_service is None:
123*9c5db199SXin Li        _logging_service = _LoggingService()
124*9c5db199SXin Li        atexit.register(_logging_service.shutdown)
125*9c5db199SXin Li    return _logging_service.create_logger(level=level, prefix=prefix)
126*9c5db199SXin Li
127*9c5db199SXin Li
128*9c5db199SXin Lidef kill_or_log_returncode(*popens):
129*9c5db199SXin Li    """Kills all the processes of the given Popens or logs the return code.
130*9c5db199SXin Li
131*9c5db199SXin Li    @param popens: The Popens to be killed.
132*9c5db199SXin Li    """
133*9c5db199SXin Li    for p in popens:
134*9c5db199SXin Li        if p.poll() is None:
135*9c5db199SXin Li            try:
136*9c5db199SXin Li                p.kill()
137*9c5db199SXin Li            except Exception as e:
138*9c5db199SXin Li                logging.warning('failed to kill %d, %s', p.pid, e)
139*9c5db199SXin Li        else:
140*9c5db199SXin Li            logging.warning('command exit (pid=%d, rc=%d): %s',
141*9c5db199SXin Li                            p.pid, p.returncode, p.command)
142*9c5db199SXin Li
143*9c5db199SXin Li
144*9c5db199SXin Lidef wait_and_check_returncode(*popens):
145*9c5db199SXin Li    """Wait for all the Popens and check the return code is 0.
146*9c5db199SXin Li
147*9c5db199SXin Li    If the return code is not 0, it raises an RuntimeError.
148*9c5db199SXin Li
149*9c5db199SXin Li    @param popens: The Popens to be checked.
150*9c5db199SXin Li    """
151*9c5db199SXin Li    error_message = None
152*9c5db199SXin Li    for p in popens:
153*9c5db199SXin Li        if p.wait() != 0:
154*9c5db199SXin Li            error_message = ('Command failed(%d, %d): %s' %
155*9c5db199SXin Li                             (p.pid, p.returncode, p.command))
156*9c5db199SXin Li            logging.error(error_message)
157*9c5db199SXin Li    if error_message:
158*9c5db199SXin Li        raise RuntimeError(error_message)
159*9c5db199SXin Li
160*9c5db199SXin Li
161*9c5db199SXin Lidef execute(args, stdin=None, stdout=TEE_TO_LOGS, stderr=TEE_TO_LOGS,
162*9c5db199SXin Li            run_as=None):
163*9c5db199SXin Li    """Executes a child command and wait for it.
164*9c5db199SXin Li
165*9c5db199SXin Li    Returns the output from standard output if 'stdout' is subprocess.PIPE.
166*9c5db199SXin Li    Raises RuntimeException if the return code of the child command is not 0.
167*9c5db199SXin Li
168*9c5db199SXin Li    @param args: the command to be executed
169*9c5db199SXin Li    @param stdin: the executed program's standard input
170*9c5db199SXin Li    @param stdout: the executed program's standard output
171*9c5db199SXin Li    @param stderr: the executed program's standard error
172*9c5db199SXin Li    @param run_as: if not None, run the command as the given user
173*9c5db199SXin Li    """
174*9c5db199SXin Li    ps = popen(args, stdin=stdin, stdout=stdout, stderr=stderr,
175*9c5db199SXin Li               run_as=run_as)
176*9c5db199SXin Li    out = ps.communicate()[0] if stdout == subprocess.PIPE else None
177*9c5db199SXin Li    wait_and_check_returncode(ps)
178*9c5db199SXin Li    return out
179*9c5db199SXin Li
180*9c5db199SXin Li
181*9c5db199SXin Lidef _run_as(user):
182*9c5db199SXin Li    """Changes the uid and gid of the running process to be that of the
183*9c5db199SXin Li    given user and configures its supplementary groups.
184*9c5db199SXin Li
185*9c5db199SXin Li    Don't call this function directly, instead wrap it in a lambda and
186*9c5db199SXin Li    pass that to the preexec_fn argument of subprocess.Popen.
187*9c5db199SXin Li
188*9c5db199SXin Li    Example usage:
189*9c5db199SXin Li    subprocess.Popen(..., preexec_fn=lambda: _run_as('chronos'))
190*9c5db199SXin Li
191*9c5db199SXin Li    @param user: the user to run as
192*9c5db199SXin Li    """
193*9c5db199SXin Li    pw = pwd.getpwnam(user)
194*9c5db199SXin Li    os.setgid(pw.pw_gid)
195*9c5db199SXin Li    os.initgroups(user, pw.pw_gid)
196*9c5db199SXin Li    os.setuid(pw.pw_uid)
197*9c5db199SXin Li
198*9c5db199SXin Li
199*9c5db199SXin Lidef popen(args, stdin=None, stdout=TEE_TO_LOGS, stderr=TEE_TO_LOGS, env=None,
200*9c5db199SXin Li          run_as=None):
201*9c5db199SXin Li    """Returns a Popen object just as subprocess.Popen does but with the
202*9c5db199SXin Li    executed command stored in Popen.command.
203*9c5db199SXin Li
204*9c5db199SXin Li    @param args: the command to be executed
205*9c5db199SXin Li    @param stdin: the executed program's standard input
206*9c5db199SXin Li    @param stdout: the executed program's standard output
207*9c5db199SXin Li    @param stderr: the executed program's standard error
208*9c5db199SXin Li    @param env: the executed program's environment
209*9c5db199SXin Li    @param run_as: if not None, run the command as the given user
210*9c5db199SXin Li    """
211*9c5db199SXin Li    command_id = next(_command_serial_number)
212*9c5db199SXin Li    prefix = '[%04d] ' % command_id
213*9c5db199SXin Li
214*9c5db199SXin Li    if stdout is TEE_TO_LOGS:
215*9c5db199SXin Li        stdout = create_logger(level=logging.DEBUG, prefix=prefix)
216*9c5db199SXin Li    if stderr is TEE_TO_LOGS:
217*9c5db199SXin Li        stderr = create_logger(level=logging.ERROR, prefix=prefix)
218*9c5db199SXin Li
219*9c5db199SXin Li    command = ' '.join(pipes.quote(x) for x in args)
220*9c5db199SXin Li    logging.info('%sRunning: %s', prefix, command)
221*9c5db199SXin Li
222*9c5db199SXin Li    preexec_fn = None
223*9c5db199SXin Li    if run_as is not None:
224*9c5db199SXin Li        preexec_fn = lambda: _run_as(run_as)
225*9c5db199SXin Li
226*9c5db199SXin Li    # The lock is required for http://crbug.com/323843.
227*9c5db199SXin Li    with _popen_lock:
228*9c5db199SXin Li        ps = subprocess.Popen(args, stdin=stdin, stdout=stdout, stderr=stderr,
229*9c5db199SXin Li                              env=env, preexec_fn=preexec_fn)
230*9c5db199SXin Li    logging.info('%spid is %d', prefix, ps.pid)
231*9c5db199SXin Li    ps.command_id = command_id
232*9c5db199SXin Li    ps.command = command
233*9c5db199SXin Li    return ps
234