Merge pull request Pyomo#3679 from jsiirola/tee-deadlock-win

Resolve `capture_output(capture_fd=True)` deadlock on Windows

Author: mrmundt

pyomo/common/tee.py

@@ -37,14 +37,17 @@
 # ~(13.1 * #threads) seconds
 _poll_timeout = 1  # 14 rounds: 0.0001 * 2**14 == 1.6384
 _poll_timeout_deadlock = 100  # seconds
-
+_pipe_buffersize = 1 << 16  # 65536
 _noop = lambda: None
 _mswindows = sys.platform.startswith('win')
 try:
     if _mswindows:
         from msvcrt import get_osfhandle
-        from win32pipe import PeekNamedPipe
         from win32file import ReadFile
+        from win32pipe import FdCreatePipe, PeekNamedPipe, SetNamedPipeHandleState
+
+        # This constant from Microsoft SetNamedPipeHandleState documentation:
+        PIPE_NOWAIT = 1
     else:
         from select import select
     _peek_available = True
@@ -59,9 +62,47 @@ def __init__(self, ostream, handle):
         super().__setattr__('_ostream', ostream)
         super().__setattr__('_handle', handle)
 
-    def flush(self):
-        self._ostream.flush()
-        self._handle.flush = True
+    if _mswindows:
+        # Because we are setting the pipe to be non-blocking in Windows,
+        # it is possible that calls to flush() and write() will raise
+        # BlockingIOError.  We will catch and retry.  In addition, we
+        # will chunk the data into pieces that should be well below the
+        # pipe buffer size (so we should avoid deadlock)
+
+        def _retry(self, fcn, *args, retries=10):
+            # Attempting to write to the pipe in the testing harness
+            # occasionally raises OSError ("No space left on disk") or
+            # BlockingIOError (when the write would be truncated).  We
+            # will re-try after a brief pause.
+            failCount = 0
+            while 1:
+                try:
+                    fcn(*args)
+                    break
+                except (OSError, BlockingIOError):
+                    failCount += 1
+                    if failCount >= retries:
+                        raise
+                    time.sleep(_poll_rampup_limit / (retries - 1))
+
+        def flush(self):
+            self._retry(self._ostream.flush)
+            self._handle.flush = True
+
+        def write(self, data):
+            chunksize = _pipe_buffersize >> 1  # 1/2 the buffer size
+            for i in range(0, len(data), chunksize):
+                self._retry(self._ostream.write, data[i : i + chunksize])
+
+        def writelines(self, data):
+            for line in data:
+                self.write(line)
+
+    else:
+
+        def flush(self):
+            self._ostream.flush()
+            self._handle.flush = True
 
     def __getattr__(self, attr):
         return getattr(self._ostream, attr)
@@ -71,13 +112,27 @@ def __setattr__(self, attr, val):
 
 
 class _AutoFlush(_SignalFlush):
-    def write(self, data):
-        self._ostream.write(data)
-        self.flush()
+    if _mswindows:
+        # Because we define write() and writelines() under windows, we
+        # need to make sure that _AutoFlush calls them
 
-    def writelines(self, data):
-        self._ostream.writelines(data)
-        self.flush()
+        def write(self, data):
+            super().write(data)
+            self.flush()
+
+        def writelines(self, data):
+            super().writelines(data)
+            self.flush()
+
+    else:
+
+        def write(self, data):
+            self._ostream.write(data)
+            self.flush()
+
+        def writelines(self, data):
+            self._ostream.writelines(data)
+            self.flush()
 
 
 class _fd_closer(object):
@@ -163,7 +218,7 @@ def __enter__(self):
         # inheritable if it is stdout/stderr
         os.dup2(out_fd, self.fd, inheritable=bool(self.std))
 
-        # We no longer need this original file descriptor
+        # We no longer need the original file descriptor
         if not isinstance(self.target, int):
             os.close(out_fd)
 
@@ -283,7 +338,9 @@ def _exit_context_stack(self, et, ev, tb):
                 cm.__exit__(et, ev, tb)
             except:
                 _stack = self.context_stack
-                FAIL.append(f"{sys.exc_info()[1]} ({len(_stack)+1}: {cm}@{id(cm):x})")
+                FAIL.append(
+                    f"{sys.exc_info()[0].__name__}: {sys.exc_info()[1]} ({len(_stack)+1}: {cm}@{id(cm):x})"
+                )
         return FAIL
 
     def __enter__(self):
@@ -492,7 +549,33 @@ def __init__(self, mode, buffering, encoding, newline):
         self.buffering = buffering
         self.newlines = newline
         self.flush = False
-        self.read_pipe, self.write_pipe = os.pipe()
+        if _peek_available and _mswindows:
+            # This is a re-implementation of os.pipe() on Windows so
+            # that we can explicitly request a larger pipe buffer (64k;
+            # matching *NIX).  Per the docs: on Windows, the pipe buffer
+            # should automatically grow if needed.  However, we have
+            # observed (see #3658) that if it happens, it can cause
+            # deadlock when clients write directly to the underlying
+            # file descriptor.  By explicitly requesting a larger buffer
+            # from the outset, we reduce the likelihood of needing to
+            # reallocate the buffer.
+            self.read_pipe, self.write_pipe = FdCreatePipe(
+                None, _pipe_buffersize, os.O_BINARY if 'b' in mode else os.O_TEXT
+            )
+            self.read_pyhandle = get_osfhandle(self.read_pipe)
+            self.write_pyhandle = get_osfhandle(self.write_pipe)
+            # Because reallocating the pipe buffer can cause deadlock
+            # (at least in the context in which we are using pipes
+            # here), we will set the write pipe to NOWAIT.  This will
+            # guarantee that we don't deadlock, but at the cost of
+            # possibly losing some output (the fprintf() to the FD will
+            # return a number of bytes written less than the string that
+            # was passed.  If the client is ignoring the return value,
+            # then *poof*: the output is truncated)
+            SetNamedPipeHandleState(self.write_pyhandle, PIPE_NOWAIT, None, None)
+        else:
+            self.read_pipe, self.write_pipe = os.pipe()
+
         if not buffering and 'b' not in mode:
             # While we support "unbuffered" behavior in text mode,
             # python does not
