largefiles: redo heads interception

The largefiles extension replaces the "heads" wire command and tries to
redirect all uses towards the custom "lheads" wire command. As seen in
issue6384, this doesn't currently work for ssh. Instead of hooking into
the _callstream interface, properly register the command for the peer
instance and monkeypatch the executor to do the redirection. This works
transparently for both all kinds of peers and both for the batch and
non-batch case.

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

hgext/largefiles/__init__.py

@@ -1,209 +1,202
 # Copyright 2009-2010 Gregory P. Ward
 # Copyright 2009-2010 Intelerad Medical Systems Incorporated
 # Copyright 2010-2011 Fog Creek Software
 # Copyright 2010-2011 Unity Technologies
 #
 # This software may be used and distributed according to the terms of the
 # GNU General Public License version 2 or any later version.
 
 '''track large binary files
 
 Large binary files tend to be not very compressible, not very
 diffable, and not at all mergeable. Such files are not handled
 efficiently by Mercurial's storage format (revlog), which is based on
 compressed binary deltas; storing large binary files as regular
 Mercurial files wastes bandwidth and disk space and increases
 Mercurial's memory usage. The largefiles extension addresses these
 problems by adding a centralized client-server layer on top of
 Mercurial: largefiles live in a *central store* out on the network
 somewhere, and you only fetch the revisions that you need when you
 need them.
 
 largefiles works by maintaining a "standin file" in .hglf/ for each
 largefile. The standins are small (41 bytes: an SHA-1 hash plus
 newline) and are tracked by Mercurial. Largefile revisions are
 identified by the SHA-1 hash of their contents, which is written to
 the standin. largefiles uses that revision ID to get/put largefile
 revisions from/to the central store. This saves both disk space and
 bandwidth, since you don't need to retrieve all historical revisions
 of large files when you clone or pull.
 
 To start a new repository or add new large binary files, just add
 --large to your :hg:`add` command. For example::
 
   $ dd if=/dev/urandom of=randomdata count=2000
   $ hg add --large randomdata
   $ hg commit -m "add randomdata as a largefile"
 
 When you push a changeset that adds/modifies largefiles to a remote
 repository, its largefile revisions will be uploaded along with it.
 Note that the remote Mercurial must also have the largefiles extension
 enabled for this to work.
 
 When you pull a changeset that affects largefiles from a remote
 repository, the largefiles for the changeset will by default not be
 pulled down. However, when you update to such a revision, any
 largefiles needed by that revision are downloaded and cached (if
 they have never been downloaded before). One way to pull largefiles
 when pulling is thus to use --update, which will update your working
 copy to the latest pulled revision (and thereby downloading any new
 largefiles).
 
 If you want to pull largefiles you don't need for update yet, then
 you can use pull with the `--lfrev` option or the :hg:`lfpull` command.
 
 If you know you are pulling from a non-default location and want to
 download all the largefiles that correspond to the new changesets at
 the same time, then you can pull with `--lfrev "pulled()"`.
 
 If you just want to ensure that you will have the largefiles needed to
 merge or rebase with new heads that you are pulling, then you can pull
 with `--lfrev "head(pulled())"` flag to pre-emptively download any largefiles
 that are new in the heads you are pulling.
 
 Keep in mind that network access may now be required to update to
 changesets that you have not previously updated to. The nature of the
 largefiles extension means that updating is no longer guaranteed to
 be a local-only operation.
 
 If you already have large files tracked by Mercurial without the
 largefiles extension, you will need to convert your repository in
 order to benefit from largefiles. This is done with the
 :hg:`lfconvert` command::
 
   $ hg lfconvert --size 10 oldrepo newrepo
 
 In repositories that already have largefiles in them, any new file
 over 10MB will automatically be added as a largefile. To change this
 threshold, set ``largefiles.minsize`` in your Mercurial config file
 to the minimum size in megabytes to track as a largefile, or use the
 --lfsize option to the add command (also in megabytes)::
 
   [largefiles]
   minsize = 2
 
   $ hg add --lfsize 2
 
 The ``largefiles.patterns`` config option allows you to specify a list
 of filename patterns (see :hg:`help patterns`) that should always be
 tracked as largefiles::
 
   [largefiles]
   patterns =
     *.jpg
     re:.*\\.(png|bmp)$
     library.zip
     content/audio/*
 
 Files that match one of these patterns will be added as largefiles
 regardless of their size.
 
 The ``largefiles.minsize`` and ``largefiles.patterns`` config options
 will be ignored for any repositories not already containing a
 largefile. To add the first largefile to a repository, you must
 explicitly do so with the --large flag passed to the :hg:`add`
 command.
 '''
 from __future__ import absolute_import
 
 from mercurial import (
     cmdutil,
     extensions,
     exthelper,
     hg,
     httppeer,
     localrepo,
     sshpeer,
     wireprotov1server,
 )
 
 from . import (
     lfcommands,
     overrides,
     proto,
     reposetup,
 )
 
 # Note for extension authors: ONLY specify testedwith = 'ships-with-hg-core' for
 # extensions which SHIP WITH MERCURIAL. Non-mainline extensions should
 # be specifying the version(s) of Mercurial they are tested with, or
 # leave the attribute unspecified.
 testedwith = b'ships-with-hg-core'
 
 eh = exthelper.exthelper()
 eh.merge(lfcommands.eh)
 eh.merge(overrides.eh)
 eh.merge(proto.eh)
 
 eh.configitem(
     b'largefiles',
     b'minsize',
     default=eh.configitem.dynamicdefault,
 )
 eh.configitem(
     b'largefiles',
     b'patterns',
     default=list,
 )
 eh.configitem(
     b'largefiles',
     b'usercache',
     default=None,
 )
 
 cmdtable = eh.cmdtable
 configtable = eh.configtable
 extsetup = eh.finalextsetup
 reposetup = reposetup.reposetup
 uisetup = eh.finaluisetup
 
 
 def featuresetup(ui, supported):
     # don't die on seeing a repo with the largefiles requirement
     supported |= {b'largefiles'}
 
 
 @eh.uisetup
 def _uisetup(ui):
     localrepo.featuresetupfuncs.add(featuresetup)
     hg.wirepeersetupfuncs.append(proto.wirereposetup)
 
     cmdutil.outgoinghooks.add(b'largefiles', overrides.outgoinghook)
     cmdutil.summaryremotehooks.add(b'largefiles', overrides.summaryremotehook)
 
     # create the new wireproto commands ...
     wireprotov1server.wireprotocommand(b'putlfile', b'sha', permission=b'push')(
         proto.putlfile
     )
     wireprotov1server.wireprotocommand(b'getlfile', b'sha', permission=b'pull')(
         proto.getlfile
     )
     wireprotov1server.wireprotocommand(
         b'statlfile', b'sha', permission=b'pull'
     )(proto.statlfile)
     wireprotov1server.wireprotocommand(b'lheads', b'', permission=b'pull')(
         wireprotov1server.heads
     )
 
     extensions.wrapfunction(
         wireprotov1server.commands[b'heads'], b'func', proto.heads
     )
     # TODO also wrap wireproto.commandsv2 once heads is implemented there.
 
