#!/usr/bin/env python
"""
Use to put a timeout on the length of time a script can run. This is
especially useful for checking for scripts that hang.

Usage
=====

NOTE: Make sure you run ``pip install -r requirements-dev.txt`` before running.

To use the script, run::

    ./timeout "./my-script-to-run" --timeout-after 5

"""
import argparse
import os
import subprocess
import sys
import time

import psutil


class TimeoutException(Exception):
    def __init__(self, timeout_len):
        msg = 'Script failed to complete within %s seconds' % timeout_len
        Exception.__init__(self, msg)


def timeout(args):
    parent_pid = os.getpid()
    child_p = run_script(args)
    try:
        run_timeout(child_p.pid, args.timeout_after)
    except (TimeoutException, KeyboardInterrupt) as e:
        proc = psutil.Process(parent_pid)
        procs = proc.children(recursive=True)

        for child in procs:
            child.terminate()

        gone, alive = psutil.wait_procs(procs, timeout=1)
        for child in alive:
            child.kill()
        raise e


def run_timeout(pid, timeout_len):
    p = psutil.Process(pid)
    start_time = time.time()
    while p.is_running():
        if p.status() == psutil.STATUS_ZOMBIE:
            p.kill()
            break
        current_time = time.time()
        # Raise a timeout if the duration of the process is longer than
        # the desired timeout.
        if current_time - start_time > timeout_len:
            raise TimeoutException(timeout_len)
        time.sleep(1)


def run_script(args):
    return subprocess.Popen(args.script, shell=True)


def main():
    parser = argparse.ArgumentParser(usage=__doc__)
    parser.add_argument('script', help='The script to run for benchmarking')
    parser.add_argument(
        '--timeout-after',
        required=True,
        type=float,
        help=(
            'The length of time in seconds allowed for the script to run '
            'before it time\'s out.'
        ),
    )
    args = parser.parse_args()
    return timeout(args)


if __name__ == '__main__':
    sys.exit(main())