largefiles: batch statlfile requests when pushing a largefiles repo (issue3386)

This implements a part of issue 3386.  It batches the request for the status of
all largefiles in the revisions that are about to be pushed into a single
request, instead of doing N separate requests.

In a real world test case, this change was verified to save 1,116 round-trips to
the server.  It only requires a client-side change; it is backwards-compatible
with an older version of the server.

hgext/largefiles/basestore.py

@@ -1,195 +1,195 b''
 # 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.
 
 '''base class for store implementations and store-related utility code'''
 
 import binascii
 import re
 
 from mercurial import util, node, hg
 from mercurial.i18n import _
 
 import lfutil
 
 class StoreError(Exception):
     '''Raised when there is a problem getting files from or putting
     files to a central store.'''
     def __init__(self, filename, hash, url, detail):
         self.filename = filename
         self.hash = hash
         self.url = url
         self.detail = detail
 
     def longmessage(self):
         if self.url:
             return ('%s: %s\n'
                     '(failed URL: %s)\n'
                     % (self.filename, self.detail, self.url))
         else:
             return ('%s: %s\n'
                     '(no default or default-push path set in hgrc)\n'
                     % (self.filename, self.detail))
 
     def __str__(self):
         return "%s: %s" % (self.url, self.detail)
 
 class basestore(object):
     def __init__(self, ui, repo, url):
         self.ui = ui
         self.repo = repo
         self.url = url
 
     def put(self, source, hash):
         '''Put source file into the store under <filename>/<hash>.'''
         raise NotImplementedError('abstract method')
 
-    def exists(self, hash):
-        '''Check to see if the store contains the given hash.'''
+    def exists(self, hashes):
+        '''Check to see if the store contains the given hashes.'''
         raise NotImplementedError('abstract method')
 
     def get(self, files):
         '''Get the specified largefiles from the store and write to local
         files under repo.root.  files is a list of (filename, hash)
         tuples.  Return (success, missing), lists of files successfuly
         downloaded and those not found in the store.  success is a list
         of (filename, hash) tuples; missing is a list of filenames that
         we could not get.  (The detailed error message will already have
         been presented to the user, so missing is just supplied as a
         summary.)'''
         success = []
         missing = []
         ui = self.ui
 
         at = 0
         for filename, hash in files:
             ui.progress(_('getting largefiles'), at, unit='lfile',
                 total=len(files))
             at += 1
             ui.note(_('getting %s:%s\n') % (filename, hash))
 
             storefilename = lfutil.storepath(self.repo, hash)
             tmpfile = util.atomictempfile(storefilename,
                                           createmode=self.repo.store.createmode)
 
             try:
                 hhash = binascii.hexlify(self._getfile(tmpfile, filename, hash))
             except StoreError, err:
                 ui.warn(err.longmessage())
                 hhash = ""
 
             if hhash != hash:
                 if hhash != "":
                     ui.warn(_('%s: data corruption (expected %s, got %s)\n')
                             % (filename, hash, hhash))
                 tmpfile.discard() # no-op if it's already closed
                 missing.append(filename)
                 continue
 
             tmpfile.close()
             lfutil.linktousercache(self.repo, hash)
             success.append((filename, hhash))
 
         ui.progress(_('getting largefiles'), None)
         return (success, missing)
 
     def verify(self, revs, contents=False):
         '''Verify the existence (and, optionally, contents) of every big
         file revision referenced by every changeset in revs.
         Return 0 if all is well, non-zero on any errors.'''
         write = self.ui.write
         failed = False
 
         write(_('searching %d changesets for largefiles\n') % len(revs))
         verified = set()                # set of (filename, filenode) tuples
 
         for rev in revs:
             cctx = self.repo[rev]
             cset = "%d:%s" % (cctx.rev(), node.short(cctx.node()))
 
             failed = util.any(self._verifyfile(
                 cctx, cset, contents, standin, verified) for standin in cctx)
 
         numrevs = len(verified)
         numlfiles = len(set([fname for (fname, fnode) in verified]))
         if contents:
             write(_('verified contents of %d revisions of %d largefiles\n')
                   % (numrevs, numlfiles))
         else:
             write(_('verified existence of %d revisions of %d largefiles\n')
                   % (numrevs, numlfiles))
 
         return int(failed)
 
     def _getfile(self, tmpfile, filename, hash):
         '''Fetch one revision of one file from the store and write it
         to tmpfile.  Compute the hash of the file on-the-fly as it
         downloads and return the binary hash.  Close tmpfile.  Raise
         StoreError if unable to download the file (e.g. it does not
         exist in the store).'''
         raise NotImplementedError('abstract method')
 
     def _verifyfile(self, cctx, cset, contents, standin, verified):
         '''Perform the actual verification of a file in the store.
         '''
         raise NotImplementedError('abstract method')
 
 import localstore, wirestore
 
 _storeprovider = {
     'file':  [localstore.localstore],
     'http':  [wirestore.wirestore],
     'https': [wirestore.wirestore],
     'ssh': [wirestore.wirestore],
     }
 
 _scheme_re = re.compile(r'^([a-zA-Z0-9+-.]+)://')
 
 # During clone this function is passed the src's ui object
 # but it needs the dest's ui object so it can read out of
 # the config file. Use repo.ui instead.
 def _openstore(repo, remote=None, put=False):
     ui = repo.ui
 
     if not remote:
         lfpullsource = getattr(repo, 'lfpullsource', None)
         if lfpullsource:
             path = ui.expandpath(lfpullsource)
         else:
             path = ui.expandpath('default-push', 'default')
 
         # ui.expandpath() leaves 'default-push' and 'default' alone if
         # they cannot be expanded: fallback to the empty string,
         # meaning the current directory.
         if path == 'default-push' or path == 'default':
             path = ''
             remote = repo
         else:
             remote = hg.peer(repo, {}, path)
 
     # The path could be a scheme so use Mercurial's normal functionality
     # to resolve the scheme to a repository and use its path
     path = util.safehasattr(remote, 'url') and remote.url() or remote.path
 
     match = _scheme_re.match(path)
     if not match:                       # regular filesystem path
         scheme = 'file'
     else:
         scheme = match.group(1)
 
     try:
         storeproviders = _storeprovider[scheme]
     except KeyError:
         raise util.Abort(_('unsupported URL scheme %r') % scheme)
 
     for classobj in storeproviders:
         try:
             return classobj(ui, repo, remote)
         except lfutil.storeprotonotcapable:
             pass
 
     raise util.Abort(_('%s does not appear to be a largefile store') % path)

