comparison mercurial/worker.py @ 42455:5ca136bbd3f6

worker: support parallelization of functions with return values

Currently worker supports running functions that return a progress
iterator. Generalize it to handle function that return a progress
iterator then a return value.

It's unused in this commit, but will be used in the next one.

Differential Revision: https://phab.mercurial-scm.org/D6515

42454:0eb8c61c306b

    linear = costperop * nops
    workers = _numworkers(ui)
    benefit = linear - (_STARTUP_COST * workers + linear / workers)
    return benefit >= 0.15

def worker(ui, costperarg, func, staticargs, args, threadsafe=True):
    '''run a function, possibly in parallel in multiple worker
    processes.

    returns a progress iterator

    costperarg - cost of a single task

    func - function to run

    staticargs - arguments to pass to every invocation of the function

    args - arguments to split into chunks, to pass to individual
    workers

    threadsafe - whether work items are thread safe and can be executed using
    a thread-based worker. Should be disabled for CPU heavy tasks that don't
    release the GIL.
    '''
    enabled = ui.configbool('worker', 'enabled')
    if enabled and worthwhile(ui, costperarg, len(args), threadsafe=threadsafe):
        return _platformworker(ui, func, staticargs, args)
    return func(*staticargs + (args,))

def _posixworker(ui, func, staticargs, args):
    workers = _numworkers(ui)
    oldhandler = signal.getsignal(signal.SIGINT)
    signal.signal(signal.SIGINT, signal.SIG_IGN)
    pids, problem = set(), [0]
    def killworkers():

            killworkers()
    oldchldhandler = signal.signal(signal.SIGCHLD, sigchldhandler)
    ui.flush()
    parentpid = os.getpid()
    pipes = []
    for pargs in partition(args, workers):
        # Every worker gets its own pipe to send results on, so we don't have to
        # implement atomic writes larger than PIPE_BUF. Each forked process has
        # its own pipe's descriptors in the local variables, and the parent
        # process has the full list of pipe descriptors (and it doesn't really
        # care what order they're in).
        rfd, wfd = os.pipe()
        pipes.append((rfd, wfd))
        # make sure we use os._exit in all worker code paths. otherwise the
        # worker may do some clean-ups which could cause surprises like
        # deadlock. see sshpeer.cleanup for example.
        # override error handling *before* fork. this is necessary because
        # exception (signal) may arrive after fork, before "pid =" assignment

                    for r, w in pipes[:-1]:
                        os.close(r)
                        os.close(w)
                    os.close(rfd)
                    for result in func(*(staticargs + (pargs,))):
                        os.write(wfd, util.pickle.dumps(result))
                    return 0

                ret = scmutil.callcatch(ui, workerfunc)
        except: # parent re-raises, child never returns
            if os.getpid() == parentpid:

    try:
        openpipes = len(pipes)
        while openpipes > 0:
            for key, events in selector.select():
                try:
                    yield util.pickle.load(key.fileobj)
                except EOFError:
                    selector.unregister(key.fileobj)
                    key.fileobj.close()
                    openpipes -= 1
                except IOError as e:

    status = cleanup()
    if status:
        if status < 0:
            os.kill(os.getpid(), -status)
        sys.exit(status)

