mercurial/worker.py

43029:c5dc122fdc2b

import threading
import time

try:
    import selectors
    selectors.BaseSelector
except ImportError:
    from .thirdparty import selectors2 as selectors

from .i18n import _

    pycompat,
    scmutil,
    util,
)

def countcpus():
    '''try to count the number of CPUs on the system'''

    # posix
    try:

    except (AttributeError, ValueError):
        pass

    # windows
    try:
        n = int(encoding.environ['NUMBER_OF_PROCESSORS'])
        if n > 0:
            return n
    except (KeyError, ValueError):
        pass

    return 1

def _numworkers(ui):
    s = ui.config('worker', 'numcpus')
    if s:
        try:
            n = int(s)
            if n >= 1:
                return n
        except ValueError:
            raise error.Abort(_('number of cpus must be an integer'))
    return min(max(countcpus(), 4), 32)

if pycompat.isposix or pycompat.iswindows:
    _STARTUP_COST = 0.01
    # The Windows worker is thread based. If tasks are CPU bound, threads
    # in the presence of the GIL result in excessive context switching and

    _DISALLOW_THREAD_UNSAFE = pycompat.iswindows
else:
    _STARTUP_COST = 1e30
    _DISALLOW_THREAD_UNSAFE = False

def worthwhile(ui, costperop, nops, threadsafe=True):
    '''try to determine whether the benefit of multiple processes can
    outweigh the cost of starting them'''

    if not threadsafe and _DISALLOW_THREAD_UNSAFE:

    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

    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():
        # unregister SIGCHLD handler as all children will be killed. This
        # function shouldn't be interrupted by another SIGCHLD; otherwise pids
        # could be updated while iterating, which would cause inconsistency.
        signal.signal(signal.SIGCHLD, oldchldhandler)

            try:
                os.kill(p, signal.SIGTERM)
            except OSError as err:
                if err.errno != errno.ESRCH:
                    raise
    def waitforworkers(blocking=True):
        for pid in pids.copy():
            p = st = 0
            while True:
                try:

                continue
            pids.discard(p)
            st = _exitstatus(st)
            if st and not problem[0]:
                problem[0] = st
    def sigchldhandler(signum, frame):
        waitforworkers(blocking=False)
        if problem[0]:
            killworkers()
    oldchldhandler = signal.signal(signal.SIGCHLD, sigchldhandler)
    ui.flush()
    parentpid = os.getpid()
    pipes = []
    retval = {}

                    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:
                raise
            exctype = sys.exc_info()[0]
            force = not issubclass(exctype, KeyboardInterrupt)
            ui.traceback(force=force)
        finally:
            if os.getpid() != parentpid:
                try:
                    ui.flush()
                except: # never returns, no re-raises
                    pass
                finally:
                    os._exit(ret & 255)
        pids.add(pid)
    selector = selectors.DefaultSelector()
    for rfd, wfd in pipes:
        os.close(wfd)
        selector.register(os.fdopen(rfd, r'rb', 0), selectors.EVENT_READ)
    def cleanup():
        signal.signal(signal.SIGINT, oldhandler)
        waitforworkers()
        signal.signal(signal.SIGCHLD, oldchldhandler)
        selector.close()
        return problem[0]
    try:
        openpipes = len(pipes)
        while openpipes > 0:
            for key, events in selector.select():
                try:

                    openpipes -= 1
                except IOError as e:
                    if e.errno == errno.EINTR:
                        continue
                    raise
    except: # re-raises
        killworkers()
        cleanup()
        raise
    status = cleanup()
    if status:

            os.kill(os.getpid(), -status)
        sys.exit(status)
    if hasretval:
        yield True, retval

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

    returns None if the process was stopped instead of exiting'''
    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
            self._resultqueue = resultqueue
            self._func = func
            self._staticargs = staticargs

                # it as if the func was running without workers.
                self.exception = e
                raise

    threads = []
    def trykillworkers():
        # Allow up to 1 second to clean worker threads nicely
        cleanupend = time.time() + 1
        for t in threads:
            t.interrupt()

            if t.is_alive():
                # pass over the workers joining failure. it is more
                # important to surface the inital exception than the
                # fact that one of workers may be processing a large
                # task and does not get to handle the interruption.
                ui.warn(_("failed to kill worker threads while "
                          "handling an exception\n"))
                return

    workers = _numworkers(ui)
    resultqueue = pycompat.queue.Queue()
    taskqueue = pycompat.queue.Queue()

            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():
        res = resultqueue.get()
        if hasretval and res[0]:

        else:
            yield res
    if hasretval:
        yield True, retval

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

def partition(lst, nslices):
    '''partition a list into N slices of roughly equal size

    The current strategy takes every Nth element from the input. If

43306:59338f956109

