view hgext/narrow/narrowbundle2.py @ 43124:100e7e0cdaad

perf: use `setup` function in `perfdirstatedirs`

The command seems to pre-date the introduction of the `setup` support in timer.
We move the line that is obviously about benchmark setup in such `setup`
function.

# narrowbundle2.py - bundle2 extensions for narrow repository support
#
# Copyright 2017 Google, Inc.
#
# This software may be used and distributed according to the terms of the
# GNU General Public License version 2 or any later version.

from __future__ import absolute_import

import errno
import struct

from mercurial.i18n import _
from mercurial.node import (
    bin,
    nullid,
)
from mercurial import (
    bundle2,
    changegroup,
    error,
    exchange,
    localrepo,
    narrowspec,
    repair,
    util,
    wireprototypes,
)
from mercurial.interfaces import repository
from mercurial.utils import stringutil

_NARROWACL_SECTION = b'narrowacl'
_CHANGESPECPART = b'narrow:changespec'
_RESSPECS = b'narrow:responsespec'
_SPECPART = b'narrow:spec'
_SPECPART_INCLUDE = b'include'
_SPECPART_EXCLUDE = b'exclude'
_KILLNODESIGNAL = b'KILL'
_DONESIGNAL = b'DONE'
_ELIDEDCSHEADER = b'>20s20s20sl'  # cset id, p1, p2, len(text)
_ELIDEDMFHEADER = b'>20s20s20s20sl'  # manifest id, p1, p2, link id, len(text)
_CSHEADERSIZE = struct.calcsize(_ELIDEDCSHEADER)
_MFHEADERSIZE = struct.calcsize(_ELIDEDMFHEADER)

# Serve a changegroup for a client with a narrow clone.
def getbundlechangegrouppart_narrow(
    bundler,
    repo,
    source,
    bundlecaps=None,
    b2caps=None,
    heads=None,
    common=None,
    **kwargs
):
    assert repo.ui.configbool(b'experimental', b'narrowservebrokenellipses')

    cgversions = b2caps.get(b'changegroup')
    cgversions = [
        v
        for v in cgversions
        if v in changegroup.supportedoutgoingversions(repo)
    ]
    if not cgversions:
        raise ValueError(_(b'no common changegroup version'))
    version = max(cgversions)

    oldinclude = sorted(filter(bool, kwargs.get(r'oldincludepats', [])))
    oldexclude = sorted(filter(bool, kwargs.get(r'oldexcludepats', [])))
    newinclude = sorted(filter(bool, kwargs.get(r'includepats', [])))
    newexclude = sorted(filter(bool, kwargs.get(r'excludepats', [])))
    known = {bin(n) for n in kwargs.get(r'known', [])}
    generateellipsesbundle2(
        bundler,
        repo,
        oldinclude,
        oldexclude,
        newinclude,
        newexclude,
        version,
        common,
        heads,
        known,
        kwargs.get(r'depth', None),
    )


