diff --git a/resallocserver/helpers.py b/resallocserver/helpers.py new file mode 100644 index 0000000..96f8191 --- /dev/null +++ b/resallocserver/helpers.py @@ -0,0 +1,54 @@ +""" +Helper utilities for the resalloc server. +""" + +import os +import select +import time + + +class LineBufferOverflow(Exception): + """Raised when a line exceeds the maximum allowed length.""" + + +def yield_lines_from_fds(fds, timeout=None, max_line_length=None): + """ + Generator that reads lines from file descriptors, yielding (fd, line) pairs. + Raises TimeoutError when the deadline is exceeded. + Raises LineBufferOverflow when a line exceeds max_line_length. + """ + timeout = timeout or 60 * 60 * 24 + deadline = time.monotonic() + timeout + buffers = {fd: b"" for fd in fds} + active_fds = set(fds) + + while active_fds: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"timeouted after {timeout}s") + + ready, _, _ = select.select(list(active_fds), [], [], remaining) + if not ready: + raise TimeoutError(f"timeouted after {timeout}s (on read)") + + for fd in ready: + if max_line_length: + capacity = max_line_length - len(buffers[fd]) + if capacity <= 0: + raise LineBufferOverflow( + f"line on fd {fd} exceeded {max_line_length} bytes") + chunk = os.read(fd, min(4096, capacity)) + else: + chunk = os.read(fd, 4096) + + if chunk: + buffers[fd] += chunk + while b'\n' in buffers[fd]: + line, buffers[fd] = buffers[fd].split(b'\n', 1) + yield (fd, line + b'\n') + continue + + # EOF + if buffers[fd]: + yield (fd, buffers[fd]) + active_fds.discard(fd) diff --git a/resallocserver/manager.py b/resallocserver/manager.py index c45235b..1a82320 100644 --- a/resallocserver/manager.py +++ b/resallocserver/manager.py @@ -22,15 +22,16 @@ import os import errno import time -import select import threading import subprocess import warnings from datetime import datetime + from resalloc import helpers from resalloc.helpers import RState from resallocserver import models from resallocserver.app import session_scope, app +from resallocserver.helpers import LineBufferOverflow, yield_lines_from_fds from resallocserver.logic import ( QResources, QTickets, assign_ticket, release_resource ) @@ -58,8 +59,7 @@ 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, - discard_partial_line=True): + catch_stdout_lines_securely=False, timeout=None): """ Run command, and log into according directory (per.app.config). If catch_stdout_bytes is specified, we read & log continuously. @@ -68,15 +68,14 @@ 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, - discard_partial_line) + data, catch_stdout_lines_securely, timeout) + 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, - discard_partial_line=True): + catch_stdout_lines_securely=False, timeout=None): """ Internal variant for _run_command() that is easier to test locally like: In [1] import logging @@ -85,7 +84,6 @@ def _run_command(log, logdir, pool_id, res_id, res_name, id_in_pool, started ; sleep 100", timeout=2) Out[3]: {'status': 124} """ - # pylint: disable=too-many-statements log.debug("running: " + command) env = command_env(pool_id, res_id, res_name, id_in_pool, named_counters, data) @@ -105,65 +103,21 @@ def _run_command(log, logdir, pool_id, res_id, res_name, id_in_pool, timeout=timeout)} except subprocess.TimeoutExpired: logfile.write(f"\n".encode()) - # default /bin/timeout utility exit status return {'status': 124} stdout_written = 0 stdout_stopped = False - # Run the sub-command to be captured. - 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"" - - 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 - - # Process all complete lines from buffer - while b'\n' in buffer: - line, buffer = buffer.split(b'\n', 1) - line += b'\n' - # Write to the log. + try: + for _, line in yield_lines_from_fds( + [sp.stdout.fileno()], + timeout=timeout, + max_line_length=10240 if catch_stdout_lines_securely else None, + ): logfile.write(line) logfile.flush() @@ -172,8 +126,6 @@ def _run_command(log, logdir, pool_id, res_id, res_name, id_in_pool, 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 @@ -185,9 +137,10 @@ def _run_command(log, logdir, pool_id, res_id, res_name, id_in_pool, stdout_written += len(line) captured_string += line - # This is to handle EOF after processing a partial line - if eof: - break + except (TimeoutError, LineBufferOverflow) as e: + logfile.write(f"<{e}>\n".encode()) + sp.kill() + return {'status': 124, 'stdout': captured_string} return { 'status': sp.wait(), diff --git a/tests/test_yield_lines_from_fds.py b/tests/test_yield_lines_from_fds.py new file mode 100644 index 0000000..daeefef --- /dev/null +++ b/tests/test_yield_lines_from_fds.py @@ -0,0 +1,163 @@ +""" +Tests for yield_lines_from_fds() +""" + +import os +import subprocess +import time + +import pytest + +from resallocserver.helpers import yield_lines_from_fds, LineBufferOverflow + + +def _pipe_with_data(data): + """Create a pipe, write data to it, close the write end, return read fd.""" + r, w = os.pipe() + os.write(w, data) + os.close(w) + return r + + +class TestYieldLinesFromFds: + + def test_single_line(self): + fd = _pipe_with_data(b"hello\n") + result = list(yield_lines_from_fds([fd])) + assert result == [(fd, b"hello\n")] + + def test_multiple_lines(self): + fd = _pipe_with_data(b"one\ntwo\nthree\n") + result = list(yield_lines_from_fds([fd])) + assert result == [ + (fd, b"one\n"), + (fd, b"two\n"), + (fd, b"three\n"), + ] + + def test_partial_line_at_eof(self): + fd = _pipe_with_data(b"no newline") + result = list(yield_lines_from_fds([fd])) + assert result == [(fd, b"no newline")] + + def test_partial_line_after_complete_lines(self): + fd = _pipe_with_data(b"first\nsecond") + result = list(yield_lines_from_fds([fd])) + assert result == [ + (fd, b"first\n"), + (fd, b"second"), + ] + + def test_empty_input(self): + fd = _pipe_with_data(b"") + result = list(yield_lines_from_fds([fd])) + assert not result + + def test_empty_lines(self): + fd = _pipe_with_data(b"\n\n\n") + result = list(yield_lines_from_fds([fd])) + assert result == [ + (fd, b"\n"), + (fd, b"\n"), + (fd, b"\n"), + ] + + def test_multiple_fds(self): + fd1 = _pipe_with_data(b"from-fd1\n") + fd2 = _pipe_with_data(b"from-fd2\n") + result = list(yield_lines_from_fds([fd1, fd2])) + assert set(result) == { + (fd1, b"from-fd1\n"), + (fd2, b"from-fd2\n"), + } + + def test_timeout_on_read(self): + r, w = os.pipe() + try: + with pytest.raises(TimeoutError, match="on read"): + list(yield_lines_from_fds([r], timeout=0.1)) + finally: + os.close(r) + os.close(w) + + def test_timeout_deadline_expired(self): + r, w = os.pipe() + os.write(w, b"line\n") + try: + with pytest.raises(TimeoutError, match="timeouted after"): + for _ in yield_lines_from_fds([r], timeout=0.01): + time.sleep(0.05) + finally: + os.close(r) + os.close(w) + + def test_max_line_length_ok(self): + fd = _pipe_with_data(b"short\n") + result = list(yield_lines_from_fds([fd], max_line_length=1000)) + assert result == [(fd, b"short\n")] + + def test_max_line_length_exceeded(self): + fd = _pipe_with_data(b"a" * 300 + b"\n") + with pytest.raises(LineBufferOverflow, match="exceeded 200 bytes"): + list(yield_lines_from_fds([fd], max_line_length=200)) + + def test_max_line_length_exact_boundary(self): + fd = _pipe_with_data(b"a" * 199 + b"\n") + result = list(yield_lines_from_fds([fd], max_line_length=200)) + assert len(result) == 1 + + def test_max_line_length_no_newline(self): + fd = _pipe_with_data(b"a" * 300) + with pytest.raises(LineBufferOverflow, match="exceeded 200 bytes"): + list(yield_lines_from_fds([fd], max_line_length=200)) + + def test_max_line_length_short_lines_pass(self): + fd = _pipe_with_data(b"ok\n" * 100) + result = list(yield_lines_from_fds([fd], max_line_length=10)) + assert len(result) == 100 + + def test_large_line_across_chunks(self): + """Line larger than 4096 read buffer works when no limit set.""" + fd = _pipe_with_data(b"x" * 10000 + b"\n") + result = list(yield_lines_from_fds([fd])) + assert result == [(fd, b"x" * 10000 + b"\n")] + + def test_binary_data(self): + data = bytes(range(256)).replace(b"\n", b"\x00") + b"\n" + fd = _pipe_with_data(data) + result = list(yield_lines_from_fds([fd])) + assert len(result) == 1 + assert result[0][1] == data + + def test_two_fds_lines_not_mixed(self): + """Lines from two fds are attributed correctly, not interleaved.""" + fd1 = _pipe_with_data(b"aaa\nbbb\nccc\n") + fd2 = _pipe_with_data(b"111\n222\n333\n") + result = list(yield_lines_from_fds([fd1, fd2])) + fd1_lines = [line for fd, line in result if fd == fd1] + fd2_lines = [line for fd, line in result if fd == fd2] + assert fd1_lines == [b"aaa\n", b"bbb\n", b"ccc\n"] + assert fd2_lines == [b"111\n", b"222\n", b"333\n"] + + def test_interleaved_stdout_stderr(self): + """Subprocess alternates stdout/stderr writes; lines stay attributed.""" + script = ( + "import sys\n" + "for i in range(5):\n" + " out = sys.stdout if i % 2 == 0 else sys.stderr\n" + " out.write(f'{out.name}:{i}\\n')\n" + " out.flush()\n" + ) + with subprocess.Popen( + ["python3", "-c", script], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) as sp: + stdout_fd = sp.stdout.fileno() + stderr_fd = sp.stderr.fileno() + result = list(yield_lines_from_fds([stdout_fd, stderr_fd])) + + stdout_lines = [line for fd, line in result if fd == stdout_fd] + stderr_lines = [line for fd, line in result if fd == stderr_fd] + assert stdout_lines == [b":0\n", b":2\n", b":4\n"] + assert stderr_lines == [b":1\n", b":3\n"]