log_output : moved from os.fork to multiprocessing.Process

1 files changed, 78 insertions, 118 deletions

diff --git a/lib/spack/llnl/util/tty/log.py b/lib/spack/llnl/util/tty/log.py
index ca82da7b17..9b9278b73c 100644
--- a/lib/spack/llnl/util/tty/log.py
+++ b/lib/spack/llnl/util/tty/log.py
@@ -29,6 +29,7 @@ import os
 import re
 import select
 import inspect
+import multiprocessing
 
 import llnl.util.tty as tty
 import llnl.util.tty.color as color
@@ -103,155 +104,114 @@ class log_output(object):
     """Redirects output and error of enclosed block to a file.
 
     Usage:
-        with log_output(open('logfile.txt', 'w')):
+        with log_output('logfile.txt', 'w'):
            # do things ... output will be logged.
 
     or:
-        with log_output(open('logfile.txt', 'w'), echo=True):
+        with log_output('logfile.txt', echo=True):
            # do things ... output will be logged
            # and also printed to stdout.
 
-    Closes the provided stream when done with the block.
+    Opens a stream in 'w' mode at instance creation and closes
+    it at instance deletion
     If echo is True, also prints the output to stdout.
     """
-    def __init__(self, stream, echo=False, force_color=False, debug=False):
-        self.stream = stream
-
-        # various output options
+    def __init__(self, filename, echo=False, force_color=False, debug=False):
+        self.filename = filename
+        # Various output options
         self.echo = echo
         self.force_color = force_color
         self.debug = debug
 
-        # Default is to try file-descriptor reassignment unless the system 
+        # Default is to try file-descriptor reassignment unless the system
         # out/err streams do not have an associated file descriptor
         self.directAssignment = False
+        self.read, self.write = os.pipe()
 
-    def trace(self, frame, event, arg):
-        """Jumps to __exit__ on the child process."""
-        raise _SkipWithBlock()
-
+        # Spawn a daemon that writes what it reads from a pipe
+        self.p = multiprocessing.Process(target=self._forward_redirected_pipe, args=(self.read,), name='logger_daemon')
+        self.p.daemon = True
+        self.p.start()
 
     def __enter__(self):
         """Redirect output from the with block to a file.
 
-        This forks the with block as a separate process, with stdout
-        and stderr redirected back to the parent via a pipe.  If
-        echo is set, also writes to standard out.
-
+        Hijacks stdout / stderr and writes to the pipe
+        connected to the logger daemon
         """
         # remember these values for later.
         self._force_color = color._force_color
         self._debug = tty._debug
-
-        read, write = os.pipe()
-
-        self.pid = os.fork()
-        if self.pid:
-            # Parent: read from child, skip the with block.
-            os.close(write)
-
-            read_file = os.fdopen(read, 'r', 0)
-            with self.stream as log_file:
-                with keyboard_input(sys.stdin):
-                    while True:
-                        rlist, w, x = select.select([read_file, sys.stdin], [], [])
-                        if not rlist:
+        # Redirect this output to a pipe
+        self._redirect_to_pipe(self.write)
+
+    def _forward_redirected_pipe(self, read):
+        # Parent: read from child, skip the with block.
+        read_file = os.fdopen(read, 'r', 0)
+        with open(self.filename, 'w') as log_file:
+            with keyboard_input(sys.stdin):
+                while True:
+                    rlist, _, _ = select.select([read_file, sys.stdin], [], [])
+                    if not rlist:
+                        break
+
+                    # Allow user to toggle echo with 'v' key.
+                    # Currently ignores other chars.
+                    if sys.stdin in rlist:
+                        if sys.stdin.read(1) == 'v':
+                            self.echo = not self.echo
+
+                    # Handle output from the with block process.
+                    if read_file in rlist:
+                        line = read_file.readline()
+                        if not line:
                             break
 