hgext/largefiles/lfcommands.py

@@ -1,541 +1,545 b''
 # 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.
 
 '''High-level command function for lfconvert, plus the cmdtable.'''
 
 import os
 import shutil
 
 from mercurial import util, match as match_, hg, node, context, error, \
     cmdutil, scmutil
 from mercurial.i18n import _
 from mercurial.lock import release
 
 import lfutil
 import basestore
 
 # -- Commands ----------------------------------------------------------
 
 def lfconvert(ui, src, dest, *pats, **opts):
     '''convert a normal repository to a largefiles repository
 
     Convert repository SOURCE to a new repository DEST, identical to
     SOURCE except that certain files will be converted as largefiles:
     specifically, any file that matches any PATTERN *or* whose size is
     above the minimum size threshold is converted as a largefile. The
     size used to determine whether or not to track a file as a
     largefile is the size of the first version of the file. The
     minimum size can be specified either with --size or in
     configuration as ``largefiles.size``.
 
     After running this command you will need to make sure that
     largefiles is enabled anywhere you intend to push the new
     repository.
 
     Use --to-normal to convert largefiles back to normal files; after
     this, the DEST repository can be used without largefiles at all.'''
 
     if opts['to_normal']:
         tolfile = False
     else:
         tolfile = True
         size = lfutil.getminsize(ui, True, opts.get('size'), default=None)
 
     if not hg.islocal(src):
         raise util.Abort(_('%s is not a local Mercurial repo') % src)
     if not hg.islocal(dest):
         raise util.Abort(_('%s is not a local Mercurial repo') % dest)
 
     rsrc = hg.repository(ui, src)
     ui.status(_('initializing destination %s\n') % dest)
     rdst = hg.repository(ui, dest, create=True)
 
     success = False
     dstwlock = dstlock = None
     try:
         # Lock destination to prevent modification while it is converted to.
         # Don't need to lock src because we are just reading from its history
         # which can't change.
         dstwlock = rdst.wlock()
         dstlock = rdst.lock()
 
         # Get a list of all changesets in the source.  The easy way to do this
         # is to simply walk the changelog, using changelog.nodesbewteen().
         # Take a look at mercurial/revlog.py:639 for more details.
         # Use a generator instead of a list to decrease memory usage
         ctxs = (rsrc[ctx] for ctx in rsrc.changelog.nodesbetween(None,
             rsrc.heads())[0])
         revmap = {node.nullid: node.nullid}
         if tolfile:
             lfiles = set()
             normalfiles = set()
             if not pats:
                 pats = ui.configlist(lfutil.longname, 'patterns', default=[])
             if pats:
                 matcher = match_.match(rsrc.root, '', list(pats))
             else:
                 matcher = None
 
             lfiletohash = {}
             for ctx in ctxs:
                 ui.progress(_('converting revisions'), ctx.rev(),
                     unit=_('revision'), total=rsrc['tip'].rev())
                 _lfconvert_addchangeset(rsrc, rdst, ctx, revmap,
                     lfiles, normalfiles, matcher, size, lfiletohash)
             ui.progress(_('converting revisions'), None)
 
             if os.path.exists(rdst.wjoin(lfutil.shortname)):
                 shutil.rmtree(rdst.wjoin(lfutil.shortname))
 
             for f in lfiletohash.keys():
                 if os.path.isfile(rdst.wjoin(f)):
                     os.unlink(rdst.wjoin(f))
                 try:
                     os.removedirs(os.path.dirname(rdst.wjoin(f)))
                 except OSError:
                     pass
 
             # If there were any files converted to largefiles, add largefiles
             # to the destination repository's requirements.
             if lfiles:
                 rdst.requirements.add('largefiles')
                 rdst._writerequirements()
         else:
             for ctx in ctxs:
                 ui.progress(_('converting revisions'), ctx.rev(),
                     unit=_('revision'), total=rsrc['tip'].rev())
                 _addchangeset(ui, rsrc, rdst, ctx, revmap)
 
             ui.progress(_('converting revisions'), None)
         success = True
     finally:
         rdst.dirstate.clear()
         release(dstlock, dstwlock)
         if not success:
             # we failed, remove the new directory
             shutil.rmtree(rdst.root)
 
 def _addchangeset(ui, rsrc, rdst, ctx, revmap):
  # Convert src parents to dst parents
     parents = _convertparents(ctx, revmap)
 
     # Generate list of changed files
     files = _getchangedfiles(ctx, parents)
 
     def getfilectx(repo, memctx, f):
         if lfutil.standin(f) in files:
             # if the file isn't in the manifest then it was removed
             # or renamed, raise IOError to indicate this
             try:
                 fctx = ctx.filectx(lfutil.standin(f))
             except error.LookupError:
                 raise IOError
             renamed = fctx.renamed()
             if renamed:
                 renamed = lfutil.splitstandin(renamed[0])
 
             hash = fctx.data().strip()
             path = lfutil.findfile(rsrc, hash)
             ### TODO: What if the file is not cached?
             data = ''
             fd = None
             try:
                 fd = open(path, 'rb')
                 data = fd.read()
             finally:
                 if fd:
                     fd.close()
             return context.memfilectx(f, data, 'l' in fctx.flags(),
                                       'x' in fctx.flags(), renamed)
         else:
             return _getnormalcontext(repo.ui, ctx, f, revmap)
 
     dstfiles = []
     for file in files:
         if lfutil.isstandin(file):
             dstfiles.append(lfutil.splitstandin(file))
         else:
             dstfiles.append(file)
     # Commit
     _commitcontext(rdst, parents, ctx, dstfiles, getfilectx, revmap)
 
 def _lfconvert_addchangeset(rsrc, rdst, ctx, revmap, lfiles, normalfiles,
         matcher, size, lfiletohash):
     # Convert src parents to dst parents
     parents = _convertparents(ctx, revmap)
 
     # Generate list of changed files
     files = _getchangedfiles(ctx, parents)
 
     dstfiles = []
     for f in files:
         if f not in lfiles and f not in normalfiles:
             islfile = _islfile(f, ctx, matcher, size)
             # If this file was renamed or copied then copy
             # the lfileness of its predecessor
             if f in ctx.manifest():
                 fctx = ctx.filectx(f)
                 renamed = fctx.renamed()
                 renamedlfile = renamed and renamed[0] in lfiles
                 islfile |= renamedlfile
                 if 'l' in fctx.flags():
                     if renamedlfile:
                         raise util.Abort(
                             _('renamed/copied largefile %s becomes symlink')
                             % f)
                     islfile = False
             if islfile:
                 lfiles.add(f)
             else:
                 normalfiles.add(f)
 
         if f in lfiles:
             dstfiles.append(lfutil.standin(f))
             # largefile in manifest if it has not been removed/renamed
             if f in ctx.manifest():
                 fctx = ctx.filectx(f)
                 if 'l' in fctx.flags():
                     renamed = fctx.renamed()
                     if renamed and renamed[0] in lfiles:
                         raise util.Abort(_('largefile %s becomes symlink') % f)
 
                 # largefile was modified, update standins
                 fullpath = rdst.wjoin(f)
                 util.makedirs(os.path.dirname(fullpath))
                 m = util.sha1('')
                 m.update(ctx[f].data())
                 hash = m.hexdigest()
                 if f not in lfiletohash or lfiletohash[f] != hash:
                     try:
                         fd = open(fullpath, 'wb')
                         fd.write(ctx[f].data())
                     finally:
                         if fd:
                             fd.close()
                     executable = 'x' in ctx[f].flags()
                     os.chmod(fullpath, lfutil.getmode(executable))
                     lfutil.writestandin(rdst, lfutil.standin(f), hash,
                         executable)
                     lfiletohash[f] = hash
         else:
             # normal file
             dstfiles.append(f)
 
     def getfilectx(repo, memctx, f):
         if lfutil.isstandin(f):
             # if the file isn't in the manifest then it was removed
             # or renamed, raise IOError to indicate this
             srcfname = lfutil.splitstandin(f)
             try:
                 fctx = ctx.filectx(srcfname)
             except error.LookupError:
                 raise IOError
             renamed = fctx.renamed()
             if renamed:
                 # standin is always a largefile because largefile-ness
                 # doesn't change after rename or copy
                 renamed = lfutil.standin(renamed[0])
 
             return context.memfilectx(f, lfiletohash[srcfname] + '\n', 'l' in
                 fctx.flags(), 'x' in fctx.flags(), renamed)
         else:
             return _getnormalcontext(repo.ui, ctx, f, revmap)
 
     # Commit
     _commitcontext(rdst, parents, ctx, dstfiles, getfilectx, revmap)
 
 def _commitcontext(rdst, parents, ctx, dstfiles, getfilectx, revmap):
     mctx = context.memctx(rdst, parents, ctx.description(), dstfiles,
                           getfilectx, ctx.user(), ctx.date(), ctx.extra())
     ret = rdst.commitctx(mctx)
     rdst.setparents(ret)
     revmap[ctx.node()] = rdst.changelog.tip()
 
 # Generate list of changed files
 def _getchangedfiles(ctx, parents):
     files = set(ctx.files())
     if node.nullid not in parents:
         mc = ctx.manifest()
         mp1 = ctx.parents()[0].manifest()
         mp2 = ctx.parents()[1].manifest()
         files |= (set(mp1) | set(mp2)) - set(mc)
         for f in mc:
             if mc[f] != mp1.get(f, None) or mc[f] != mp2.get(f, None):
                 files.add(f)
     return files
 
 # Convert src parents to dst parents
 def _convertparents(ctx, revmap):
     parents = []
     for p in ctx.parents():
         parents.append(revmap[p.node()])
     while len(parents) < 2:
         parents.append(node.nullid)
     return parents
 
 # Get memfilectx for a normal file
 def _getnormalcontext(ui, ctx, f, revmap):
     try:
         fctx = ctx.filectx(f)
     except error.LookupError:
         raise IOError
     renamed = fctx.renamed()
     if renamed:
         renamed = renamed[0]
 
     data = fctx.data()
     if f == '.hgtags':
         data = _converttags (ui, revmap, data)
     return context.memfilectx(f, data, 'l' in fctx.flags(),
                               'x' in fctx.flags(), renamed)
 
 # Remap tag data using a revision map
 def _converttags(ui, revmap, data):
     newdata = []
     for line in data.splitlines():
         try:
             id, name = line.split(' ', 1)
         except ValueError:
             ui.warn(_('skipping incorrectly formatted tag %s\n'
                 % line))
             continue
         try:
             newid = node.bin(id)
         except TypeError:
             ui.warn(_('skipping incorrectly formatted id %s\n'
                 % id))
             continue
         try:
             newdata.append('%s %s\n' % (node.hex(revmap[newid]),
                 name))
         except KeyError:
             ui.warn(_('no mapping for id %s\n') % id)
             continue
     return ''.join(newdata)
 
 def _islfile(file, ctx, matcher, size):
     '''Return true if file should be considered a largefile, i.e.
     matcher matches it or it is larger than size.'''
     # never store special .hg* files as largefiles
     if file == '.hgtags' or file == '.hgignore' or file == '.hgsigs':
         return False
     if matcher and matcher(file):
         return True
     try:
         return ctx.filectx(file).size() >= size * 1024 * 1024
     except error.LookupError:
         return False
 
 def uploadlfiles(ui, rsrc, rdst, files):
     '''upload largefiles to the central store'''
 
     if not files:
         return
 
     store = basestore._openstore(rsrc, rdst, put=True)
 
     at = 0
