mitogen: _first_stage: Add timeout handling

Do not wait/block forever for data to be read.

Add a test for this.

The test can be run using the following command:

PYTHONPATH=$(pwd)/tests:$PYTHONPATH python -m unittest -v tests.first_stage_test

Signed-off-by: Marc Hartmayer <[email protected]>

Author: mhartmay

mitogen/parent.py

@@ -1378,6 +1378,9 @@ class Connection(object):
     #: user.
     exception = None
 
+    #: First stage timeout in seconds,
+    _first_stage_timeout = 10
+
     #: Extra text appended to :class:`EofError` if that exception is raised on
     #: a failed connection attempt. May be used in subclasses to hint at common
     #: problems with a particular connection method.
@@ -1418,6 +1421,9 @@ def __repr__(self):
     #   C: the decompressed core source.
     #   n: size of the compressed core source to be read
     #   V: data chunk
+    #   rl: list of FDs ready for reading
+    #   t: timeout value in seconds
+    #   _: throw away variable
 
     # Final os.close(STDOUT_FILENO) to avoid --py-debug build corrupting stream with
     # "[1234 refs]" during exit.
@@ -1439,8 +1445,16 @@ def _first_stage():
             os.environ['ARGV0']=sys.executable
             os.execl(sys.executable,sys.executable+'(mitogen:%s)'%sys.argv[2])
         os.write(1,'MITO000\n'.encode())
-        n=int(sys.argv[3]);C=''.encode();V='V'
-        while n>len(C) and V:select.select([0],[],[]);V=os.read(0,n-len(C));C+=V
+        n=int(sys.argv[3])
+        t=float(sys.argv[4])
+        C=''.encode()
+        V='V'
+        while n>len(C) and V:
+            rl,_,_=select.select([0],[],[],t)
+            if not rl:
+                raise Exception("TimeoutError")
+            V=os.read(0,n-len(C))
+            C+=V
         C=zlib.decompress(C)
         f=os.fdopen(W,'wb',0)
         f.write(C)
@@ -1486,6 +1500,7 @@ def get_boot_command(self):
             encoded.decode(),
             self.options.remote_name,
             str(len(self.get_preamble())),
+            str(self._first_stage_timeout),
         ]
 
     def get_econtext_config(self):

tests/first_stage_test.py

@@ -1,4 +1,5 @@
 import subprocess
+import sys
 
 import mitogen.core
 import mitogen.parent
@@ -50,3 +51,36 @@ def test_valid_syntax(self):
             )
         finally:
             fp.close()
+
+    def test_timeout_error(self):
+        options = mitogen.parent.Options(max_message_size=123)
+        conn = mitogen.parent.Connection(options, self.router)
+        conn.context = mitogen.core.Context(None, 123)
+        # We do not want to wait the default of 10s, change it to 0.1s
+        conn._first_stage_timeout = 0.1
+        args = conn.get_boot_command()
+
+        # The boot command should write an ECO marker to stdout, read the
+        # preamble from stdin, then execute it.
+
+        # This test attaches closes stdin to create a specific failure
+        # 1. Fork child tries to read from STDIN, but fails as it is closed
+        # 2. Fork child raises TimeoutError
+        # 3. Fork child's file descriptors (write pipes) are closed by the OS
+        # 4. Fork parent does `dup(<read pipe>, <stdin>)` and `exec(<python>)`
+        # 5. Python reads `b''` (i.e. EOF) from stdin (a closed pipe)
+        # 6. Python runs `''` (a valid script) and exits with success
+
+        proc = mitogen.parent.popen(args=args,
+                                    stdout=subprocess.PIPE,
+                                    stderr=subprocess.PIPE,
+                                    preexec_fn=sys.stdin.close
+                                    )
+        stdout, stderr = proc.communicate()
+        self.assertEqual(0, proc.returncode)
+        self.assertEqual(stdout,
+                         mitogen.parent.BootstrapProtocol.EC0_MARKER+b('\n'))
+        self.assertIn(
+            b("TimeoutError"),
+            stderr,
+        )