def generateellipsesbundle2(
    bundler,
    repo,
    oldinclude,
    oldexclude,
    newinclude,
    newexclude,
    version,
    common,
    heads,
    known,
    depth,
):
    newmatch = narrowspec.match(
        repo.root, include=newinclude, exclude=newexclude
    )
    if depth is not None:
        depth = int(depth)
        if depth < 1:
            raise error.Abort(_(b'depth must be positive, got %d') % depth)

    heads = set(heads or repo.heads())
    common = set(common or [nullid])
    if known and (oldinclude != newinclude or oldexclude != newexclude):
        # Steps:
        # 1. Send kill for "$known & ::common"
        #
        # 2. Send changegroup for ::common
        #
        # 3. Proceed.
        #
        # In the future, we can send kills for only the specific
        # nodes we know should go away or change shape, and then
        # send a data stream that tells the client something like this:
        #
        # a) apply this changegroup
        # b) apply nodes XXX, YYY, ZZZ that you already have
        # c) goto a
        #
        # until they've built up the full new state.
        # Convert to revnums and intersect with "common". The client should
        # have made it a subset of "common" already, but let's be safe.
        known = set(repo.revs(b"%ln & ::%ln", known, common))
        # TODO: we could send only roots() of this set, and the
        # list of nodes in common, and the client could work out
        # what to strip, instead of us explicitly sending every
        # single node.
        deadrevs = known

        def genkills():
            for r in deadrevs:
                yield _KILLNODESIGNAL
                yield repo.changelog.node(r)
            yield _DONESIGNAL

        bundler.newpart(_CHANGESPECPART, data=genkills())
        newvisit, newfull, newellipsis = exchange._computeellipsis(
            repo, set(), common, known, newmatch
        )
        if newvisit:
            packer = changegroup.getbundler(
                version,
                repo,
                matcher=newmatch,
                ellipses=True,
                shallow=depth is not None,
                ellipsisroots=newellipsis,
                fullnodes=newfull,
            )
            cgdata = packer.generate(common, newvisit, False, b'narrow_widen')

            part = bundler.newpart(b'changegroup', data=cgdata)
            part.addparam(b'version', version)
            if b'treemanifest' in repo.requirements:
                part.addparam(b'treemanifest', b'1')

    visitnodes, relevant_nodes, ellipsisroots = exchange._computeellipsis(
        repo, common, heads, set(), newmatch, depth=depth
    )

    repo.ui.debug(b'Found %d relevant revs\n' % len(relevant_nodes))
    if visitnodes:
        packer = changegroup.getbundler(
            version,
            repo,
            matcher=newmatch,
            ellipses=True,
            shallow=depth is not None,
            ellipsisroots=ellipsisroots,
            fullnodes=relevant_nodes,
        )
        cgdata = packer.generate(common, visitnodes, False, b'narrow_widen')

        part = bundler.newpart(b'changegroup', data=cgdata)
        part.addparam(b'version', version)
        if b'treemanifest' in repo.requirements:
            part.addparam(b'treemanifest', b'1')


@bundle2.parthandler(_SPECPART, (_SPECPART_INCLUDE, _SPECPART_EXCLUDE))
def _handlechangespec_2(op, inpart):
    # XXX: This bundle2 handling is buggy and should be removed after hg5.2 is
    # released. New servers will send a mandatory bundle2 part named
    # 'Narrowspec' and will send specs as data instead of params.
    # Refer to issue5952 and 6019
    includepats = set(inpart.params.get(_SPECPART_INCLUDE, b'').splitlines())
    excludepats = set(inpart.params.get(_SPECPART_EXCLUDE, b'').splitlines())
    narrowspec.validatepatterns(includepats)
    narrowspec.validatepatterns(excludepats)

    if not repository.NARROW_REQUIREMENT in op.repo.requirements:
        op.repo.requirements.add(repository.NARROW_REQUIREMENT)
        op.repo._writerequirements()
    op.repo.setnarrowpats(includepats, excludepats)
    narrowspec.copytoworkingcopy(op.repo)


@bundle2.parthandler(_RESSPECS)
def _handlenarrowspecs(op, inpart):
    data = inpart.read()
    inc, exc = data.split(b'\0')
    includepats = set(inc.splitlines())
    excludepats = set(exc.splitlines())
    narrowspec.validatepatterns(includepats)
    narrowspec.validatepatterns(excludepats)

    if repository.NARROW_REQUIREMENT not in op.repo.requirements:
        op.repo.requirements.add(repository.NARROW_REQUIREMENT)
        op.repo._writerequirements()
    op.repo.setnarrowpats(includepats, excludepats)
    narrowspec.copytoworkingcopy(op.repo)


@bundle2.parthandler(_CHANGESPECPART)
def _handlechangespec(op, inpart):
    repo = op.repo
    cl = repo.changelog

    # changesets which need to be stripped entirely. either they're no longer
    # needed in the new narrow spec, or the server is sending a replacement
    # in the changegroup part.
    clkills = set()

    # A changespec part contains all the updates to ellipsis nodes
    # that will happen as a result of widening or narrowing a
    # repo. All the changes that this block encounters are ellipsis
    # nodes or flags to kill an existing ellipsis.
    chunksignal = changegroup.readexactly(inpart, 4)
    while chunksignal != _DONESIGNAL:
        if chunksignal == _KILLNODESIGNAL:
            # a node used to be an ellipsis but isn't anymore
            ck = changegroup.readexactly(inpart, 20)
            if cl.hasnode(ck):
                clkills.add(ck)
        else:
            raise error.Abort(
                _(b'unexpected changespec node chunk type: %s') % chunksignal
            )
        chunksignal = changegroup.readexactly(inpart, 4)

    if clkills:
        # preserve bookmarks that repair.strip() would otherwise strip
        op._bookmarksbackup = repo._bookmarks

        class dummybmstore(dict):
            def applychanges(self, repo, tr, changes):
                pass

        localrepo.localrepository._bookmarks.set(repo, dummybmstore())
        chgrpfile = repair.strip(
            op.ui, repo, list(clkills), backup=True, topic=b'widen'
        )
        if chgrpfile:
            op._widen_uninterr = repo.ui.uninterruptible()
            op._widen_uninterr.__enter__()
            # presence of _widen_bundle attribute activates widen handler later
            op._widen_bundle = chgrpfile
    # Set the new narrowspec if we're widening. The setnewnarrowpats() method
    # will currently always be there when using the core+narrowhg server, but
    # other servers may include a changespec part even when not widening (e.g.
    # because we're deepening a shallow repo).
    if util.safehasattr(repo, 'setnewnarrowpats'):
        repo.setnewnarrowpats()


