Source code for utah.timeout

# Ubuntu Testing Automation Harness
# Copyright 2012 Canonical Ltd.

# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License version 3, as published
# by the Free Software Foundation.

# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranties of
# MERCHANTABILITY, SATISFACTORY QUALITY, or FITNESS FOR A PARTICULAR
# PURPOSE.  See the GNU General Public License for more details.

# You should have received a copy of the GNU General Public License along
# with this program.  If not, see <http://www.gnu.org/licenses/>.

"""Provide functionality to execute command with timeouts."""


import os
import signal
import subprocess
import sys

from utah.commandstr import commandstr


[docs]class UTAHTimeout(SystemExit):

    """Provide a special exception to indicate an operation timed out."""

    pass


# Mostly pulled from utah/client/common.py
# Inspired by:
# http://stackoverflow.com/a/3326559
[docs]def timeout(timeout, command, *args, **kw):
    """Run a command for up to ``timeout`` seconds.

    :param timeout:
        Maximum amount of time to wait for the command to complete in seconds.
    :type timeout: int
    :param command:
        Command whose execution should complete before timeout expires.
    :type command: callable
    :param args: Positional arguments to be passed to the callable.
    :param kwargs: Keyword arguments to be passed to the callable.
    :returns: The value returned by the callable.
    :raises UTAHTimeout:
        If command execution hasn't finished before ``timeout`` seconds.

    .. seealso:: :func:`utah.retry.retry`

    """
    #TODO: Better support for nested timeouts.
    if command is None:
        return

    def alarm_handler(_signum, _frame):
        raise UTAHTimeout('{} timed out after {} seconds'
                          .format(commandstr(command, *args, **kw),
                                  str(timeout)))

    if timeout is None:
        return command(*args, **kw)
    elif timeout != 0:
        oldtimeout = signal.alarm(0)
        if oldtimeout > 0:
            # We're in a nested timeout.  Use the outermost one for now.
            signal.alarm(oldtimeout)
            return command(*args, **kw)
        else:
            signal.signal(signal.SIGALRM, alarm_handler)
            signal.alarm(timeout)
            retval = command(*args, **kw)
            if timeout != 0:
                signal.alarm(0)
            return retval
    else:
        raise UTAHTimeout('0 timeout specified')


[docs]def subprocesstimeout(timeout, *args, **kw):
    """Run command through ``subprocess.Popen`` for up to ``timeout`` seconds.

    Command termination is checked using ``subprocess.Popen.poll``.

    :param timeout:
        Maximum amount of time to wait for the command to complete in seconds.
    :type timeout: int
    :param args: Positional arguments to be passed to ``subprocess.Popen``.
    :param kwargs: Keyword arguments to be passed to the ``subprocess.Popen``.
    :returns: The subprocess object
    :rtype: ``subprocess.Popen``
    :raises UTAHTimeout:
        If command execution hasn't finished before ``timeout`` seconds.

    """
    if args is None:
        return

    class TimeoutAlarm(Exception):
        pass

    def alarm_handler(_signum, _frame):
        raise TimeoutAlarm

    if timeout != 0:
        signal.signal(signal.SIGALRM, alarm_handler)
        oldtimeout = signal.alarm(timeout)
        if oldtimeout > 0:
            signal.alarm(oldtimeout)

    try:
        p = subprocess.Popen(*args, **kw)
        while p.poll() is None:
            pass
        if timeout != 0:
            signal.alarm(0)
        return p
    except TimeoutAlarm:
        pids = [p.pid]
        # Kill p's children too.
        pids.extend(get_process_children(p.pid))

        for pid in pids:
            # process might have died before getting to this line
            # so wrap to avoid OSError: no such process
            try:
                os.kill(pid, signal.SIGKILL)
            except OSError:
                pass

        msg = ('{} timed out after {} seconds'
               .format(' '.join(args[0]), str(timeout)))
        raise UTAHTimeout(msg)


[docs]def get_process_children(pid, logmethod=sys.stderr.write):
    """Find process children so they can be killed when the timeout expires.

    :param pid: Process ID for the parent process
    :type pid: int
    :returns: The pid for each children process
    :rtype: list(int)

    """
    try:
        pids = subprocess.check_output(['ps', '--no-headers', '-o', 'pid',
                                        '--ppid', pid]).split()
        return [int(p) for p in pids]
    except subprocess.CalledProcessError:
        return []
    except OSError as err:
        logmethod('Could not kill process children: {}'.format(err))
        return []