def _posixexitstatus(code):
    '''convert a posix exit status into the same form returned by
    os.spawnv

    if os.WIFEXITED(code):
        return os.WEXITSTATUS(code)
    elif os.WIFSIGNALED(code):
        return -os.WTERMSIG(code)

def _windowsworker(ui, func, staticargs, args):
    class Worker(threading.Thread):
        def __init__(self, taskqueue, resultqueue, func, staticargs, *args,
                     **kwargs):
            threading.Thread.__init__(self, *args, **kwargs)
            self._taskqueue = taskqueue

        def run(self):
            try:
                while not self._taskqueue.empty():
                    try:
                        args = self._taskqueue.get_nowait()
                        for res in self._func(*self._staticargs + (args,)):
                            self._resultqueue.put(res)
                            # threading doesn't provide a native way to
                            # interrupt execution. handle it manually at every
                            # iteration.
                            if self._interrupted:
                                return

                return

    workers = _numworkers(ui)
    resultqueue = pycompat.queue.Queue()
    taskqueue = pycompat.queue.Queue()
    # partition work to more pieces than workers to minimize the chance
    # of uneven distribution of large tasks between the workers
    for pargs in partition(args, workers * 20):
        taskqueue.put(pargs)
    for _i in range(workers):
        t = Worker(taskqueue, resultqueue, func, staticargs)
        threads.append(t)
        t.start()
    try:
        while len(threads) > 0:
            while not resultqueue.empty():
                yield resultqueue.get()
            threads[0].join(0.05)
            finishedthreads = [_t for _t in threads if not _t.is_alive()]
            for t in finishedthreads:
                if t.exception is not None:
                    raise t.exception
                threads.remove(t)
    except (Exception, KeyboardInterrupt): # re-raises
        trykillworkers()
        raise
    while not resultqueue.empty():
        yield resultqueue.get()

if pycompat.iswindows:
    _platformworker = _windowsworker
else:
    _platformworker = _posixworker

42455:5ca136bbd3f6

    linear = costperop * nops
    workers = _numworkers(ui)
    benefit = linear - (_STARTUP_COST * workers + linear / workers)
    return benefit >= 0.15

def worker(ui, costperarg, func, staticargs, args, hasretval=False,
           threadsafe=True):
    '''run a function, possibly in parallel in multiple worker
    processes.

    returns a progress iterator

    costperarg - cost of a single task

    func - function to run. It is expected to return a progress iterator.

    staticargs - arguments to pass to every invocation of the function

    args - arguments to split into chunks, to pass to individual
    workers

    hasretval - when True, func and the current function return an progress
    iterator then a list (encoded as an iterator that yield many (False, ..)
    then a (True, list)). The resulting list is in the natural order.

    threadsafe - whether work items are thread safe and can be executed using
    a thread-based worker. Should be disabled for CPU heavy tasks that don't
    release the GIL.
    '''
    enabled = ui.configbool('worker', 'enabled')
    if enabled and worthwhile(ui, costperarg, len(args), threadsafe=threadsafe):
        return _platformworker(ui, func, staticargs, args, hasretval)
    return func(*staticargs + (args,))

def _posixworker(ui, func, staticargs, args, hasretval):
    workers = _numworkers(ui)
    oldhandler = signal.getsignal(signal.SIGINT)
    signal.signal(signal.SIGINT, signal.SIG_IGN)
    pids, problem = set(), [0]
    def killworkers():

            killworkers()
    oldchldhandler = signal.signal(signal.SIGCHLD, sigchldhandler)
    ui.flush()
    parentpid = os.getpid()
    pipes = []
    retvals = []
    for i, pargs in enumerate(partition(args, workers)):
        # Every worker gets its own pipe to send results on, so we don't have to
        # implement atomic writes larger than PIPE_BUF. Each forked process has
        # its own pipe's descriptors in the local variables, and the parent
        # process has the full list of pipe descriptors (and it doesn't really
        # care what order they're in).
        rfd, wfd = os.pipe()
        pipes.append((rfd, wfd))
        retvals.append(None)
        # make sure we use os._exit in all worker code paths. otherwise the
        # worker may do some clean-ups which could cause surprises like
        # deadlock. see sshpeer.cleanup for example.
        # override error handling *before* fork. this is necessary because
        # exception (signal) may arrive after fork, before "pid =" assignment

                    for r, w in pipes[:-1]:
                        os.close(r)
                        os.close(w)
                    os.close(rfd)
                    for result in func(*(staticargs + (pargs,))):
                        os.write(wfd, util.pickle.dumps((i, result)))
                    return 0

                ret = scmutil.callcatch(ui, workerfunc)
        except: # parent re-raises, child never returns
            if os.getpid() == parentpid:

    try:
        openpipes = len(pipes)
        while openpipes > 0:
            for key, events in selector.select():
                try:
                    i, res = util.pickle.load(key.fileobj)
                    if hasretval and res[0]:
                        retvals[i] = res[1]
                    else:
                        yield res
                except EOFError:
                    selector.unregister(key.fileobj)
                    key.fileobj.close()
                    openpipes -= 1
                except IOError as e:

    status = cleanup()
    if status:
        if status < 0:
            os.kill(os.getpid(), -status)
        sys.exit(status)
    if hasretval:
        yield True, sum(retvals, [])

def _posixexitstatus(code):
    '''convert a posix exit status into the same form returned by
    os.spawnv

    if os.WIFEXITED(code):
        return os.WEXITSTATUS(code)
    elif os.WIFSIGNALED(code):
        return -os.WTERMSIG(code)

def _windowsworker(ui, func, staticargs, args, hasretval):
    class Worker(threading.Thread):
        def __init__(self, taskqueue, resultqueue, func, staticargs, *args,
                     **kwargs):
            threading.Thread.__init__(self, *args, **kwargs)
            self._taskqueue = taskqueue

        def run(self):
            try:
                while not self._taskqueue.empty():
                    try:
                        i, args = self._taskqueue.get_nowait()
                        for res in self._func(*self._staticargs + (args,)):
                            self._resultqueue.put((i, res))
                            # threading doesn't provide a native way to
                            # interrupt execution. handle it manually at every
                            # iteration.
                            if self._interrupted:
                                return

                return

    workers = _numworkers(ui)
    resultqueue = pycompat.queue.Queue()
    taskqueue = pycompat.queue.Queue()
    retvals = []
    # partition work to more pieces than workers to minimize the chance
    # of uneven distribution of large tasks between the workers
    for pargs in enumerate(partition(args, workers * 20)):
        retvals.append(None)
        taskqueue.put(pargs)
    for _i in range(workers):
        t = Worker(taskqueue, resultqueue, func, staticargs)
        threads.append(t)
        t.start()
    try:
        while len(threads) > 0:
            while not resultqueue.empty():
                (i, res) = resultqueue.get()
                if hasretval and res[0]:
                    retvals[i] = res[1]
                else:
                    yield res
            threads[0].join(0.05)
            finishedthreads = [_t for _t in threads if not _t.is_alive()]
            for t in finishedthreads:
                if t.exception is not None:
                    raise t.exception
                threads.remove(t)
    except (Exception, KeyboardInterrupt): # re-raises
        trykillworkers()
        raise
    while not resultqueue.empty():
        (i, res) = resultqueue.get()
        if hasretval and res[0]:
            retvals[i] = res[1]
        else:
            yield res
    if hasretval:
        yield True, sum(retvals, [])

if pycompat.iswindows:
    _platformworker = _windowsworker
else:
    _platformworker = _posixworker