1# narrowbundle2.py - bundle2 extensions for narrow repository support
2#
3# Copyright 2017 Google, Inc.
4#
5# This software may be used and distributed according to the terms of the
6# GNU General Public License version 2 or any later version.
7
8from __future__ import absolute_import
9
10import errno
11import struct
12
13from mercurial.i18n import _
14from mercurial import (
15    bundle2,
16    changegroup,
17    error,
18    exchange,
19    localrepo,
20    narrowspec,
21    repair,
22    requirements,
23    scmutil,
24    util,
25    wireprototypes,
26)
27from mercurial.utils import stringutil
28
29_NARROWACL_SECTION = b'narrowacl'
30_CHANGESPECPART = b'narrow:changespec'
31_RESSPECS = b'narrow:responsespec'
32_SPECPART = b'narrow:spec'
33_SPECPART_INCLUDE = b'include'
34_SPECPART_EXCLUDE = b'exclude'
35_KILLNODESIGNAL = b'KILL'
36_DONESIGNAL = b'DONE'
37_ELIDEDCSHEADER = b'>20s20s20sl'  # cset id, p1, p2, len(text)
38_ELIDEDMFHEADER = b'>20s20s20s20sl'  # manifest id, p1, p2, link id, len(text)
39_CSHEADERSIZE = struct.calcsize(_ELIDEDCSHEADER)
40_MFHEADERSIZE = struct.calcsize(_ELIDEDMFHEADER)
41
42# Serve a changegroup for a client with a narrow clone.
43def getbundlechangegrouppart_narrow(
44    bundler,
45    repo,
46    source,
47    bundlecaps=None,
48    b2caps=None,
49    heads=None,
50    common=None,
51    **kwargs
52):
53    assert repo.ui.configbool(b'experimental', b'narrowservebrokenellipses')
54
55    cgversions = b2caps.get(b'changegroup')
56    cgversions = [
57        v
58        for v in cgversions
59        if v in changegroup.supportedoutgoingversions(repo)
60    ]
61    if not cgversions:
62        raise ValueError(_(b'no common changegroup version'))
63    version = max(cgversions)
64
65    include = sorted(filter(bool, kwargs.get('includepats', [])))
66    exclude = sorted(filter(bool, kwargs.get('excludepats', [])))
67    generateellipsesbundle2(
68        bundler,
69        repo,
70        include,
71        exclude,
72        version,
73        common,
74        heads,
75        kwargs.get('depth', None),
76    )
77
78
79def generateellipsesbundle2(
80    bundler,
81    repo,
82    include,
83    exclude,
84    version,
85    common,
86    heads,
87    depth,
88):
89    match = narrowspec.match(repo.root, include=include, exclude=exclude)
90    if depth is not None:
91        depth = int(depth)
92        if depth < 1:
93            raise error.Abort(_(b'depth must be positive, got %d') % depth)
94
95    heads = set(heads or repo.heads())
96    common = set(common or [repo.nullid])
97
98    visitnodes, relevant_nodes, ellipsisroots = exchange._computeellipsis(
99        repo, common, heads, set(), match, depth=depth
100    )
101
102    repo.ui.debug(b'Found %d relevant revs\n' % len(relevant_nodes))
103    if visitnodes:
104        packer = changegroup.getbundler(
105            version,
106            repo,
107            matcher=match,
108            ellipses=True,
109            shallow=depth is not None,
110            ellipsisroots=ellipsisroots,
111            fullnodes=relevant_nodes,
112        )
113        cgdata = packer.generate(common, visitnodes, False, b'narrow_widen')
114
115        part = bundler.newpart(b'changegroup', data=cgdata)
116        part.addparam(b'version', version)
117        if scmutil.istreemanifest(repo):
118            part.addparam(b'treemanifest', b'1')
119
120
121def generate_ellipses_bundle2_for_widening(
122    bundler,
123    repo,
124    oldmatch,
125    newmatch,
126    version,
127    common,
128    known,
129):
130    common = set(common or [repo.nullid])
131    # Steps:
132    # 1. Send kill for "$known & ::common"
133    #
134    # 2. Send changegroup for ::common
135    #
136    # 3. Proceed.
137    #
138    # In the future, we can send kills for only the specific
139    # nodes we know should go away or change shape, and then
140    # send a data stream that tells the client something like this:
141    #
142    # a) apply this changegroup
143    # b) apply nodes XXX, YYY, ZZZ that you already have
144    # c) goto a
145    #
146    # until they've built up the full new state.
147    knownrevs = {repo.changelog.rev(n) for n in known}
148    # TODO: we could send only roots() of this set, and the
149    # list of nodes in common, and the client could work out
150    # what to strip, instead of us explicitly sending every
151    # single node.
152    deadrevs = knownrevs
153
154    def genkills():
155        for r in deadrevs:
156            yield _KILLNODESIGNAL
157            yield repo.changelog.node(r)
158        yield _DONESIGNAL
159
160    bundler.newpart(_CHANGESPECPART, data=genkills())
161    newvisit, newfull, newellipsis = exchange._computeellipsis(
162        repo, set(), common, knownrevs, newmatch
163    )
164    if newvisit:
165        packer = changegroup.getbundler(
166            version,
167            repo,
168            matcher=newmatch,
169            ellipses=True,
170            shallow=False,
171            ellipsisroots=newellipsis,
172            fullnodes=newfull,
173        )
174        cgdata = packer.generate(common, newvisit, False, b'narrow_widen')
175
176        part = bundler.newpart(b'changegroup', data=cgdata)
177        part.addparam(b'version', version)
178        if scmutil.istreemanifest(repo):
179            part.addparam(b'treemanifest', b'1')
180
181
182@bundle2.parthandler(_SPECPART, (_SPECPART_INCLUDE, _SPECPART_EXCLUDE))
183def _handlechangespec_2(op, inpart):
184    # XXX: This bundle2 handling is buggy and should be removed after hg5.2 is
185    # released. New servers will send a mandatory bundle2 part named
186    # 'Narrowspec' and will send specs as data instead of params.
187    # Refer to issue5952 and 6019
188    includepats = set(inpart.params.get(_SPECPART_INCLUDE, b'').splitlines())
189    excludepats = set(inpart.params.get(_SPECPART_EXCLUDE, b'').splitlines())
190    narrowspec.validatepatterns(includepats)
191    narrowspec.validatepatterns(excludepats)
192
193    if not requirements.NARROW_REQUIREMENT in op.repo.requirements:
194        op.repo.requirements.add(requirements.NARROW_REQUIREMENT)
195        scmutil.writereporequirements(op.repo)
196    op.repo.setnarrowpats(includepats, excludepats)
197    narrowspec.copytoworkingcopy(op.repo)
198
199
200@bundle2.parthandler(_RESSPECS)
201def _handlenarrowspecs(op, inpart):
202    data = inpart.read()
203    inc, exc = data.split(b'\0')
204    includepats = set(inc.splitlines())
205    excludepats = set(exc.splitlines())
206    narrowspec.validatepatterns(includepats)
207    narrowspec.validatepatterns(excludepats)
208
209    if requirements.NARROW_REQUIREMENT not in op.repo.requirements:
210        op.repo.requirements.add(requirements.NARROW_REQUIREMENT)
211        scmutil.writereporequirements(op.repo)
212    op.repo.setnarrowpats(includepats, excludepats)
213    narrowspec.copytoworkingcopy(op.repo)
214
215
216@bundle2.parthandler(_CHANGESPECPART)
217def _handlechangespec(op, inpart):
218    repo = op.repo
219    cl = repo.changelog
220
221    # changesets which need to be stripped entirely. either they're no longer
222    # needed in the new narrow spec, or the server is sending a replacement
223    # in the changegroup part.
224    clkills = set()
225
226    # A changespec part contains all the updates to ellipsis nodes
227    # that will happen as a result of widening or narrowing a
228    # repo. All the changes that this block encounters are ellipsis
229    # nodes or flags to kill an existing ellipsis.
230    chunksignal = changegroup.readexactly(inpart, 4)
231    while chunksignal != _DONESIGNAL:
232        if chunksignal == _KILLNODESIGNAL:
233            # a node used to be an ellipsis but isn't anymore
234            ck = changegroup.readexactly(inpart, 20)
235            if cl.hasnode(ck):
236                clkills.add(ck)
237        else:
238            raise error.Abort(
239                _(b'unexpected changespec node chunk type: %s') % chunksignal
240            )
241        chunksignal = changegroup.readexactly(inpart, 4)
242
243    if clkills:
244        # preserve bookmarks that repair.strip() would otherwise strip
245        op._bookmarksbackup = repo._bookmarks
246
247        class dummybmstore(dict):
248            def applychanges(self, repo, tr, changes):
249                pass
250
251        localrepo.localrepository._bookmarks.set(repo, dummybmstore())
252        chgrpfile = repair.strip(
253            op.ui, repo, list(clkills), backup=True, topic=b'widen'
254        )
255        if chgrpfile:
256            op._widen_uninterr = repo.ui.uninterruptible()
257            op._widen_uninterr.__enter__()
258            # presence of _widen_bundle attribute activates widen handler later
259            op._widen_bundle = chgrpfile
260    # Set the new narrowspec if we're widening. The setnewnarrowpats() method
261    # will currently always be there when using the core+narrowhg server, but
262    # other servers may include a changespec part even when not widening (e.g.
263    # because we're deepening a shallow repo).
264    if util.safehasattr(repo, 'setnewnarrowpats'):
265        repo.setnewnarrowpats()
266
267
268def handlechangegroup_widen(op, inpart):
269    """Changegroup exchange handler which restores temporarily-stripped nodes"""
270    # We saved a bundle with stripped node data we must now restore.
271    # This approach is based on mercurial/repair.py@6ee26a53c111.
272    repo = op.repo
273    ui = op.ui
274
275    chgrpfile = op._widen_bundle
276    del op._widen_bundle
277    vfs = repo.vfs
278
279    ui.note(_(b"adding branch\n"))
280    f = vfs.open(chgrpfile, b"rb")
281    try:
282        gen = exchange.readbundle(ui, f, chgrpfile, vfs)
283        # silence internal shuffling chatter
284        maybe_silent = (
285            ui.silent() if not ui.verbose else util.nullcontextmanager()
286        )
287        with maybe_silent:
288            if isinstance(gen, bundle2.unbundle20):
289                with repo.transaction(b'strip') as tr:
290                    bundle2.processbundle(repo, gen, lambda: tr)
291            else:
292                gen.apply(
293                    repo, b'strip', b'bundle:' + vfs.join(chgrpfile), True
294                )
295    finally:
296        f.close()
297
298    # remove undo files
299    for undovfs, undofile in repo.undofiles():
300        try:
301            undovfs.unlink(undofile)
302        except OSError as e:
303            if e.errno != errno.ENOENT:
304                ui.warn(
305                    _(b'error removing %s: %s\n')
306                    % (undovfs.join(undofile), stringutil.forcebytestr(e))
307                )
308
309    # Remove partial backup only if there were no exceptions
310    op._widen_uninterr.__exit__(None, None, None)
311    vfs.unlink(chgrpfile)
312
313
314def setup():
315    """Enable narrow repo support in bundle2-related extension points."""
316    getbundleargs = wireprototypes.GETBUNDLE_ARGUMENTS
317
318    getbundleargs[b'narrow'] = b'boolean'
319    getbundleargs[b'depth'] = b'plain'
320    getbundleargs[b'oldincludepats'] = b'csv'
321    getbundleargs[b'oldexcludepats'] = b'csv'
322    getbundleargs[b'known'] = b'csv'
323
324    # Extend changegroup serving to handle requests from narrow clients.
325    origcgfn = exchange.getbundle2partsmapping[b'changegroup']
326
327    def wrappedcgfn(*args, **kwargs):
328        repo = args[1]
329        if repo.ui.has_section(_NARROWACL_SECTION):
330            kwargs = exchange.applynarrowacl(repo, kwargs)
331
332        if kwargs.get('narrow', False) and repo.ui.configbool(
333            b'experimental', b'narrowservebrokenellipses'
334        ):
335            getbundlechangegrouppart_narrow(*args, **kwargs)
336        else:
337            origcgfn(*args, **kwargs)
338
339    exchange.getbundle2partsmapping[b'changegroup'] = wrappedcgfn
340
341    # Extend changegroup receiver so client can fixup after widen requests.
342    origcghandler = bundle2.parthandlermapping[b'changegroup']
343
344    def wrappedcghandler(op, inpart):
345        origcghandler(op, inpart)
346        if util.safehasattr(op, '_widen_bundle'):
347            handlechangegroup_widen(op, inpart)
348        if util.safehasattr(op, '_bookmarksbackup'):
349            localrepo.localrepository._bookmarks.set(
350                op.repo, op._bookmarksbackup
351            )
352            del op._bookmarksbackup
353
354    wrappedcghandler.params = origcghandler.params
355    bundle2.parthandlermapping[b'changegroup'] = wrappedcghandler
356