Fix spot check hang (#1242).

borgmatic/execute.py

@@ -3,6 +3,7 @@ import contextlib
 import enum
 import json
 import logging
+import os
 import select
 import subprocess
 import textwrap
@@ -196,6 +197,50 @@ def handle_log_record(log_record, last_lines):
     return log_record
 
 
+READ_CHUNK_SIZE = 4096
+
+
+def read_lines(buffer, process, line_separator='\n'):
+    '''
+    Given a Python buffer (like stdout) ready for reading, its process, and a line separator,
+    repeatedly yield a tuple of (decoded) lines from the buffer until the process has exited.
+
+    It is assumed that this function's generator is used in conjunction with an external select()
+    call to know when to read more lines. Otherwise, the generator will busywait if it's called in a
+    tight loop.
+    '''
+    data = ''
+
+    while True:
+        #chunk = buffer.read(READ_CHUNK_SIZE).decode()
+        chunk = os.read(buffer.fileno(), READ_CHUNK_SIZE).decode()
+
+        if not chunk:  # EOF
+            # The process is still running, so we keep running too.
+            if process.poll() is None:
+                continue
+
+            break
+
+        data += chunk
+        lines = []
+
+        while True:
+            separator_position = data.find(line_separator)
+
+            if separator_position == -1:
+                break
+
+            lines.append(data[:separator_position].rstrip())
+            data = data[separator_position + 1 :]
+
+        yield tuple(lines)
+
+    # Yield any leftover data from the end of the buffer.
+    if data:
+        yield tuple(data.rstrip().splitlines())
+
+
 def log_outputs(  # noqa: PLR0912
     processes,
     exclude_stdouts,
@@ -231,6 +276,12 @@ def log_outputs(  # noqa: PLR0912
     }
     output_buffers = list(process_for_output_buffer.keys())
     process_to_capture = processes[-1]
+    reader_for_output_buffer = {
+        buffer: read_lines(buffer, process)
+        for process in processes
+        if process.stdout or process.stderr
+        for buffer in output_buffers_for_process(process, exclude_stdouts)
+    }
 
     # Log output for each process until they all exit.
     while any(process.poll() is None for process in processes):
@@ -252,40 +303,44 @@ def log_outputs(  # noqa: PLR0912
                         ):
                             # Add the process's output to output_buffers to ensure it'll get read.
                             output_buffers.append(other_process.stdout)
+                            process_for_output_buffer[other_process.stdout] = other_process
+                            reader_for_output_buffer[other_process.stdout] = read_lines(
+                                other_process.stdout, other_process
+                            )
 
-                line = ready_buffer.readline().rstrip().decode()
-                if not line or not ready_process:
+                try:
+                    lines = next(reader_for_output_buffer[ready_buffer])
+                except StopIteration:
                     continue
 
-                command = (
-                    ready_process.args.split(' ')
-                    if isinstance(ready_process.args, str)
-                    else ready_process.args
-                )
+                for line in lines:
+                    if not line or not ready_process:
+                        continue
 
-                # Keep the last few lines of output in case the process errors and we need the
-                # output for the exception below.
-                log_record = handle_log_record(
-                    parse_log_line(
-                        line=line,
-                        log_level=output_log_level,
-                        elevate_stderr=(
-                            ready_buffer == ready_process.stderr and not capture_stderr
+                    command = (
+                        ready_process.args.split(' ')
+                        if isinstance(ready_process.args, str)
+                        else ready_process.args
+                    )
+
+                    # Keep the last few lines of output in case the process errors and we need the
+                    # output for the exception below.
+                    log_record = handle_log_record(
+                        parse_log_line(
+                            line=line,
+                            log_level=output_log_level,
+                            elevate_stderr=(
+                                ready_buffer == ready_process.stderr and not capture_stderr
+                            ),
+                            borg_local_path=borg_local_path,
+                            command=command,
                         ),
-                        borg_local_path=borg_local_path,
-                        command=command,
-                    ),
-                    last_lines=process_last_lines[ready_process],
-                )
+                        last_lines=process_last_lines[ready_process],
+                    )
 
-                if log_record.levelno is None and ready_process == process_to_capture:
-                    print('***', log_record.getMessage())
-                    yield log_record.getMessage()
-                else:
-                    # Yield an empty string just to keep the generator executing. Without this,
-                    # logged lines might not get processed by the caller until the next yielded
-                    # message.
-                    yield ''
+                    if log_record.levelno is None and ready_process == process_to_capture:
+                        print('***', log_record.getMessage())
+                        yield log_record.getMessage()
 
         for process in processes:
             exit_code = process.poll() if output_buffers else process.wait()
@@ -301,32 +356,6 @@ def log_outputs(  # noqa: PLR0912
 
                 # If an error occurs, include its output in the raised exception so that we don't
                 # inadvertently hide error output.
-                for output_buffer in output_buffers_for_process(process, exclude_stdouts):
-                    # Collect any straggling output lines that came in since we last gathered output.
-                    while output_buffer:  # pragma: no cover
-                        line = output_buffer.readline().rstrip().decode()
-                        if not line:
-                            break
-
-                        log_record = handle_log_record(
-                            parse_log_line(
-                                line=line,
-                                log_level=output_log_level,
-                                elevate_stderr=(
-                                    output_buffer == process.stderr and not capture_stderr
-                                ),
-                                borg_local_path=borg_local_path,
-                                command=command,
-                            ),
-                            last_lines=last_lines,
-                        )
-
-                        if log_record.levelno is None and process == process_to_capture:
-                            print('**e', log_record.getMessage())
-                            yield log_record.getMessage()
-                        else:
-                            yield
-
                 if len(last_lines) == ERROR_OUTPUT_MAX_LINE_COUNT:
                     last_lines.insert(0, '...')
 
@@ -349,27 +378,29 @@ def log_outputs(  # noqa: PLR0912
     # Now that all processes have exited, drain and consume any last output.
     for output_buffer in output_buffers:
         process = process_for_output_buffer.get(output_buffer)
+
+        if not process:
+            continue
+
+        command = process.args.split(' ') if isinstance(process.args, str) else process.args
         last_lines = process_last_lines[process]
 
-        for line in output_buffer.readlines():
-            log_record = handle_log_record(
-                parse_log_line(
-                    line=line.rstrip().decode(),
-                    log_level=output_log_level,
-                    elevate_stderr=(
-                        output_buffer == process.stderr and not capture_stderr
+        for lines in reader_for_output_buffer[output_buffer]:
+            for line in lines:
+                log_record = handle_log_record(
+                    parse_log_line(
+                        line=line.rstrip(),
+                        log_level=output_log_level,
+                        elevate_stderr=(output_buffer == process.stderr and not capture_stderr),
+                        borg_local_path=borg_local_path,
+                        command=command,
                     ),
-                    borg_local_path=borg_local_path,
-                    command=command,
-                ),
-                last_lines=last_lines,
-            )
+                    last_lines=last_lines,
+                )
 
-            if log_record.levelno is None and process == process_to_capture:
-                yield log_record.getMessage()
-                print('**d', log_record.getMessage())
-            else:
-                yield ''
+                if log_record.levelno is None and process == process_to_capture:
+                    print('***', log_record.getMessage())
+                    yield log_record.getMessage()
 
 
 SECRET_COMMAND_FLAG_NAMES = {'--password'}

tests/integration/test_execute.py

@@ -269,7 +269,7 @@ def test_log_outputs_kills_other_processes_and_raises_when_one_errors():
         other_process,
         (),
     ).and_return((other_process.stdout,))
-    flexmock(other_process).should_receive('kill').once()
+    flexmock(other_process).should_call('kill').once()
 
     with pytest.raises(subprocess.CalledProcessError) as error:
         tuple(
@@ -291,12 +291,6 @@ def test_log_outputs_kills_other_processes_and_returns_when_one_exits_with_warni
     flexmock(module).should_receive('command_for_process').and_return('grep')
 
     process = subprocess.Popen(['grep'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
-    flexmock(module).should_receive('interpret_exit_code').with_args(
-        ['grep'],
-        None,
-        'borg',
-        None,
-    ).and_return(module.Exit_status.SUCCESS)
     flexmock(module).should_receive('interpret_exit_code').with_args(
         ['grep'],
         2,
@@ -313,7 +307,7 @@ def test_log_outputs_kills_other_processes_and_returns_when_one_exits_with_warni
         None,
         'borg',
         None,
-    ).and_return(module.Exit_status.SUCCESS)
+    ).and_return(module.Exit_status.STILL_RUNNING)
     flexmock(module).should_receive('output_buffers_for_process').with_args(process, ()).and_return(
         (process.stdout,),
     )
@@ -321,7 +315,7 @@ def test_log_outputs_kills_other_processes_and_returns_when_one_exits_with_warni
         other_process,
         (),
     ).and_return((other_process.stdout,))
-    flexmock(other_process).should_receive('kill').once()
+    flexmock(other_process).should_call('kill').once()
 
     assert (
         tuple(
@@ -372,18 +366,15 @@ def test_log_outputs_vents_other_processes_when_one_exits():
     ).and_return((other_process.stdout,))
     flexmock(process.stdout).should_call('readline').at_least().once()
 
-    assert (
-        tuple(
-            module.log_outputs(
-                (process, other_process),
-                exclude_stdouts=(process.stdout,),
-                output_log_level=logging.INFO,
-                borg_local_path='borg',
-                borg_exit_codes=None,
-            )
+    assert tuple(
+        module.log_outputs(
+            (process, other_process),
+            exclude_stdouts=(process.stdout,),
+            output_log_level=logging.INFO,
+            borg_local_path='borg',
+            borg_exit_codes=None,
         )
-        == ()
-    )
+    ) == ('',)
 
 
 def test_log_outputs_does_not_error_when_one_process_exits():
@@ -433,12 +424,6 @@ def test_log_outputs_truncates_long_error_output():
     flexmock(module).should_receive('command_for_process').and_return('grep')
 
     process = subprocess.Popen(['grep'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
-    flexmock(module).should_receive('interpret_exit_code').with_args(
-        ['grep'],
-        None,
-        'borg',
-        None,
-    ).and_return(module.Exit_status.SUCCESS)
     flexmock(module).should_receive('interpret_exit_code').with_args(
         ['grep'],
         2,
@@ -458,8 +443,8 @@ def test_log_outputs_truncates_long_error_output():
             )
         )
 
-    assert error.value.returncode == 2
-    assert error.value.output.startswith('...')
+        assert error.value.returncode == 2
+        assert error.value.output.startswith('...')
 
 
 def test_log_outputs_with_no_output_logs_nothing():
@@ -487,8 +472,8 @@ def test_log_outputs_with_unfinished_process_re_polls():
     flexmock(module.logger).should_receive('log').never()
     flexmock(module).should_receive('interpret_exit_code').and_return(module.Exit_status.SUCCESS)
 
-    process = subprocess.Popen(['true'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
-    flexmock(process).should_receive('poll').and_return(None).and_return(0).times(3)
+    process = subprocess.Popen(['sleep', '0.001'], stdout=subprocess.PIPE)
+    flexmock(process).should_call('poll').at_least().times(3)
     flexmock(module).should_receive('output_buffers_for_process').and_return((process.stdout,))
 
     assert (