diff --git a/resallocserver/manager.py b/resallocserver/manager.py index e3a1d16..53020a2 100644 --- a/resallocserver/manager.py +++ b/resallocserver/manager.py @@ -22,6 +22,7 @@ import os import errno import time +import select import threading import subprocess import warnings @@ -57,7 +58,8 @@ def command_env(pool_id=None, res_id=None, res_name=None, def run_command(pool_id, res_id, res_name, id_in_pool, named_counters, command, ltype='alloc', catch_stdout_bytes=None, data=None, - catch_stdout_lines_securely=False, timeout=None): + catch_stdout_lines_securely=False, timeout=None, + discard_partial_line=True): """ Run command, and log into according directory (per.app.config). If catch_stdout_bytes is specified, we read & log continuously. @@ -66,13 +68,15 @@ def run_command(pool_id, res_id, res_name, id_in_pool, named_counters, command, logdir = os.path.join(config['logdir'], 'hooks') return _run_command(app.log, logdir, pool_id, res_id, res_name, id_in_pool, named_counters, command, ltype, catch_stdout_bytes, - data, catch_stdout_lines_securely, timeout) + data, catch_stdout_lines_securely, timeout, + discard_partial_line) def _run_command(log, logdir, pool_id, res_id, res_name, id_in_pool, named_counters, command, ltype='alloc', catch_stdout_bytes=None, data=None, - catch_stdout_lines_securely=False, timeout=None): + catch_stdout_lines_securely=False, timeout=None, + discard_partial_line=True): """ Internal variant for _run_command() that is easier to test locally like: In [1] import logging @@ -81,10 +85,7 @@ def _run_command(log, logdir, pool_id, res_id, res_name, id_in_pool, started ; sleep 100", timeout=2) Out[3]: {'status': 124} """ - - # we do not support timeout when catching lines - assert not (timeout and catch_stdout_lines_securely) - + # pylint: disable=too-many-statements log.debug("running: " + command) env = command_env(pool_id, res_id, res_name, id_in_pool, named_counters, data) @@ -115,29 +116,78 @@ def _run_command(log, logdir, pool_id, res_id, res_name, id_in_pool, with subprocess.Popen(command, env=env, shell=True, stdout=subprocess.PIPE, stderr=logfile) as sp: captured_string = b"" + timeout = timeout or 60 * 60 * 24 + deadline = time.monotonic() + timeout + buffer = b"" - for line in iter(sp.stdout.readline, b''): - # Write to the log. - logfile.write(line) - logfile.flush() + while True: + # How many seconds are left before we want to timeout + remaining = deadline - time.monotonic() + if remaining <= 0: + logfile.write(f"\n".encode()) + sp.kill() + return {'status': 124, 'stdout': captured_string} + + # Either we get a readable file descriptors back or we get nothing + # because the `select` timeouted on `read` + rlist = [sp.stdout.fileno()] + ready, _, _ = select.select(rlist, [], [], remaining) + if not ready: + logfile.write(f"\n".encode()) + sp.kill() + return {'status': 124, 'stdout': captured_string} + + # To support multiple input streams, we should iterate over + # `ready` instead of reading stdout directly. But there also + # need to be some additional changes like separate EOF handling + # for every input stream + chunk = os.read(sp.stdout.fileno(), 4096) + eof = chunk == b'' + + # We successfully read a chunk of data + if not eof: + buffer += chunk + # We encountered EOF but we still have unprocessed bytes + elif buffer: + # Discard the partial line or treat it as a complete line? + if discard_partial_line: + buffer = b'' + else: + buffer += b'\n' + # We encountered EOF and we already processed everything + else: + break - if stdout_stopped: - continue + # Process all complete lines from buffer + while b'\n' in buffer: + line, buffer = buffer.split(b'\n', 1) + line += b'\n' - if stdout_written + len(line) > catch_stdout_bytes: - if stdout_written == 0 and not catch_stdout_lines_securely: - # Even the first line is too long for this buffer. Catch at - # least part of it. - line = line[:catch_stdout_bytes] - captured_string += line + # Write to the log. + logfile.write(line) + logfile.flush() - stdout_stopped = True - if not catch_stdout_lines_securely: - captured_string += b"<< trimmed >>\n" - continue + if stdout_stopped: + continue + + if stdout_written + len(line) > catch_stdout_bytes: + if stdout_written == 0 and not catch_stdout_lines_securely: + # Even the first line is too long for this buffer. + # Catch at # least part of it. + line = line[:catch_stdout_bytes] + captured_string += line + + stdout_stopped = True + if not catch_stdout_lines_securely: + captured_string += b"<< trimmed >>\n" + continue + + stdout_written += len(line) + captured_string += line - stdout_written += len(line) - captured_string += line + # This is to handle EOF after processing a partial line + if eof: + break return { 'status': sp.wait(), @@ -436,6 +486,7 @@ def _list_all_resources(self): ltype="list", catch_stdout_bytes=5120, catch_stdout_lines_securely=True, + timeout=120, ) return result["stdout"].decode("utf-8").strip().split()