import threading
import time

try:
    import selectors

    selectors.BaseSelector
except ImportError:
    from .thirdparty import selectors2 as selectors

from .i18n import _

    pycompat,
    scmutil,
    util,
)


def countcpus():
    '''try to count the number of CPUs on the system'''

    # posix
    try:

    except (AttributeError, ValueError):
        pass

    # windows
    try:
        n = int(encoding.environ[b'NUMBER_OF_PROCESSORS'])
        if n > 0:
            return n
    except (KeyError, ValueError):
        pass

    return 1


def _numworkers(ui):
    s = ui.config(b'worker', b'numcpus')
    if s:
        try:
            n = int(s)
            if n >= 1:
                return n
        except ValueError:
            raise error.Abort(_(b'number of cpus must be an integer'))
    return min(max(countcpus(), 4), 32)


if pycompat.isposix or pycompat.iswindows:
    _STARTUP_COST = 0.01
    # The Windows worker is thread based. If tasks are CPU bound, threads
    # in the presence of the GIL result in excessive context switching and

    _DISALLOW_THREAD_UNSAFE = pycompat.iswindows
else:
    _STARTUP_COST = 1e30
    _DISALLOW_THREAD_UNSAFE = False


def worthwhile(ui, costperop, nops, threadsafe=True):
    '''try to determine whether the benefit of multiple processes can
    outweigh the cost of starting them'''

    if not threadsafe and _DISALLOW_THREAD_UNSAFE:

    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

    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(b'worker', b'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():
        # unregister SIGCHLD handler as all children will be killed. This
        # function shouldn't be interrupted by another SIGCHLD; otherwise pids
        # could be updated while iterating, which would cause inconsistency.
        signal.signal(signal.SIGCHLD, oldchldhandler)

            try:
                os.kill(p, signal.SIGTERM)
            except OSError as err:
                if err.errno != errno.ESRCH:
                    raise

    def waitforworkers(blocking=True):
        for pid in pids.copy():
            p = st = 0
            while True:
                try:

                continue
            pids.discard(p)
            st = _exitstatus(st)
            if st and not problem[0]:
                problem[0] = st

    def sigchldhandler(signum, frame):
        waitforworkers(blocking=False)
        if problem[0]:
            killworkers()

    oldchldhandler = signal.signal(signal.SIGCHLD, sigchldhandler)
    ui.flush()
    parentpid = os.getpid()
    pipes = []
    retval = {}

                    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:
                raise
            exctype = sys.exc_info()[0]
            force = not issubclass(exctype, KeyboardInterrupt)
            ui.traceback(force=force)
        finally:
            if os.getpid() != parentpid:
                try:
                    ui.flush()
                except:  # never returns, no re-raises
                    pass
                finally:
                    os._exit(ret & 255)
        pids.add(pid)
    selector = selectors.DefaultSelector()
    for rfd, wfd in pipes:
        os.close(wfd)
        selector.register(os.fdopen(rfd, r'rb', 0), selectors.EVENT_READ)

    def cleanup():
        signal.signal(signal.SIGINT, oldhandler)
        waitforworkers()
        signal.signal(signal.SIGCHLD, oldchldhandler)
        selector.close()
        return problem[0]

    try:
        openpipes = len(pipes)
        while openpipes > 0:
            for key, events in selector.select():
                try:

                    openpipes -= 1
                except IOError as e:
                    if e.errno == errno.EINTR:
                        continue
                    raise
    except:  # re-raises
        killworkers()
        cleanup()
        raise
    status = cleanup()
    if status:

            os.kill(os.getpid(), -status)
        sys.exit(status)
    if hasretval:
        yield True, retval


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

    returns None if the process was stopped instead of exiting'''
    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
            self._resultqueue = resultqueue
            self._func = func
            self._staticargs = staticargs

                # it as if the func was running without workers.
                self.exception = e
                raise

    threads = []

    def trykillworkers():
        # Allow up to 1 second to clean worker threads nicely
        cleanupend = time.time() + 1
        for t in threads:
            t.interrupt()

            if t.is_alive():
                # pass over the workers joining failure. it is more
                # important to surface the inital exception than the
                # fact that one of workers may be processing a large
                # task and does not get to handle the interruption.
                ui.warn(
                    _(
                        b"failed to kill worker threads while "
                        b"handling an exception\n"
                    )
                )
                return

    workers = _numworkers(ui)
    resultqueue = pycompat.queue.Queue()
    taskqueue = pycompat.queue.Queue()

            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():
        res = resultqueue.get()
        if hasretval and res[0]:

        else:
            yield res
    if hasretval:
        yield True, retval


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


def partition(lst, nslices):
    '''partition a list into N slices of roughly equal size

    The current strategy takes every Nth element from the input. If