Spawn with timeout and with the same behaviour across operating systems

spawn.py

"""
Execute external processes.
"""
import os

from twisted.internet import reactor, error, defer
from twisted.internet.protocol import ProcessProtocol

SPAWN_PROCESS_TIMEOUT = 5


class ProcessTimedOut(error.ConnectionLost):
    """
    Process was killed as is took to much to execute.
    """


class ProcessSpawnFailed(error.ConnectionLost):
    """
    Failed to spawn the process.
    """

    def __init__(self, message='Failed to launch the process'):
        self.message = message

    def __str__(self):
        return self.message


class ExecuteWithTimeout(ProcessProtocol):
    """
    Allow command to execute, ignoring stdout and stderrr.

    Result is returned via the `deferred`.
     * Success if command exit code is 0.
     * failure if command exit code is not 0.
     * ProcessTimedOut failure when command takes more than `timeout`
       to execute.
    * ProcessSpawnFailed failure when failed to launch the executable
    """
    def __init__(self, deferred, scheduler, timeout):
        self._timeout = timeout
        self._deferred = deferred
        self._scheduler = scheduler
        self._killed = False
        self._timeout_call = None
        self._first_err_received = False
        self._spawn_failure = None

    def errReceived(self, data):
        """
        Data was received on standard error.
        """
        # To detect spawn errors on Unix, we check the first error message
        # and see if it is from execvpe.
        if os.name != 'posix':
            return

        if self._first_err_received:
            return
        self._first_err_received = True

        if not data.startswith('Upon execvpe'):
            return

        # Keep last error line as error message.
        self._spawn_failure = to_unicode(data.splitlines()[-1])

    def connectionMade(self):
        """
        Called when process was started.
        """
        def kill_if_alive():
            try:
                self.transport.signalProcess('KILL')
            except error.ProcessExitedAlready:
                pass
            else:
                self._killed = True
                self.transport.loseConnection()

        self._timeout_call = self._scheduler.callLater(
            self._timeout, kill_if_alive)

    def processEnded(self, reason):
        """
        Called when process terminates.
        """
        exit_code = reason.value.exitCode

        try:
            self._timeout_call.cancel()
        except error.AlreadyCalled:
            pass

        # Sometimes, even when we kill a process, GetExitCodeProcess
        # still returns a zero exit status.
        # http://stackoverflow.com/q/2061735/539264
        if self._killed and exit_code == 0:
            exit_code = 1

        if exit_code == 0:
            self._deferred.callback(None)
        elif self._spawn_failure:
            self._deferred.errback(ProcessSpawnFailed(self._spawn_failure))
        elif self._killed:
            self._deferred.errback(ProcessTimedOut())
        else:
            self._deferred.errback(reason)


def get_process_exit_code(
        command, arguments, reactor=reactor, timeout=SPAWN_PROCESS_TIMEOUT):
    """
    Execute `command` with `arguments` and return the exit code.
    """
    deferred = defer.Deferred()
    protocol = ExecuteWithTimeout(deferred, scheduler=reactor, timeout=timeout)

    final_command = command.encode('utf-8')
    arguments = [a.encode('utf-8') for a in arguments]
    final_arguments = (final_command,) + tuple(arguments)

    try:
        reactor.spawnProcess(protocol, final_command, final_arguments)
    except OSError, error:
        deferred.errback(ProcessSpawnFailed(unicode(error)))

    return deferred