-                        # Allow user to toggle echo with 'v' key.
-                        # Currently ignores other chars.
-                        if sys.stdin in rlist:
-                            if sys.stdin.read(1) == 'v':
-                                self.echo = not self.echo
-
-                        # handle output from the with block process.
-                        if read_file in rlist:
-                            line = read_file.readline()
-                            if not line:
-                                break
-
-                            # Echo to stdout if requested.
-                            if self.echo:
-                                sys.stdout.write(line)
+                        # Echo to stdout if requested.
+                        if self.echo:
+                            sys.stdout.write(line)
 
-                            # Stripped output to log file.
-                            log_file.write(_strip(line))
-
-            read_file.flush()
-            read_file.close()
-
-            # Set a trace function to skip the with block.
-            sys.settrace(lambda *args, **keys: None)
-            frame = inspect.currentframe(1)
-            frame.f_trace = self.trace
-
-        else:
-            # Child: redirect output, execute the with block.
-            os.close(read)
-
-            try:
-                # Save old stdout and stderr
-                self._stdout = os.dup(sys.stdout.fileno())
-                self._stderr = os.dup(sys.stderr.fileno())
-
-                # redirect to the pipe.
-                os.dup2(write, sys.stdout.fileno())
-                os.dup2(write, sys.stderr.fileno())
-            except AttributeError:
-                self.directAssignment = True
-                self._stdout = sys.stdout
-                self._stderr = sys.stderr
-                output_redirect = os.fdopen(write, 'w')
-                sys.stdout = output_redirect
-                sys.stderr = output_redirect
-
-            if self.force_color:
-                color._force_color = True
-
-            if self.debug:
-                tty._debug = True
+                        # Stripped output to log file.
+                        log_file.write(_strip(line))
+                        log_file.flush()
 
+    def _redirect_to_pipe(self, write):
+        try:
+            # Save old stdout and stderr
+            self._stdout = os.dup(sys.stdout.fileno())
+            self._stderr = os.dup(sys.stderr.fileno())
+
+            # redirect to the pipe.
+            os.dup2(write, sys.stdout.fileno())
+            os.dup2(write, sys.stderr.fileno())
+        except AttributeError:
+            self.directAssignment = True
+            self._stdout = sys.stdout
+            self._stderr = sys.stderr
+            output_redirect = os.fdopen(write, 'w')
+            sys.stdout = output_redirect
+            sys.stderr = output_redirect
+        if self.force_color:
+            color._force_color = True
+        if self.debug:
+            tty._debug = True
 
     def __exit__(self, exc_type, exception, traceback):
-        """Exits on child, handles skipping the with block on parent."""
-        # Child should just exit here.
-        if self.pid == 0:
-            # Flush the log to disk.
-            sys.stdout.flush()
-            sys.stderr.flush()
-
-            if exception:
-                # Restore stdout on the child if there's an exception,
-                # and let it be raised normally.
-                #
-                # This assumes that even if the exception is caught,
-                # the child will exit with a nonzero return code.  If
-                # it doesn't, the child process will continue running.
-                #
-                # TODO: think about how this works outside install.
-                # TODO: ideally would propagate exception to parent...
-                if self.directAssignment:
-                    sys.stdout = self._stdout
-                    sys.stderr = self._stderr
-                else:
-                    os.dup2(self._stdout, sys.stdout.fileno())
-                    os.dup2(self._stderr, sys.stderr.fileno())                    
-
-                return False
-
-            else:
-                # Die quietly if there was no exception.
-                os._exit(0)
-
-        else:
-            # If the child exited badly, parent also should exit.
-            pid, returncode = os.waitpid(self.pid, 0)
-            if returncode != 0:
-                os._exit(1)
+        """Plugs back the original file descriptors
+        for stdout and stderr
+        """
+        # Flush the log to disk.
+        sys.stdout.flush()
+        sys.stderr.flush()
+        os.dup2(self._stdout, sys.stdout.fileno())
+        os.dup2(self._stderr, sys.stderr.fileno())
 
         # restore output options.
         color._force_color = self._force_color
         tty._debug = self._debug
 
-        # Suppresses exception if it's our own.
-        return exc_type is _SkipWithBlock
+    def __del__(self):
+        """Closes the pipes and joins the daemon"""
+        os.close(self.write)
+        self.p.join()
+        os.close(self.read)