-    # can't do this in reposetup because it needs to have happened before
-    # wirerepo.__init__ is called
-    proto.ssholdcallstream = sshpeer.sshv1peer._callstream
-    proto.httpoldcallstream = httppeer.httppeer._callstream
-    sshpeer.sshv1peer._callstream = proto.sshrepocallstream
-    httppeer.httppeer._callstream = proto.httprepocallstream
-
     # override some extensions' stuff as well
     for name, module in extensions.extensions():
         if name == b'rebase':
             # TODO: teach exthelper to handle this
             extensions.wrapfunction(
                 module, b'rebase', overrides.overriderebasecmd
             )
 
 
 revsetpredicate = eh.revsetpredicate

hgext/largefiles/proto.py

@@ -1,221 +1,218
 # Copyright 2011 Fog Creek Software
 #
 # 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 os
 import re
 
 from mercurial.i18n import _
 from mercurial.pycompat import open
 
 from mercurial import (
     error,
     exthelper,
     httppeer,
     util,
     wireprototypes,
     wireprotov1peer,
     wireprotov1server,
 )
 
 from . import lfutil
 
 urlerr = util.urlerr
 urlreq = util.urlreq
 
 LARGEFILES_REQUIRED_MSG = (
     b'\nThis repository uses the largefiles extension.'
     b'\n\nPlease enable it in your Mercurial config '
     b'file.\n'
 )
 
 eh = exthelper.exthelper()
 
