Impose time/timeout threshold on command running time (POSIX & Windows)

test_timeout.py

#! /usr/bin/env python

import sys, time, os
from timeout import timeOut                                                         

def testNoisyGraceful():
    sys.stdout.write('NoisyGraceful:'); sys.stdout.flush()
    t0 = time.time()
    #cmd = r'python -c "import time; time.sleep(10); print 1 * 4096"'
    cmd = 'python sleep.py 1 4096'
    timedOut, es, junkOut, junkErr = timeOut(cmd, 3)
    #print timedOut, es, junkOut, junkErr
    if (time.time() - t0) > 1.6:
        print 'Fail: took too long to gracefully complete'
    else:
        print 'Pass'

def testNoisyRunaway():
    sys.stdout.write('NoisyRunaway:'); sys.stdout.flush()
    t0 = time.time()
    cmd = 'python sleep.py 10 4096'
    timedOut, es, junkOut, junkErr = timeOut(cmd, 1)
    if (time.time() - t0) > 1.6:
        print 'Fail: took too long to kill'
    else:
        print 'Pass'

def testGrandChild():
    # Don't know how to check for existence of grandchildren
    # programmatically.
    sys.stdout.write('GrandChild:'); sys.stdout.flush()
    cmd = 'python sleep.py 10 0 grandchild'
    timeOut(cmd, 2)
    print '<watch process table and make sure grandchild is killed>'

if __name__ == '__main__':
    testNoisyGraceful()
    testNoisyRunaway()
    testGrandChild()

timeout.py

#! /usr/bin/env python

"""Impose time/timeout threshold on command running time (POSIX & Windows).

Monitor and possibly kill a command that exceeds the user-provided time
threshold.

Two interfaces are supported: command-line, and timeout() function.

The objectives are 1) to be able to only require that standard Python
modules be used; and 2) both Posix and NT OSes are supported.  There are
three main components to monitoring and killing a child (sub)process on
standard Python:

    1.  **Must be able to get the child PID.**
        On Windows this is only possible by using the *subprocess*
        module.  On Linux it is also possible via the popen2 module's
        Popen3 class.  (But see below why subprocess is really the only
        solution.)

    2.  **Must be able to kill the child and all its descendants.**
        On Windows this can be done with ctypes or the win32api modules,
        as demonstrated in Recipe #347462.  However both of these are
        not part of standard Python (though ctypes should be in v2.5).
        Thus we are relegated to using NT's *taskkill* utility.  It has
        an option (/T) to kill descendant processes.  On Linux the kill
        can only affect all descendants if the child is made into a
        *group leader*, and killed with a negated PID.  This is only
        possible by using subprocess.Popen's *preexec_fn* parameter.

    3.  **Must have a timing mechanism to signal an expired threshold.**
        The simplest tool for this job is the threading module's *Timer*
        convenience class.  Unfortunately, this behaves differently
        between Windows and Linux when a kill is involved -- on Linux
        behavior is as expected; on Windows the non-timer thread's call
        to the timer's "isAlive" reports True even after the expiration
        function has been called on Windows. Thus we are again relegated
        to an ad-hoc sleep loop. The loop is actually quite
        small/trivial and is probably as good as Timer anyway.

    4.  **Must be able to handle voluminous amounts of child data.**
        The subprocess module (and popen-family) will get into a
        deadlock state when trying to read >4095 bytes from the child's
        stdout/stderr.  A problem is that we must assume that we have no
        control over the behavior of the child's output.  So the
        simplest way to not deadlock is to have the child's
        stdout/stderr sent to a temporary file.  This is only possible
        by using subprocess.Popen's *stdout/stderr* parameters.

    5.  **Must preserve child return code.**


Limitations:
    * Requires Python 2.4+.

    * Use of this module has potential to add one jiffy to each run.

    * Time overhead of using of temp files (this is only ~20ms).

    * Overhead Python startup (up to ~500ms).

See also:
    Another subprocess module with timeout (*nix only):
        http://www.pixelbeat.org/libs/subProcess.py

    Timer discussion:
        http://mail.python.org/pipermail/python-list/2002-October/127761.html

    Timer doc:
        http://www.python.org/doc/2.2/lib/timer-objects.html

    Subprocess module for older Python (windows only):
        http://effbot.org/downloads/#subprocess

    Kill recipe (Windows):
        http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/347462

    Capturing stdout/stderr streams recipe:
        http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/52296

    Subprocess module doc:
        http://docs.python.org/lib/module-subprocess.html

    Flow controls issues doc:
        http://www.python.org/doc/lib/popen2-flow-control.html

"""