-    files = filter(lambda h: not store.exists(h), files)
+    ui.debug("sending statlfile command for %d largefiles\n" % len(files))
+    retval = store.exists(files)
+    files = filter(lambda h: not retval[h], files)
+    ui.debug("%d largefiles need to be uploaded\n" % len(files))
+
     for hash in files:
         ui.progress(_('uploading largefiles'), at, unit='largefile',
                     total=len(files))
         source = lfutil.findfile(rsrc, hash)
         if not source:
             raise util.Abort(_('largefile %s missing from store'
                                ' (needs to be uploaded)') % hash)
         # XXX check for errors here
         store.put(source, hash)
         at += 1
     ui.progress(_('uploading largefiles'), None)
 
 def verifylfiles(ui, repo, all=False, contents=False):
     '''Verify that every big file revision in the current changeset
     exists in the central store.  With --contents, also verify that
     the contents of each big file revision are correct (SHA-1 hash
     matches the revision ID).  With --all, check every changeset in
     this repository.'''
     if all:
         # Pass a list to the function rather than an iterator because we know a
         # list will work.
         revs = range(len(repo))
     else:
         revs = ['.']
 
     store = basestore._openstore(repo)
     return store.verify(revs, contents=contents)
 
 def cachelfiles(ui, repo, node, filelist=None):
     '''cachelfiles ensures that all largefiles needed by the specified revision
     are present in the repository's largefile cache.
 
     returns a tuple (cached, missing).  cached is the list of files downloaded
     by this operation; missing is the list of files that were needed but could
     not be found.'''
     lfiles = lfutil.listlfiles(repo, node)
     if filelist:
         lfiles = set(lfiles) & set(filelist)
     toget = []
 
     for lfile in lfiles:
         # If we are mid-merge, then we have to trust the standin that is in the
         # working copy to have the correct hashvalue.  This is because the
         # original hg.merge() already updated the standin as part of the normal
         # merge process -- we just have to udpate the largefile to match.
         if (getattr(repo, "_ismerging", False) and
              os.path.exists(repo.wjoin(lfutil.standin(lfile)))):
             expectedhash = lfutil.readstandin(repo, lfile)
         else:
             expectedhash = repo[node][lfutil.standin(lfile)].data().strip()
 
         # if it exists and its hash matches, it might have been locally
         # modified before updating and the user chose 'local'.  in this case,
         # it will not be in any store, so don't look for it.
         if ((not os.path.exists(repo.wjoin(lfile)) or
              expectedhash != lfutil.hashfile(repo.wjoin(lfile))) and
             not lfutil.findfile(repo, expectedhash)):
             toget.append((lfile, expectedhash))
 
     if toget:
         store = basestore._openstore(repo)
         ret = store.get(toget)
         return ret
 
     return ([], [])
 
 def downloadlfiles(ui, repo, rev=None):
     matchfn = scmutil.match(repo[None],
                             [repo.wjoin(lfutil.shortname)], {})
     def prepare(ctx, fns):
         pass
     totalsuccess = 0
     totalmissing = 0
     for ctx in cmdutil.walkchangerevs(repo, matchfn, {'rev' : rev},
                                       prepare):
         success, missing = cachelfiles(ui, repo, ctx.node())
         totalsuccess += len(success)
         totalmissing += len(missing)
     ui.status(_("%d additional largefiles cached\n") % totalsuccess)
     if totalmissing > 0:
         ui.status(_("%d largefiles failed to download\n") % totalmissing)
     return totalsuccess, totalmissing
 
 def updatelfiles(ui, repo, filelist=None, printmessage=True):
     wlock = repo.wlock()
     try:
         lfdirstate = lfutil.openlfdirstate(ui, repo)
         lfiles = set(lfutil.listlfiles(repo)) | set(lfdirstate)
 
         if filelist is not None:
             lfiles = [f for f in lfiles if f in filelist]
 
         printed = False
         if printmessage and lfiles:
             ui.status(_('getting changed largefiles\n'))
             printed = True
             cachelfiles(ui, repo, '.', lfiles)
 
         updated, removed = 0, 0
         for i in map(lambda f: _updatelfile(repo, lfdirstate, f), lfiles):
             # increment the appropriate counter according to _updatelfile's
             # return value
             updated += i > 0 and i or 0
             removed -= i < 0 and i or 0
             if printmessage and (removed or updated) and not printed:
                 ui.status(_('getting changed largefiles\n'))
                 printed = True
 
         lfdirstate.write()
         if printed and printmessage:
             ui.status(_('%d largefiles updated, %d removed\n') % (updated,
                 removed))
     finally:
         wlock.release()
 
 def _updatelfile(repo, lfdirstate, lfile):
     '''updates a single largefile and copies the state of its standin from
     the repository's dirstate to its state in the lfdirstate.
 
     returns 1 if the file was modified, -1 if the file was removed, 0 if the
     file was unchanged, and None if the needed largefile was missing from the
     cache.'''
     ret = 0
     abslfile = repo.wjoin(lfile)
     absstandin = repo.wjoin(lfutil.standin(lfile))
     if os.path.exists(absstandin):
         if os.path.exists(absstandin+'.orig'):
             shutil.copyfile(abslfile, abslfile+'.orig')
         expecthash = lfutil.readstandin(repo, lfile)
         if (expecthash != '' and
             (not os.path.exists(abslfile) or
              expecthash != lfutil.hashfile(abslfile))):
             if not lfutil.copyfromcache(repo, expecthash, lfile):
                 # use normallookup() to allocate entry in largefiles dirstate,
                 # because lack of it misleads lfilesrepo.status() into
                 # recognition that such cache missing files are REMOVED.
                 lfdirstate.normallookup(lfile)
                 return None # don't try to set the mode
             ret = 1
         mode = os.stat(absstandin).st_mode
         if mode != os.stat(abslfile).st_mode:
             os.chmod(abslfile, mode)
             ret = 1
     else:
         # Remove lfiles for which the standin is deleted, unless the
         # lfile is added to the repository again. This happens when a
         # largefile is converted back to a normal file: the standin
         # disappears, but a new (normal) file appears as the lfile.
         if os.path.exists(abslfile) and lfile not in repo[None]:
             util.unlinkpath(abslfile)
             ret = -1
     state = repo.dirstate[lfutil.standin(lfile)]
     if state == 'n':
         # When rebasing, we need to synchronize the standin and the largefile,
         # because otherwise the largefile will get reverted.  But for commit's
         # sake, we have to mark the file as unclean.
         if getattr(repo, "_isrebasing", False):
            lfdirstate.normallookup(lfile)
         else:
             lfdirstate.normal(lfile)
     elif state == 'r':
         lfdirstate.remove(lfile)
     elif state == 'a':
         lfdirstate.add(lfile)
     elif state == '?':
         lfdirstate.drop(lfile)
     return ret
 
 def catlfile(repo, lfile, rev, filename):
     hash = lfutil.readstandin(repo, lfile, rev)
     if not lfutil.inusercache(repo.ui, hash):
         store = basestore._openstore(repo)
         success, missing = store.get([(lfile, hash)])
         if len(success) != 1:
             raise util.Abort(
                 _('largefile %s is not in cache and could not be downloaded')
                     % lfile)
     path = lfutil.usercachepath(repo.ui, hash)
     fpout = cmdutil.makefileobj(repo, filename)
     fpin = open(path, "rb")
     fpout.write(fpin.read())
     fpout.close()
     fpin.close()
     return 0
 
 # -- hg commands declarations ------------------------------------------------
 
 cmdtable = {
     'lfconvert': (lfconvert,
                   [('s', 'size', '',
                     _('minimum size (MB) for files to be converted '
                       'as largefiles'),
                     'SIZE'),
                   ('', 'to-normal', False,
                    _('convert from a largefiles repo to a normal repo')),
                   ],
                   _('hg lfconvert SOURCE DEST [FILE ...]')),
     }

