# 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

_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

):

    cgversions = [

        v

        for v in cgversions

        if v in changegroup.supportedoutgoingversions(repo)

    ]

    if not cgversions:

    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:

    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.

        # 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,

            )

    visitnodes, relevant_nodes, ellipsisroots = exchange._computeellipsis(

        repo, common, heads, set(), newmatch, depth=depth

    )

    if visitnodes:

        packer = changegroup.getbundler(

            version,

            repo,

            matcher=newmatch,

            ellipses=True,

            shallow=depth is not None,

            ellipsisroots=ellipsisroots,

            fullnodes=relevant_nodes,

        )

@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

    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()

    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(

            )

        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(

        )

        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).

        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

    try:

        gen = exchange.readbundle(ui, f, chgrpfile, vfs)

        # silence internal shuffling chatter

        if ui.verbose:

            override = {}

        with ui.configoverride(override):

            if isinstance(gen, bundle2.unbundle20):

                    bundle2.processbundle(repo, gen, lambda: tr)

            else:

    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(

                    % (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

    # Extend changegroup serving to handle requests from narrow clients.

    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(

        ):

            getbundlechangegrouppart_narrow(*args, **kwargs)

        else:

            origcgfn(*args, **kwargs)

    # Extend changegroup receiver so client can fixup after widen requests.

    def wrappedcghandler(op, inpart):

        origcghandler(op, inpart)

            handlechangegroup_widen(op, inpart)

            localrepo.localrepository._bookmarks.set(

                op.repo, op._bookmarksbackup

            )

            del op._bookmarksbackup

    wrappedcghandler.params = origcghandler.params