1#!/usr/bin/env python3
2#
3#   Copyright 2020 - The Android Open Source Project
4#
5#   Licensed under the Apache License, Version 2.0 (the "License");
6#   you may not use this file except in compliance with the License.
7#   You may obtain a copy of the License at
8#
9#       http://www.apache.org/licenses/LICENSE-2.0
10#
11#   Unless required by applicable law or agreed to in writing, software
12#   distributed under the License is distributed on an "AS IS" BASIS,
13#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14#   See the License for the specific language governing permissions and
15#   limitations under the License.
16
17from contextlib import ExitStack
18import concurrent.futures
19import logging
20import re
21import subprocess
22from cert.os_utils import TerminalColor
23
24
25class AsyncSubprocessLogger:
26    """
27    An asynchronous logger for subprocesses.Popen object's STDOUT
28
29    Contains threading functionality that allows asynchronous handling of lines
30    from STDOUT from subprocess.Popen
31    """
32    WAIT_TIMEOUT_SECONDS = 10
33    PROCESS_TAG_MIN_WIDTH = 24
34
35    def __init__(self,
36                 process: subprocess.Popen,
37                 log_file_paths,
38                 log_to_stdout=False,
39                 tag=None,
40                 color: TerminalColor = None):
41        """
42        :param process: a subprocess.Popen object with STDOUT
43        :param log_file_paths: list of log files to redirect log to
44        :param log_to_stdout: whether to dump logs to stdout in the format of
45                              "[tag] logline"
46        :param tag: tag to be used in above format
47        :param color: when dumping to stdout, what color to use for tag
48        """
49        if not process:
50            raise ValueError("process cannot be None")
51        if not process.stdout:
52            raise ValueError("process.stdout cannot be None")
53        if log_to_stdout:
54            if not tag or type(tag) is not str:
55                raise ValueError("When logging to stdout, log tag must be set")
56        self.log_file_paths = log_file_paths
57        self.log_to_stdout = log_to_stdout
58        self.tag = tag
59        self.color = color
60        self.process = process
61        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
62        self.future = self.executor.submit(self.__logging_loop)
63
64    def stop(self):
65        """
66        Stop this logger and this object can no longer be used after this call
67        """
68        try:
69            result = self.future.result(timeout=self.WAIT_TIMEOUT_SECONDS)
70            if result:
71                logging.error("logging thread %s produced an error when executing: %s" % (self.tag, str(result)))
72        except concurrent.futures.TimeoutError:
73            logging.error("logging thread %s failed to finish after %d seconds" % (self.tag, self.WAIT_TIMEOUT_SECONDS))
74        self.executor.shutdown(wait=False)
75
76    def __logging_loop(self):
77        if self.color:
78            loggableTag = "[%s%s%s]" % (self.color, self.tag, TerminalColor.END)
79        else:
80            loggableTag = "[%s]" % self.tag
81        tagLength = len(re.sub('[^\w\s]', '', loggableTag))
82        if tagLength < self.PROCESS_TAG_MIN_WIDTH:
83            loggableTag += " " * (self.PROCESS_TAG_MIN_WIDTH - tagLength)
84        with ExitStack() as stack:
85            log_files = [stack.enter_context(open(file_path, 'w')) for file_path in self.log_file_paths]
86            for line in self.process.stdout:
87                for log_file in log_files:
88                    log_file.write(line)
89                if self.log_to_stdout:
90                    print("{}{}".format(loggableTag, line.strip()))
91