-# these will all be replaced by largefiles.uisetup
-ssholdcallstream = None
-httpoldcallstream = None
-
 
 def putlfile(repo, proto, sha):
     """Server command for putting a largefile into a repository's local store
     and into the user cache."""
     with proto.mayberedirectstdio() as output:
         path = lfutil.storepath(repo, sha)
         util.makedirs(os.path.dirname(path))
         tmpfp = util.atomictempfile(path, createmode=repo.store.createmode)
 
         try:
             for p in proto.getpayload():
                 tmpfp.write(p)
             tmpfp._fp.seek(0)
             if sha != lfutil.hexsha1(tmpfp._fp):
                 raise IOError(0, _(b'largefile contents do not match hash'))
             tmpfp.close()
             lfutil.linktousercache(repo, sha)
         except IOError as e:
             repo.ui.warn(
                 _(b'largefiles: failed to put %s into store: %s\n')
                 % (sha, e.strerror)
             )
             return wireprototypes.pushres(
                 1, output.getvalue() if output else b''
             )
         finally:
             tmpfp.discard()
 
     return wireprototypes.pushres(0, output.getvalue() if output else b'')
 
 
 def getlfile(repo, proto, sha):
     """Server command for retrieving a largefile from the repository-local
     cache or user cache."""
     filename = lfutil.findfile(repo, sha)
     if not filename:
         raise error.Abort(
             _(b'requested largefile %s not present in cache') % sha
         )
     f = open(filename, b'rb')
     length = os.fstat(f.fileno())[6]
 
     # Since we can't set an HTTP content-length header here, and
     # Mercurial core provides no way to give the length of a streamres
     # (and reading the entire file into RAM would be ill-advised), we
     # just send the length on the first line of the response, like the
     # ssh proto does for string responses.
     def generator():
         yield b'%d\n' % length
         for chunk in util.filechunkiter(f):
             yield chunk
 
     return wireprototypes.streamreslegacy(gen=generator())
 
 
 def statlfile(repo, proto, sha):
     """Server command for checking if a largefile is present - returns '2\n' if
     the largefile is missing, '0\n' if it seems to be in good condition.
 
     The value 1 is reserved for mismatched checksum, but that is too expensive
     to be verified on every stat and must be caught be running 'hg verify'
     server side."""
     filename = lfutil.findfile(repo, sha)
     if not filename:
         return wireprototypes.bytesresponse(b'2\n')
     return wireprototypes.bytesresponse(b'0\n')
 
 
 def wirereposetup(ui, repo):
+    orig_commandexecutor = repo.commandexecutor
+
     class lfileswirerepository(repo.__class__):