def handlechangegroup_widen(op, inpart):
    """Changegroup exchange handler which restores temporarily-stripped nodes"""
    # We saved a bundle with stripped node data we must now restore.
    # This approach is based on mercurial/repair.py@6ee26a53c111.
    repo = op.repo
    ui = op.ui

    chgrpfile = op._widen_bundle
    del op._widen_bundle
    vfs = repo.vfs

    ui.note(_(b"adding branch\n"))
    f = vfs.open(chgrpfile, b"rb")
    try:
        gen = exchange.readbundle(ui, f, chgrpfile, vfs)
        # silence internal shuffling chatter
        override = {(b'ui', b'quiet'): True}
        if ui.verbose:
            override = {}
        with ui.configoverride(override):
            if isinstance(gen, bundle2.unbundle20):
                with repo.transaction(b'strip') as tr:
                    bundle2.processbundle(repo, gen, lambda: tr)
            else:
                gen.apply(
                    repo, b'strip', b'bundle:' + vfs.join(chgrpfile), True
                )
    finally:
        f.close()

    # remove undo files
    for undovfs, undofile in repo.undofiles():
        try:
            undovfs.unlink(undofile)
        except OSError as e:
            if e.errno != errno.ENOENT:
                ui.warn(
                    _(b'error removing %s: %s\n')
                    % (undovfs.join(undofile), stringutil.forcebytestr(e))
                )

    # Remove partial backup only if there were no exceptions
    op._widen_uninterr.__exit__(None, None, None)
    vfs.unlink(chgrpfile)


def setup():
    """Enable narrow repo support in bundle2-related extension points."""
    getbundleargs = wireprototypes.GETBUNDLE_ARGUMENTS

    getbundleargs[b'narrow'] = b'boolean'
    getbundleargs[b'depth'] = b'plain'
    getbundleargs[b'oldincludepats'] = b'csv'
    getbundleargs[b'oldexcludepats'] = b'csv'
    getbundleargs[b'known'] = b'csv'

    # Extend changegroup serving to handle requests from narrow clients.
    origcgfn = exchange.getbundle2partsmapping[b'changegroup']

    def wrappedcgfn(*args, **kwargs):
        repo = args[1]
        if repo.ui.has_section(_NARROWACL_SECTION):
            kwargs = exchange.applynarrowacl(repo, kwargs)

        if kwargs.get(r'narrow', False) and repo.ui.configbool(
            b'experimental', b'narrowservebrokenellipses'
        ):
            getbundlechangegrouppart_narrow(*args, **kwargs)
        else:
            origcgfn(*args, **kwargs)

    exchange.getbundle2partsmapping[b'changegroup'] = wrappedcgfn

    # Extend changegroup receiver so client can fixup after widen requests.
    origcghandler = bundle2.parthandlermapping[b'changegroup']

    def wrappedcghandler(op, inpart):
        origcghandler(op, inpart)
        if util.safehasattr(op, '_widen_bundle'):
            handlechangegroup_widen(op, inpart)
        if util.safehasattr(op, '_bookmarksbackup'):
            localrepo.localrepository._bookmarks.set(
                op.repo, op._bookmarksbackup
            )
            del op._bookmarksbackup

    wrappedcghandler.params = origcghandler.params
    bundle2.parthandlermapping[b'changegroup'] = wrappedcghandler