import fcntl, os, select, time
from subprocess import Popen, PIPE

# Run a series of subprocesses. Try to keep up to a certain number going in
# parallel at any given time. Enforce time limits.
#
# This is implemented using non-blocking I/O, and so is Unix-specific.
#
# We assume that, if a task closes its standard error, then it's safe to
# wait for it to terminate. So an ill-behaved task that closes its standard
# output and then hangs will hang us, as well. However, as it takes special
# effort to close one's standard output, this seems unlikely to be a
# problem in practice.
class TaskPool(object):

    # A task we should run in a subprocess. Users should subclass this and
    # fill in the methods as given.
    class Task(object):
        def __init__(self):
            self.pipe = None
            self.start_time = None

        # Record that this task is running, with |pipe| as its Popen object,
        # and should time out at |deadline|.
        def start(self, pipe, deadline):
            self.pipe = pipe
            self.deadline = deadline

        # Return a shell command (a string or sequence of arguments) to be
        # passed to Popen to run the task. The command will be given
        # /dev/null as its standard input, and pipes as its standard output
        # and error.
        def cmd(self):
            raise NotImplementedError

        # TaskPool calls this method to report that the process wrote
        # |string| to its standard output.
        def onStdout(self, string):
            raise NotImplementedError

        # TaskPool calls this method to report that the process wrote
        # |string| to its standard error.
        def onStderr(self, string):
            raise NotImplementedError

        # TaskPool calls this method to report that the process terminated,
        # yielding |returncode|.
        def onFinished(self, returncode):
            raise NotImplementedError

        # TaskPool calls this method to report that the process timed out and
        # was killed.
        def onTimeout(self):
            raise NotImplementedError

    # If a task output handler (onStdout, onStderr) throws this, we terminate
    # the task.
    class TerminateTask(Exception):
        pass

    def __init__(self, tasks, cwd='.', job_limit=4, timeout=150):
        self.pending = iter(tasks)
        self.cwd = cwd
        self.job_limit = job_limit
        self.timeout = timeout
        self.next_pending = self.get_next_pending()

    # Set self.next_pending to the next task that has not yet been executed.
    def get_next_pending(self):
        try:
            return self.pending.next()
        except StopIteration:
            return None

    def run_all(self):
        # The currently running tasks: a set of Task instances.
        running = set()
        with open(os.devnull, 'r') as devnull:
            while True:
                while len(running) < self.job_limit and self.next_pending:
                    t = self.next_pending
                    p = Popen(t.cmd(), bufsize=16384,
                              stdin=devnull, stdout=PIPE, stderr=PIPE,
                              cwd=self.cwd)

                    # Put the stdout and stderr pipes in non-blocking mode. See
                    # the post-'select' code below for details.
                    flags = fcntl.fcntl(p.stdout, fcntl.F_GETFL)
                    fcntl.fcntl(p.stdout, fcntl.F_SETFL, flags | os.O_NONBLOCK)
                    flags = fcntl.fcntl(p.stderr, fcntl.F_GETFL)
                    fcntl.fcntl(p.stderr, fcntl.F_SETFL, flags | os.O_NONBLOCK)

                    t.start(p, time.time() + self.timeout)
                    running.add(t)
                    self.next_pending = self.get_next_pending()

                # If we have no tasks running, and the above wasn't able to
                # start any new ones, then we must be done!
                if not running:
                    break

                # How many seconds do we have until the earliest deadline?
                now = time.time()
                secs_to_next_deadline = max(min([t.deadline for t in running]) - now, 0)

                # Wait for output or a timeout.
                stdouts_and_stderrs = ([t.pipe.stdout for t in running]
                                     + [t.pipe.stderr for t in running])
                (readable,w,x) = select.select(stdouts_and_stderrs, [], [], secs_to_next_deadline)
                finished = set()
                terminate = set()
                for t in running:
                    # Since we've placed the pipes in non-blocking mode, these
                    # 'read's will simply return as many bytes as are available,
                    # rather than blocking until they have accumulated the full
                    # amount requested (or reached EOF). The 'read's should
                    # never throw, since 'select' has told us there was
                    # something available.
                    if t.pipe.stdout in readable:
                        output = t.pipe.stdout.read(16384)
                        if output != "":
                            try:
                                t.onStdout(output)
                            except TerminateTask:
                                terminate.add(t)
                    if t.pipe.stderr in readable:
                        output = t.pipe.stderr.read(16384)
                        if output != "":
                            try:
                                t.onStderr(output)
                            except TerminateTask:
                                terminate.add(t)
                        else:
                            # We assume that, once a task has closed its stderr,
                            # it will soon terminate. If a task closes its
                            # stderr and then hangs, we'll hang too, here.
                            t.pipe.wait()
                            t.onFinished(t.pipe.returncode)
                            finished.add(t)
                # Remove the finished tasks from the running set. (Do this here
                # to avoid mutating the set while iterating over it.)
                running -= finished

                # Terminate any tasks whose handlers have asked us to do so.
                for t in terminate:
                    t.pipe.terminate()
                    t.pipe.wait()
                    running.remove(t)

                # Terminate any tasks which have missed their deadline.
                finished = set()
                for t in running:
                    if now >= t.deadline:
                        t.pipe.terminate()
                        t.pipe.wait()
                        t.onTimeout()
                        finished.add(t)
                # Remove the finished tasks from the running set. (Do this here
                # to avoid mutating the set while iterating over it.)
                running -= finished
        return None

def get_cpu_count():
    """
    Guess at a reasonable parallelism count to set as the default for the
    current machine and run.
    """
    # Python 2.6+
    try:
        import multiprocessing
        return multiprocessing.cpu_count()
    except (ImportError,NotImplementedError):
        pass

    # POSIX
    try:
        res = int(os.sysconf('SC_NPROCESSORS_ONLN'))
        if res > 0:
            return res
    except (AttributeError,ValueError):
        pass

    # Windows
    try:
        res = int(os.environ['NUMBER_OF_PROCESSORS'])
        if res > 0:
            return res
    except (KeyError, ValueError):
        pass

    return 1

if __name__ == '__main__':
    # Test TaskPool by using it to implement the unique 'sleep sort' algorithm.
    def sleep_sort(ns, timeout):
        sorted=[]
        class SortableTask(TaskPool.Task):
            def __init__(self, n):
                super(SortableTask, self).__init__()
                self.n = n
            def start(self, pipe, deadline):
                super(SortableTask, self).start(pipe, deadline)
            def cmd(self):
                return ['sh', '-c', 'echo out; sleep %d; echo err>&2' % (self.n,)]
            def onStdout(self, text):
                print '%d stdout: %r' % (self.n, text)
            def onStderr(self, text):
                print '%d stderr: %r' % (self.n, text)
            def onFinished(self, returncode):
                print '%d (rc=%d)' % (self.n, returncode)
                sorted.append(self.n)
            def onTimeout(self):
                print '%d timed out' % (self.n,)

        p = TaskPool([SortableTask(_) for _ in ns], job_limit=len(ns), timeout=timeout)
        p.run_all()
        return sorted

    print repr(sleep_sort([1,1,2,3,5,8,13,21,34], 15))