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

-SSH command size: 838
+SSH command size: 906

                         Original           Minimized           Compressed
-mitogen.parent        98827  96.5KiB  51219 50.0KiB 51.8%  12942 12.6KiB 13.1%
+mitogen.parent        99063  96.7KiB  51318 50.1KiB 51.8%  12979 12.7KiB 13.1%

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

Author: mhartmay

mitogen/parent.py

@@ -1418,6 +1418,8 @@ 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
+    #   _: throw away variable
 
     # Final os.close(STDOUT_FILENO) to avoid --py-debug build corrupting stream with
     # "[1234 refs]" during exit.
@@ -1439,8 +1441,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],[],[],10)
+            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)