__author__    = 'Micah Elliott'
__date__      = '$Date: 2006-07-25 10:55:50 -0700 (Tue, 25 Jul 2006) $'
__copyright__ = 'WTFPL http://sam.zoy.org/wtfpl/'

#-----------------------------------------------------------------------

import time, os, sys, tempfile

try:
    import subprocess
except ImportError, err:
    print 'ERROR:', err
    print 'timeout: subprocess module required, introduced in Python 2.4.'
    if os.name == 'nt':
        print 'Downloadable from: http://effbot.org/downloads/#subprocess'
    sys.exit(1)

#import pdb
#def runPdb(*args): pdb.pm()
#sys.excepthook = runPdb

GRANULARITY = 10

def timeOut(cmd, thresh):
    timedOut = False # for when called from or other util
    # Popen only needs to set child as group leader on posix.
    if os.name == 'posix':
        setProcessGroup = os.setpgrp
    else:
        setProcessGroup = None
    # Could do this twice if need to separate stdout from stderr.
    tmpOut, tmpOutName = tempfile.mkstemp()
    tmpErr, tmpErrName = tempfile.mkstemp()
    # Spawn child command.
    cmd = cmd.split()
    #print '*** cmd:', cmd
    p = subprocess.Popen(cmd,
                         stdout=tmpOut,  # file descriptor
                         stderr=tmpErr,
                         preexec_fn=setProcessGroup
                        )
    for jiffy in range( int(thresh * GRANULARITY) ):
        pollStatus = p.poll()
        # Child still running.
        if pollStatus is None:
            time.sleep(1./GRANULARITY)
        # Child finished.
        else:
            #print 'Child finished'
            output = p.communicate()
            # output should be empty since writing is to temp file.
            status = p.returncode
            break
    # Time threshold exceeded.
    else:
        #print 'Timer expired'
        timedOut = True
        # kill the child
        if os.name == 'posix':
            import signal
            #print 'killing with', signal.SIGTERM
            #status = os.system('kill -%s %s' % (signal.SIGTERM, p.pid))
            #os.kill(p.pid, signal.SIGTERM) # Doesn't kill children.
                # those two don't kill descendants.
            os.killpg(p.pid, signal.SIGTERM)  # or kill(-pid, sig)
            s = os.waitpid(p.pid, 0)[1] # wait retires zombie child
            status = os.WTERMSIG(s) + 128 # emulate shell
            #print( 'Expired! Killed pid %s WEXITSTATUS %s WTERMSIG %s' %
                   #(p.pid, os.WEXITSTATUS(s), os.WTERMSIG(s)) )
        else:
            killcmd = 'taskkill /t /f /pid %s' % p.pid
            #print 'killing with "%s"' % killcmd
            #status = os.system(killcmd)
            pi, poe = os.popen4(killcmd)
            pi.close()
            poe.read()  # ignore output
            status = poe.close()
            # taskkill simply returns 0, so fabricate bogus silly code to
            # indicate an expiration.
            status = 143
    # Windows requires closure before a second open!
    os.close(tmpOut)
    os.close(tmpErr)
    tmpOut2 = open(tmpOutName)
    tmpErr2 = open(tmpErrName)
    out = tmpOut2.read()
    err = tmpErr2.read()
    tmpOut2.close()
    tmpErr2.close()
    os.remove(tmpOutName)
    os.remove(tmpErrName)
    return (timedOut, status, out, err)

def timeout():
    """Command line parsing and usage.
    """
    usage = 'usage: timeout THRESHOLD COMMAND'
    def tryHelp(msg):
        print>>sys.stderr, 'timeout: error:', msg
        print 'Try "timeout --help" for more information.'
        sys.exit(1)
    if len(sys.argv) < 2:
        tryHelp('missing arg.')
    if sys.argv[1] == '-h' or sys.argv[1] == '--help':
        print usage
        sys.exit(0)
    if len(sys.argv) < 3:
        tryHelp('missing args.')
    try:
        thresh = float(sys.argv[1])
    except ValueError:
        tryHelp('first arg must be integer or float value.')
    cmd = ' '.join(sys.argv[2:])
    tc, status, out, err = timeOut(cmd, thresh)
    if out: print out,
    if err: print>>sys.stderr, err,
    #print 'status', status
    sys.exit(status)

if __name__ == '__main__':
    timeout()