@@ -535,9 +618,11 @@ def close(self):
             self.write_file = None
 
         if self.write_pipe is not None:
-            # If someone else has closed the file descriptor, then
-            # python raises an OSError
+            # Close the write side of the pipe: the reader thread is
+            # waiting for the EOF so that it can shut down.
             try:
+                # If someone else has closed the file descriptor, then
+                # python raises an OSError
                 os.close(self.write_pipe)
             except OSError:
                 pass
@@ -550,8 +635,8 @@ def finalize(self, ostreams):
         self.decodeIncomingBuffer()
         if ostreams:
             self.writeOutputBuffer(ostreams, True)
+        # Close the read side of the pipe.
         os.close(self.read_pipe)
-
         if self.decoder_buffer:
             logger.error(
                 "Stream handle closed with un-decoded characters "
@@ -827,7 +912,7 @@ def _mergedReader(self):
             if _mswindows:
                 for handle in list(handles):
                     try:
-                        pipe = get_osfhandle(handle.read_pipe)
+                        pipe = handle.read_pyhandle
                         numAvail = PeekNamedPipe(pipe, 0)[1]
                         if numAvail:
                             result, new_data = ReadFile(pipe, numAvail, None)

pyomo/common/tests/test_tee.py

@@ -51,16 +51,16 @@ def flush(self):
 
     def check(self, *bases):
         """Map the recorded times to {0, 1} based on the range of times
-        recorded: anything in the first half of the range is mapped to
-        0, and anything in the second half is mapped to 1.  This
+        recorded: anything in the first two-thirds of the range is mapped to
+        0, and anything in the last third is mapped to 1.  This
         "discretizes" the times so that we can reliably compare to
         baselines.
 
         """
 
         n = list(itertools.chain(*self.buf))
-        mid = (min(n) + max(n)) / 2.0
-        result = [tuple(0 if i < mid else 1 for i in _) for _ in self.buf]
+        cutoff = min(n) + (max(n) - min(n)) * 2.0 / 3.0
+        result = [tuple(0 if i < cutoff else 1 for i in _) for _ in self.buf]
         if result not in bases:
             base = ' or '.join(str(_) for _ in bases)
             self.error = f"result {result} != baseline {base}\nRaw timing: {self.buf}"
@@ -303,6 +303,7 @@ def test_handle_prematurely_closed(self):
 class TestCapture(unittest.TestCase):
     def setUp(self):
         self.streams = sys.stdout, sys.stderr
+        self.fd = [os.dup(stream.fileno()) for stream in self.streams]
         self.reenable_gc = gc.isenabled()
         gc.disable()
         gc.collect()
@@ -313,6 +314,8 @@ def setUp(self):
 
     def tearDown(self):
         sys.stdout, sys.stderr = self.streams
+        os.dup2(self.fd[0], self.streams[0].fileno())
+        os.dup2(self.fd[1], self.streams[1].fileno())
         sys.setswitchinterval(self.switchinterval)
         if self.reenable_gc:
             gc.enable()
@@ -531,14 +534,14 @@ def test_no_fileno_stdout(self):
             self.assertEqual(len(T.context_stack), 2)
         T = tee.capture_output(capture_fd=True)
         # out & err point to something other than fd 1 and 2
-        sys.stdout = os.fdopen(os.dup(1), closefd=True)
-        sys.stderr = os.fdopen(os.dup(2), closefd=True)
+        sys.stdout = os.fdopen(os.dup(1), 'w', closefd=True)
+        sys.stderr = os.fdopen(os.dup(2), 'w', closefd=True)
         with sys.stdout, sys.stderr:
             with T:
                 self.assertEqual(len(T.context_stack), 8)
         # out & err point to fd 1 and 2
-        sys.stdout = os.fdopen(1, closefd=False)
-        sys.stderr = os.fdopen(2, closefd=False)
+        sys.stdout = os.fdopen(1, 'w', closefd=False)
+        sys.stderr = os.fdopen(2, 'w', closefd=False)
         with sys.stdout, sys.stderr:
             with T:
                 self.assertEqual(len(T.context_stack), 6)
@@ -781,7 +784,8 @@ def test_buffered_stdout_long_message(self):
         ts = timestamper()
         ts.write(f"{time.time()}")
         with tee.TeeStream(ts, ts) as t, tee.capture_output(t.STDOUT, capture_fd=fd):
-            sys.stdout.write(f"{time.time()}" + '    ' * 4096 + "\n")
+            # Note: bigger than the buffer we allocate on Windows.
+            sys.stdout.write(f"{time.time()}" + ' ' * tee._pipe_buffersize + "\n")
             time.sleep(self.dt)
         ts.write(f"{time.time()}")
         if not ts.check([(0, 0), (0, 0), (0, 0), (1, 1)]):