Update Tornado lib

libs/tornado/process.py

@@ -19,14 +19,19 @@
 from __future__ import absolute_import, division, with_statement
 
 import errno
-import logging
+import functools
 import os
+import signal
+import subprocess
 import sys
 import time
 
 from binascii import hexlify
 
 from tornado import ioloop
+from tornado.iostream import PipeIOStream
+from tornado.log import gen_log
+from tornado import stack_context
 
 try:
     import multiprocessing  # Python 2.6+
@@ -45,7 +50,7 @@ def cpu_count():
         return os.sysconf("SC_NPROCESSORS_CONF")
     except ValueError:
         pass
-    logging.error("Could not detect number of processors; assuming 1")
+    gen_log.error("Could not detect number of processors; assuming 1")
     return 1
 
 
@@ -98,7 +103,7 @@ def fork_processes(num_processes, max_restarts=100):
         raise RuntimeError("Cannot run in multiple processes: IOLoop instance "
                            "has already been initialized. You cannot call "
                            "IOLoop.instance() before calling start_processes()")
-    logging.info("Starting %d processes", num_processes)
+    gen_log.info("Starting %d processes", num_processes)
     children = {}
 
     def start_child(i):
@@ -128,13 +133,13 @@ def fork_processes(num_processes, max_restarts=100):
             continue
         id = children.pop(pid)
         if os.WIFSIGNALED(status):
-            logging.warning("child %d (pid %d) killed by signal %d, restarting",
+            gen_log.warning("child %d (pid %d) killed by signal %d, restarting",
                             id, pid, os.WTERMSIG(status))
         elif os.WEXITSTATUS(status) != 0:
-            logging.warning("child %d (pid %d) exited with status %d, restarting",
+            gen_log.warning("child %d (pid %d) exited with status %d, restarting",
                             id, pid, os.WEXITSTATUS(status))
         else:
-            logging.info("child %d (pid %d) exited normally", id, pid)
+            gen_log.info("child %d (pid %d) exited normally", id, pid)
             continue
         num_restarts += 1
         if num_restarts > max_restarts:
@@ -156,3 +161,122 @@ def task_id():
     """
     global _task_id
     return _task_id
+
+class Subprocess(object):
+    """Wraps ``subprocess.Popen`` with IOStream support.
+
+    The constructor is the same as ``subprocess.Popen`` with the following
+    additions:
+
+    * ``stdin``, ``stdout``, and ``stderr`` may have the value
+      `tornado.process.Subprocess.STREAM`, which will make the corresponding
+      attribute of the resulting Subprocess a `PipeIOStream`.
+    * A new keyword argument ``io_loop`` may be used to pass in an IOLoop.
+    """
+    STREAM = object()
+
+    _initialized = False
+    _waiting = {}
+
+    def __init__(self, *args, **kwargs):
+        self.io_loop = kwargs.pop('io_loop', None)
+        to_close = []
+        if kwargs.get('stdin') is Subprocess.STREAM:
+            in_r, in_w = os.pipe()
+            kwargs['stdin'] = in_r
+            to_close.append(in_r)
+            self.stdin = PipeIOStream(in_w, io_loop=self.io_loop)
+        if kwargs.get('stdout') is Subprocess.STREAM:
+            out_r, out_w = os.pipe()
+            kwargs['stdout'] = out_w
+            to_close.append(out_w)
+            self.stdout = PipeIOStream(out_r, io_loop=self.io_loop)
+        if kwargs.get('stderr') is Subprocess.STREAM:
+            err_r, err_w = os.pipe()
+            kwargs['stderr'] = err_w
+            to_close.append(err_w)
+            self.stdout = PipeIOStream(err_r, io_loop=self.io_loop)
+        self.proc = subprocess.Popen(*args, **kwargs)
+        for fd in to_close:
+            os.close(fd)
+        for attr in ['stdin', 'stdout', 'stderr', 'pid']:
+            if not hasattr(self, attr):  # don't clobber streams set above
+                setattr(self, attr, getattr(self.proc, attr))
+        self._exit_callback = None
+        self.returncode = None
+
+    def set_exit_callback(self, callback):
+        """Runs ``callback`` when this process exits.
+
+        The callback takes one argument, the return code of the process.
+
+        This method uses a ``SIGCHILD`` handler, which is a global setting
+        and may conflict if you have other libraries trying to handle the
+        same signal.  If you are using more than one ``IOLoop`` it may
+        be necessary to call `Subprocess.initialize` first to designate
+        one ``IOLoop`` to run the signal handlers.
+
+        In many cases a close callback on the stdout or stderr streams
+        can be used as an alternative to an exit callback if the
+        signal handler is causing a problem.
+        """
+        self._exit_callback = stack_context.wrap(callback)
+        Subprocess.initialize(self.io_loop)
+        Subprocess._waiting[self.pid] = self
+        Subprocess._try_cleanup_process(self.pid)
+
+    @classmethod
+    def initialize(cls, io_loop=None):
+        """Initializes the ``SIGCHILD`` handler.
+
+        The signal handler is run on an IOLoop to avoid locking issues.
+        Note that the IOLoop used for signal handling need not be the
+        same one used by individual Subprocess objects (as long as the
+        IOLoops are each running in separate threads).
+        """
+        if cls._initialized:
+            return
+        if io_loop is None:
+            io_loop = ioloop.IOLoop.instance()
+        cls._old_sigchld = signal.signal(
+            signal.SIGCHLD,
+            lambda sig, frame: io_loop.add_callback_from_signal(cls._cleanup))
+        cls._initialized = True
+
+    @classmethod
+    def uninitialize(cls):
+        """Removes the ``SIGCHILD`` handler."""
+        if not cls._initialized:
+            return
+        signal.signal(signal.SIGCHLD, cls._old_sigchld)
+        cls._initialized = False
+
+    @classmethod
+    def _cleanup(cls):
+        for pid in cls._waiting.keys():
+            cls._try_cleanup_process(pid)
+
+    @classmethod
+    def _try_cleanup_process(cls, pid):
+        try:
+            ret_pid, status = os.waitpid(pid, os.WNOHANG)
+        except OSError, e:
+            if e.args[0] == errno.ECHILD:
+                return
+        if ret_pid == 0:
+            return
+        assert ret_pid == pid
+        subproc = cls._waiting.pop(pid)
+        subproc.io_loop.add_callback_from_signal(
+            functools.partial(subproc._set_returncode, status))
+
+    def _set_returncode(self, status):
+        if os.WIFSIGNALED(status):
+            self.returncode = -os.WTERMSIG(status)
+        else:
+            assert os.WIFEXITED(status)
+            self.returncode = os.WEXITSTATUS(status)
+        if self._exit_callback:
+            callback = self._exit_callback
+            self._exit_callback = None
+            callback(self.returncode)