Dedicate the line-by-line reader from FD a separate method

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)

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"<timeouted after {timeout}s>\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"<timeouted after {timeout}s (too much output)>\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"<timeouted after {timeout}s (on read)>\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(),

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"<stdout>:0\n", b"<stdout>:2\n", b"<stdout>:4\n"]
+        assert stderr_lines == [b"<stderr>:1\n", b"<stderr>:3\n"]