# Copyright (c) 2023 Egor Tensin <Egor.Tensin@gmail.com>
# This file is part of the "cimple" project.
# For details, see https://github.com/egor-tensin/cimple.
# Distributed under the MIT License.
from contextlib import contextmanager
import logging
import os
import shutil
import subprocess
from threading import Event, Lock, Thread
class LoggingEvent(Event):
def __init__(self, timeout=10):
self.timeout = timeout
super().__init__()
def log_line_matches(self, line):
return False
def wait(self):
if not super().wait(self.timeout):
raise RuntimeError('timed out while waiting for an event')
class LoggingThread(Thread):
def __init__(self, process, events=None):
self.process = process
self.events_lock = Lock()
if events is None:
events = []
self.events = events
super().__init__(target=lambda: self.process_output_lines())
self.start()
def add_event(self, event):
with self.events_lock:
self.events.append(event)
def process_output_lines(self):
for line in self.process.stdout:
line = line.removesuffix('\n')
logging.info('%s: %s', self.process.log_id, line)
with self.events_lock:
for event in self.events:
if event.is_set():
continue
if not event.log_line_matches(line):
continue
event.set()
self.events = [event for event in self.events if not event.is_set()]
class CmdLine:
@staticmethod
def which(binary):
if os.path.split(binary)[0]:
# shutil.which('bin/bash') doesn't work.
return os.path.abspath(binary)
path = shutil.which(binary)
if path is None:
raise RuntimeError("couldn't find a binary: " + binary)
return path
@staticmethod
def unbuffered():
return CmdLine('stdbuf', '-o0')
def __init__(self, binary, *args, name=None):
binary = self.which(binary)
argv = [binary] + list(args)
self.binary = binary
self.argv = argv
if name is None:
name = os.path.basename(binary)
self.process_name = name
def log_line_means_process_ready(self, line):
return True
@classmethod
def wrap(cls, outer, inner):
return cls(outer.argv[0], *outer.argv[1:], *inner.argv, name=inner.process_name)
def run(self, *argv):
return Process.run(*self.argv, *argv)
@contextmanager
def run_async(self, *argv):
with Process(self, *argv) as process:
yield process
class LoggingEventProcessReady(LoggingEvent):
def __init__(self, process):
self.process = process
super().__init__()
def set(self):
logging.info('Process %s is ready', self.process.log_id)
super().set()
def log_line_matches(self, line):
return self.process.cmd_line.log_line_means_process_ready(line)
class Process(subprocess.Popen):
_COMMON_ARGS = {
'text': True,
'stdin': subprocess.DEVNULL,
'stdout': subprocess.PIPE,
'stderr': subprocess.STDOUT,
}
@staticmethod
def _log_process_start(argv):
logging.info('Executing command: %s', argv)
@staticmethod
def _log_process_end(argv, ec, output):
log = logging.info
if ec:
log = logging.error
if ec:
log('Command %s exited with code %s', argv, ec)
else:
log('Command %s completed successfully', argv)
if output:
log('Output:\n%s', output)
@staticmethod
def run(*args, **kwargs):
argv = list(args)
Process._log_process_start(argv)
try:
result = subprocess.run(argv, check=True, **Process._COMMON_ARGS, **kwargs)
ec, output = result.returncode, result.stdout
Process._log_process_end(argv, ec, output)
return output
except subprocess.CalledProcessError as e:
ec, output = e.returncode, e.stdout
Process._log_process_end(argv, ec, output)
raise
def __init__(self, cmd_line, *args):
self.cmd_line = cmd_line
argv = cmd_line.argv + list(args)
self._log_process_start(argv)
super().__init__(argv, **Process._COMMON_ARGS)
logging.info('Process %s has started', self.log_id)
ready_event = LoggingEventProcessReady(self)
self.logger = LoggingThread(self, [ready_event])
ready_event.wait()
@property
def log_id(self):
return f'{self.pid}/{self.cmd_line.process_name}'
def __exit__(self, *args):
try:
self.shut_down()
self.logger.join()
except Exception as e:
logging.exception(e)
# Postpone closing the pipes until after the logging thread is finished
# so that it doesn't attempt to read from closed descriptors.
super().__exit__(*args)
SHUT_DOWN_TIMEOUT_SEC = 3
def shut_down(self):
ec = self.poll()
if ec is not None:
return
logging.info('Terminating process %s', self.log_id)
self.terminate()
try:
self.wait(timeout=Process.SHUT_DOWN_TIMEOUT_SEC)
return
except subprocess.TimeoutExpired:
pass
logging.info('Process %s failed to terminate in time, killing it', self.log_id)
self.kill()
self.wait(timeout=Process.SHUT_DOWN_TIMEOUT_SEC)