hgext/largefiles/proto.py

@@ -1,168 +1,173 b''
 # 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.
 
 import os
 import urllib2
 
 from mercurial import error, httprepo, util, wireproto
+from mercurial.wireproto import batchable, future
 from mercurial.i18n import _
 
 import lfutil
 
 LARGEFILES_REQUIRED_MSG = ('\nThis repository uses the largefiles extension.'
                            '\n\nPlease enable it in your Mercurial config '
                            'file.\n')
 
 def putlfile(repo, proto, sha):
     '''Put a largefile into a repository's local store and into the
     user cache.'''
     proto.redirect()
 
     path = lfutil.storepath(repo, sha)
     util.makedirs(os.path.dirname(path))
     tmpfp = util.atomictempfile(path, createmode=repo.store.createmode)
 
     try:
         try:
             proto.getfile(tmpfp)
             tmpfp._fp.seek(0)
             if sha != lfutil.hexsha1(tmpfp._fp):
                 raise IOError(0, _('largefile contents do not match hash'))
             tmpfp.close()
             lfutil.linktousercache(repo, sha)
         except IOError, e:
             repo.ui.warn(_('largefiles: failed to put %s into store: %s') %
                          (sha, e.strerror))
             return wireproto.pushres(1)
     finally:
         tmpfp.discard()
 
     return wireproto.pushres(0)
 
 def getlfile(repo, proto, sha):
     '''Retrieve a largefile from the repository-local cache or system
     cache.'''
     filename = lfutil.findfile(repo, sha)
     if not filename:
         raise util.Abort(_('requested largefile %s not present in cache') % sha)
     f = open(filename, '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 '%d\n' % length
         for chunk in f:
             yield chunk
     return wireproto.streamres(generator())
 
 def statlfile(repo, proto, sha):
     '''Return '2\n' if the largefile is missing, '1\n' if it has a
     mismatched checksum, or '0\n' if it is in good condition'''
     filename = lfutil.findfile(repo, sha)
     if not filename:
         return '2\n'
     fd = None
     try:
         fd = open(filename, 'rb')
         return lfutil.hexsha1(fd) == sha and '0\n' or '1\n'
     finally:
         if fd:
             fd.close()
 
 def wirereposetup(ui, repo):
     class lfileswirerepository(repo.__class__):
         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__, httprepo.httprepository):
                 res = None
                 try:
                     res = self._call('putlfile', data=fd, sha=sha,
                         headers={'content-type':'application/mercurial-0.1'})
                     d, output = res.split('\n', 1)
                     for l in output.splitlines(True):
                         self.ui.warn(_('remote: '), l, '\n')
                     return int(d)
                 except (ValueError, urllib2.HTTPError):
                     self.ui.warn(_('unexpected putlfile response: %s') % 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("putlfile", fd, sha=sha)
                     if ret == "":
                         raise error.ResponseError(_('putlfile failed:'),
                                 output)
                     return int(ret)
                 except IOError:
                     return 1
                 except ValueError:
                     raise error.ResponseError(
                         _('putlfile failed (unexpected response):'), ret)
 
         def getlfile(self, sha):
             stream = self._callstream("getlfile", sha=sha)
             length = stream.readline()
             try:
                 length = int(length)
             except ValueError:
                 self._abort(error.ResponseError(_("unexpected response:"),
                                                 length))
             return (length, stream)
 
+        @batchable
         def statlfile(self, sha):
+            f = future()
+            result = {'sha': sha}
+            yield result, f
             try:
-                return int(self._call("statlfile", sha=sha))
+                yield int(f.value)
             except (ValueError, urllib2.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.
-                return 2
+                yield 2
 
     repo.__class__ = lfileswirerepository
 
 # advertise the largefiles=serve capability
 def capabilities(repo, proto):
     return capabilitiesorig(repo, proto) + ' largefiles=serve'
 
 # duplicate what Mercurial's new out-of-band errors mechanism does, because
 # clients old and new alike both handle it well
 def webprotorefuseclient(self, message):
     self.req.header([('Content-Type', 'application/hg-error')])
     return message
 
 def sshprotorefuseclient(self, message):
     self.ui.write_err('%s\n-\n' % message)
     self.fout.write('\n')
     self.fout.flush()
 
     return ''
 
 def heads(repo, proto):
     if lfutil.islfilesrepo(repo):
         return wireproto.ooberror(LARGEFILES_REQUIRED_MSG)
     return wireproto.heads(repo, proto)
 
 def sshrepocallstream(self, cmd, **args):
     if cmd == 'heads' and self.capable('largefiles'):
         cmd = 'lheads'
     if cmd == 'batch' and self.capable('largefiles'):
         args['cmds'] = args['cmds'].replace('heads ', 'lheads ')
     return ssholdcallstream(self, cmd, **args)
 
 def httprepocallstream(self, cmd, **args):
     if cmd == 'heads' and self.capable('largefiles'):
         cmd = 'lheads'
     if cmd == 'batch' and self.capable('largefiles'):
         args['cmds'] = args['cmds'].replace('heads ', 'lheads ')
     return httpoldcallstream(self, cmd, **args)

hgext/largefiles/remotestore.py

@@ -1,106 +1,110 b''
 # 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.
 
 '''remote largefile store; the base class for servestore'''
 
 import urllib2
 
 from mercurial import util
 from mercurial.i18n import _
+from mercurial.wireproto import remotebatch
 
 import lfutil
 import basestore
 
 class remotestore(basestore.basestore):
     '''a largefile store accessed over a network'''
     def __init__(self, ui, repo, url):
         super(remotestore, self).__init__(ui, repo, url)
 
     def put(self, source, hash):
-        if self._verify(hash):
-            return
         if self.sendfile(source, hash):
             raise util.Abort(
                 _('remotestore: could not put %s to remote store %s')
                 % (source, self.url))
         self.ui.debug(
             _('remotestore: put %s to remote store %s') % (source, self.url))
 
-    def exists(self, hash):
-        return self._verify(hash)
+    def exists(self, hashes):
+        return self._verify(hashes)
 
     def sendfile(self, filename, hash):
         self.ui.debug('remotestore: sendfile(%s, %s)\n' % (filename, hash))
         fd = None
         try:
             try:
                 fd = lfutil.httpsendfile(self.ui, filename)
             except IOError, e:
                 raise util.Abort(
                     _('remotestore: could not open file %s: %s')
                     % (filename, str(e)))
             return self._put(hash, fd)
         finally:
             if fd:
                 fd.close()
 
     def _getfile(self, tmpfile, filename, hash):
         # quit if the largefile isn't there
         stat = self._stat(hash)
         if stat == 1:
             raise util.Abort(_('remotestore: largefile %s is invalid') % hash)
         elif stat == 2:
             raise util.Abort(_('remotestore: largefile %s is missing') % hash)
 
         try:
             length, infile = self._get(hash)
         except urllib2.HTTPError, e:
             # 401s get converted to util.Aborts; everything else is fine being
             # turned into a StoreError
             raise basestore.StoreError(filename, hash, self.url, str(e))
         except urllib2.URLError, e:
             # This usually indicates a connection problem, so don't
             # keep trying with the other files... they will probably
             # all fail too.
             raise util.Abort('%s: %s' % (self.url, e.reason))
         except IOError, e:
             raise basestore.StoreError(filename, hash, self.url, str(e))
 
         # Mercurial does not close its SSH connections after writing a stream
         if length is not None:
             infile = lfutil.limitreader(infile, length)
         return lfutil.copyandhash(lfutil.blockstream(infile), tmpfile)
 
-    def _verify(self, hash):
-        return not self._stat(hash)
+    def _verify(self, hashes):
+        return self._stat(hashes)
 
     def _verifyfile(self, cctx, cset, contents, standin, verified):
         filename = lfutil.splitstandin(standin)
         if not filename:
             return False
         fctx = cctx[standin]
         key = (filename, fctx.filenode())
         if key in verified:
             return False
 
         verified.add(key)
 
         stat = self._stat(hash)
         if not stat:
             return False
         elif stat == 1:
             self.ui.warn(
                 _('changeset %s: %s: contents differ\n')
                 % (cset, filename))
             return True # failed
         elif stat == 2:
             self.ui.warn(
                 _('changeset %s: %s missing\n')
                 % (cset, filename))
             return True # failed
         else:
             raise RuntimeError('verify failed: unexpected response from '
                                'statlfile (%r)' % stat)
+
+    def batch(self):
+        '''Support for remote batching.'''
+        return remotebatch(self)
+

hgext/largefiles/wirestore.py

@@ -1,29 +1,37 b''
 # Copyright 2010-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.
 
 '''largefile store working over Mercurial's wire protocol'''
 
 import lfutil
 import remotestore
 
 class wirestore(remotestore.remotestore):
     def __init__(self, ui, repo, remote):
         cap = remote.capable('largefiles')
         if not cap:
             raise lfutil.storeprotonotcapable([])
         storetypes = cap.split(',')
         if 'serve' not in storetypes:
             raise lfutil.storeprotonotcapable(storetypes)
         self.remote = remote
         super(wirestore, self).__init__(ui, repo, remote.url())
 
     def _put(self, hash, fd):
         return self.remote.putlfile(hash, fd)
 
     def _get(self, hash):
         return self.remote.getlfile(hash)
 
-    def _stat(self, hash):
-        return self.remote.statlfile(hash)
+    def _stat(self, hashes):
+        batch = self.remote.batch()
+        futures = {}
+        for hash in hashes:
+            futures[hash] = batch.statlfile(hash)
+        batch.submit()
+        retval = {}
+        for hash in hashes:
+            retval[hash] = not futures[hash].value
+        return retval