+        def commandexecutor(self):
+            executor = orig_commandexecutor()
+            if self.capable(b'largefiles'):
+                orig_callcommand = executor.callcommand
+
+                class lfscommandexecutor(executor.__class__):
+                    def callcommand(self, command, args):
+                        if command == b'heads':
+                            command = b'lheads'
+                        return orig_callcommand(command, args)
+
+                executor.__class__ = lfscommandexecutor
+            return executor
+
+        @wireprotov1peer.batchable
+        def lheads(self):
+            return self.heads.batchable(self)
+
         def putlfile(self, sha, fd):
             # unfortunately, httprepository._callpush tries to convert its
             # input file-like into a bundle before sending it, so we can't use
             # it ...
             if issubclass(self.__class__, httppeer.httppeer):
                 res = self._call(
                     b'putlfile',
                     data=fd,
                     sha=sha,
                     headers={'content-type': 'application/mercurial-0.1'},
                 )
                 try:
                     d, output = res.split(b'\n', 1)
                     for l in output.splitlines(True):
                         self.ui.warn(_(b'remote: '), l)  # assume l ends with \n
                     return int(d)
                 except ValueError:
                     self.ui.warn(_(b'unexpected putlfile response: %r\n') % res)
                     return 1
             # ... but we can't use sshrepository._call because the data=
             # argument won't get sent, and _callpush does exactly what we want
             # in this case: send the data straight through
             else:
                 try:
                     ret, output = self._callpush(b"putlfile", fd, sha=sha)
                     if ret == b"":
                         raise error.ResponseError(
                             _(b'putlfile failed:'), output
                         )
                     return int(ret)
                 except IOError:
                     return 1
                 except ValueError:
                     raise error.ResponseError(
                         _(b'putlfile failed (unexpected response):'), ret
                     )
 
         def getlfile(self, sha):
             """returns an iterable with the chunks of the file with sha sha"""
             stream = self._callstream(b"getlfile", sha=sha)
             length = stream.readline()
             try:
                 length = int(length)
             except ValueError:
                 self._abort(
                     error.ResponseError(_(b"unexpected response:"), length)
                 )
 
             # SSH streams will block if reading more than length
             for chunk in util.filechunkiter(stream, limit=length):
                 yield chunk
             # HTTP streams must hit the end to process the last empty
             # chunk of Chunked-Encoding so the connection can be reused.
             if issubclass(self.__class__, httppeer.httppeer):
                 chunk = stream.read(1)
                 if chunk:
                     self._abort(
                         error.ResponseError(_(b"unexpected response:"), chunk)
                     )
 
         @wireprotov1peer.batchable
         def statlfile(self, sha):
             f = wireprotov1peer.future()
             result = {b'sha': sha}
             yield result, f
             try:
                 yield int(f.value)
             except (ValueError, urlerr.httperror):
                 # If the server returns anything but an integer followed by a
                 # newline, newline, it's not speaking our language; if we get
                 # an HTTP error, we can't be sure the largefile is present;
                 # either way, consider it missing.
                 yield 2
 
     repo.__class__ = lfileswirerepository
 
 
 # advertise the largefiles=serve capability
 @eh.wrapfunction(wireprotov1server, b'_capabilities')
 def _capabilities(orig, repo, proto):
     '''announce largefile server capability'''
     caps = orig(repo, proto)
     caps.append(b'largefiles=serve')
     return caps
 
 
 def heads(orig, repo, proto):
     """Wrap server command - largefile capable clients will know to call
     lheads instead"""
     if lfutil.islfilesrepo(repo):
         return wireprototypes.ooberror(LARGEFILES_REQUIRED_MSG)
 
     return orig(repo, proto)
-
-
-def sshrepocallstream(self, cmd, **args):
-    if cmd == b'heads' and self.capable(b'largefiles'):
-        cmd = b'lheads'
-    if cmd == b'batch' and self.capable(b'largefiles'):
-        args['cmds'] = args[r'cmds'].replace(b'heads ', b'lheads ')
-    return ssholdcallstream(self, cmd, **args)
-
-
-headsre = re.compile(br'(^|;)heads\b')
-
-
-def httprepocallstream(self, cmd, **args):
-    if cmd == b'heads' and self.capable(b'largefiles'):
-        cmd = b'lheads'
-    if cmd == b'batch' and self.capable(b'largefiles'):
-        args['cmds'] = headsre.sub(b'lheads', args['cmds'])
-    return httpoldcallstream(self, cmd, **args)