wireprotov1peer: update all rpcs to use the new batchable scheme

If desired, we could keep the future class and the function that
upgrades an old style rpc instead of a new style, for extensions.

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

hgext/fastannotate/protocol.py

@@ -1,263 +1,261 b''
 # Copyright 2016-present Facebook. All Rights Reserved.
 #
 # protocol: logic for a server providing fastannotate support
 #
 # 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 contextlib
 import os
 
 from mercurial.i18n import _
 from mercurial.pycompat import open
 from mercurial import (
     error,
     extensions,
     hg,
     pycompat,
     util,
     wireprotov1peer,
     wireprotov1server,
 )
 from mercurial.utils import (
     urlutil,
 )
 from . import context
 
 # common
 
 
 def _getmaster(ui):
     """get the mainbranch, and enforce it is set"""
     master = ui.config(b'fastannotate', b'mainbranch')
     if not master:
         raise error.Abort(
             _(
                 b'fastannotate.mainbranch is required '
                 b'for both the client and the server'
             )
         )
     return master
 
 
 # server-side
 
 
 def _capabilities(orig, repo, proto):
     result = orig(repo, proto)
     result.append(b'getannotate')
     return result
 
 
 def _getannotate(repo, proto, path, lastnode):
     # output:
     #   FILE := vfspath + '\0' + str(size) + '\0' + content
     #   OUTPUT := '' | FILE + OUTPUT
     result = b''
     buildondemand = repo.ui.configbool(
         b'fastannotate', b'serverbuildondemand', True
     )
     with context.annotatecontext(repo, path) as actx:
         if buildondemand:
             # update before responding to the client
             master = _getmaster(repo.ui)
             try:
                 if not actx.isuptodate(master):
                     actx.annotate(master, master)
             except Exception:
                 # non-fast-forward move or corrupted. rebuild automically.
                 actx.rebuild()
                 try:
                     actx.annotate(master, master)
                 except Exception:
                     actx.rebuild()  # delete files
             finally:
                 # although the "with" context will also do a close/flush, we
                 # need to do it early so we can send the correct respond to
                 # client.
                 actx.close()
         # send back the full content of revmap and linelog, in the future we
         # may want to do some rsync-like fancy updating.
         # the lastnode check is not necessary if the client and the server
         # agree where the main branch is.
         if actx.lastnode != lastnode:
             for p in [actx.revmappath, actx.linelogpath]:
                 if not os.path.exists(p):
                     continue
                 with open(p, b'rb') as f:
                     content = f.read()
                 vfsbaselen = len(repo.vfs.base + b'/')
                 relpath = p[vfsbaselen:]
                 result += b'%s\0%d\0%s' % (relpath, len(content), content)
     return result
 
 
 def _registerwireprotocommand():
     if b'getannotate' in wireprotov1server.commands:
         return
     wireprotov1server.wireprotocommand(b'getannotate', b'path lastnode')(
         _getannotate
     )
 
 
 def serveruisetup(ui):
     _registerwireprotocommand()
     extensions.wrapfunction(wireprotov1server, b'_capabilities', _capabilities)
 
 
 # client-side
 
 
 def _parseresponse(payload):
     result = {}
     i = 0
     l = len(payload) - 1
     state = 0  # 0: vfspath, 1: size
     vfspath = size = b''
     while i < l:
         ch = payload[i : i + 1]
         if ch == b'\0':
             if state == 1:
                 result[vfspath] = payload[i + 1 : i + 1 + int(size)]
                 i += int(size)
                 state = 0
                 vfspath = size = b''
             elif state == 0:
                 state = 1
         else:
             if state == 1:
                 size += ch
             elif state == 0:
                 vfspath += ch
         i += 1
     return result
 
 
 def peersetup(ui, peer):
     class fastannotatepeer(peer.__class__):
         @wireprotov1peer.batchable
         def getannotate(self, path, lastnode=None):
             if not self.capable(b'getannotate'):
                 ui.warn(_(b'remote peer cannot provide annotate cache\n'))
-                yield None, None
+                return None, None
             else:
                 args = {b'path': path, b'lastnode': lastnode or b''}
-                f = wireprotov1peer.future()
-                yield args, f
-                yield _parseresponse(f.value)
+                return args, _parseresponse
 
     peer.__class__ = fastannotatepeer
 
 
 @contextlib.contextmanager
 def annotatepeer(repo):
     ui = repo.ui
 
     remotedest = ui.config(b'fastannotate', b'remotepath', b'default')
     r = urlutil.get_unique_pull_path(b'fastannotate', repo, ui, remotedest)
     remotepath = r[0]
     peer = hg.peer(ui, {}, remotepath)
 
     try:
         yield peer
     finally:
         peer.close()
 
 
 def clientfetch(repo, paths, lastnodemap=None, peer=None):
     """download annotate cache from the server for paths"""
     if not paths:
         return
 
     if peer is None:
         with annotatepeer(repo) as peer:
             return clientfetch(repo, paths, lastnodemap, peer)
 
     if lastnodemap is None:
         lastnodemap = {}
 
     ui = repo.ui
     results = []
     with peer.commandexecutor() as batcher:
         ui.debug(b'fastannotate: requesting %d files\n' % len(paths))
         for p in paths:
             results.append(
                 batcher.callcommand(
                     b'getannotate',
                     {b'path': p, b'lastnode': lastnodemap.get(p)},
                 )
             )
 
         for result in results:
             r = result.result()
             # TODO: pconvert these paths on the server?
             r = {util.pconvert(p): v for p, v in pycompat.iteritems(r)}
             for path in sorted(r):
                 # ignore malicious paths
                 if not path.startswith(b'fastannotate/') or b'/../' in (
                     path + b'/'
                 ):
                     ui.debug(
                         b'fastannotate: ignored malicious path %s\n' % path
                     )
                     continue
                 content = r[path]
                 if ui.debugflag:
                     ui.debug(
                         b'fastannotate: writing %d bytes to %s\n'
                         % (len(content), path)
                     )
                 repo.vfs.makedirs(os.path.dirname(path))
                 with repo.vfs(path, b'wb') as f:
                     f.write(content)
 
 
 def _filterfetchpaths(repo, paths):
     """return a subset of paths whose history is long and need to fetch linelog
     from the server. works with remotefilelog and non-remotefilelog repos.
     """
     threshold = repo.ui.configint(b'fastannotate', b'clientfetchthreshold', 10)
     if threshold <= 0:
         return paths
 
     result = []
     for path in paths:
         try:
             if len(repo.file(path)) >= threshold:
                 result.append(path)
         except Exception:  # file not found etc.
             result.append(path)
 
     return result
 
 
 def localreposetup(ui, repo):
     class fastannotaterepo(repo.__class__):
         def prefetchfastannotate(self, paths, peer=None):
             master = _getmaster(self.ui)
             needupdatepaths = []
             lastnodemap = {}
             try:
                 for path in _filterfetchpaths(self, paths):
                     with context.annotatecontext(self, path) as actx:
                         if not actx.isuptodate(master, strict=False):
                             needupdatepaths.append(path)
                             lastnodemap[path] = actx.lastnode
                 if needupdatepaths:
                     clientfetch(self, needupdatepaths, lastnodemap, peer)
             except Exception as ex:
                 # could be directory not writable or so, not fatal
                 self.ui.debug(b'fastannotate: prefetch failed: %r\n' % ex)
 
     repo.__class__ = fastannotaterepo
 
 
 def clientreposetup(ui, repo):
     _registerwireprotocommand()
     if repo.local():
         localreposetup(ui, repo)
     # TODO: this mutates global state, but only if at least one repo
     # has the extension enabled. This is probably bad for hgweb.
     if peersetup not in hg.wirepeersetupfuncs:
         hg.wirepeersetupfuncs.append(peersetup)

hgext/infinitepush/__init__.py

@@ -1,1389 +1,1390 b''
 # Infinite push
 #
 # Copyright 2016 Facebook, Inc.
 #
 # This software may be used and distributed according to the terms of the
 # GNU General Public License version 2 or any later version.
 """ store some pushes in a remote blob store on the server (EXPERIMENTAL)
 
 IMPORTANT: if you use this extension, please contact
 [e-mail] ASAP. This extension is believed to
 be unused and barring learning of users of this functionality, we will
 delete this code at the end of 2020.
 
     [infinitepush]
     # Server-side and client-side option. Pattern of the infinitepush bookmark
     branchpattern = PATTERN
 
     # Server or client
     server = False
 
     # Server-side option. Possible values: 'disk' or 'sql'. Fails if not set
     indextype = disk
 
     # Server-side option. Used only if indextype=sql.
     # Format: 'IP:PORT:DB_NAME:USER:PASSWORD'
     sqlhost = IP:PORT:DB_NAME:USER:PASSWORD
 
     # Server-side option. Used only if indextype=disk.
     # Filesystem path to the index store
     indexpath = PATH
 
     # Server-side option. Possible values: 'disk' or 'external'
     # Fails if not set
     storetype = disk
 
     # Server-side option.
     # Path to the binary that will save bundle to the bundlestore
     # Formatted cmd line will be passed to it (see `put_args`)
     put_binary = put
 
     # Serser-side option. Used only if storetype=external.
     # Format cmd-line string for put binary. Placeholder: {filename}
     put_args = {filename}
 
     # Server-side option.
     # Path to the binary that get bundle from the bundlestore.
     # Formatted cmd line will be passed to it (see `get_args`)
     get_binary = get
 
     # Serser-side option. Used only if storetype=external.
     # Format cmd-line string for get binary. Placeholders: {filename} {handle}
     get_args = {filename} {handle}
 
     # Server-side option
     logfile = FIlE
 
     # Server-side option
     loglevel = DEBUG
 
     # Server-side option. Used only if indextype=sql.
     # Sets mysql wait_timeout option.
     waittimeout = 300
 
     # Server-side option. Used only if indextype=sql.
     # Sets mysql innodb_lock_wait_timeout option.
     locktimeout = 120
 
     # Server-side option. Used only if indextype=sql.
     # Name of the repository
     reponame = ''
 
     # Client-side option. Used by --list-remote option. List of remote scratch
     # patterns to list if no patterns are specified.
     defaultremotepatterns = ['*']
 
     # Instructs infinitepush to forward all received bundle2 parts to the
     # bundle for storage. Defaults to False.
     storeallparts = True
 
     # routes each incoming push to the bundlestore. defaults to False
     pushtobundlestore = True
 
     [remotenames]
     # Client-side option
     # This option should be set only if remotenames extension is enabled.
     # Whether remote bookmarks are tracked by remotenames extension.
     bookmarks = True
 """
 
 from __future__ import absolute_import
 
 import collections
 import contextlib
 import errno
 import functools
 import logging
 import os
 import random
 import re
 import socket
 import subprocess
 import time
 
 from mercurial.node import (
     bin,
     hex,
 )
 
 from mercurial.i18n import _
 
 from mercurial.pycompat import (
     getattr,
     open,
 )
 
 from mercurial.utils import (
     procutil,
     stringutil,
     urlutil,
 )
 
 from mercurial import (
     bundle2,
     changegroup,
     commands,
     discovery,
     encoding,
     error,
     exchange,
     extensions,
     hg,
     localrepo,
     phases,
     pushkey,
     pycompat,
     registrar,
     util,
     wireprototypes,
     wireprotov1peer,
     wireprotov1server,
 )
 
 from . import (
     bundleparts,
     common,
 )
 
 # 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'
 
 configtable = {}
 configitem = registrar.configitem(configtable)
 
 configitem(
     b'infinitepush',
     b'server',
     default=False,
 )
 configitem(
     b'infinitepush',
     b'storetype',
     default=b'',
 )
 configitem(
     b'infinitepush',
     b'indextype',
     default=b'',
 )
 configitem(
     b'infinitepush',
     b'indexpath',
     default=b'',
 )
 configitem(
     b'infinitepush',
     b'storeallparts',
     default=False,
 )
 configitem(
     b'infinitepush',
     b'reponame',
     default=b'',
 )
 configitem(
     b'scratchbranch',
     b'storepath',
     default=b'',
 )
 configitem(
     b'infinitepush',
     b'branchpattern',
     default=b'',
 )
 configitem(
     b'infinitepush',
     b'pushtobundlestore',
     default=False,
 )
 configitem(
     b'experimental',
     b'server-bundlestore-bookmark',
     default=b'',
 )
 configitem(
     b'experimental',
     b'infinitepush-scratchpush',
     default=False,
 )
 
 experimental = b'experimental'
 configbookmark = b'server-bundlestore-bookmark'
 configscratchpush = b'infinitepush-scratchpush'
 
 scratchbranchparttype = bundleparts.scratchbranchparttype
 revsetpredicate = registrar.revsetpredicate()
 templatekeyword = registrar.templatekeyword()
 _scratchbranchmatcher = lambda x: False
 _maybehash = re.compile('^[a-f0-9]+$').search
 
 
 def _buildexternalbundlestore(ui):
     put_args = ui.configlist(b'infinitepush', b'put_args', [])
     put_binary = ui.config(b'infinitepush', b'put_binary')
     if not put_binary:
         raise error.Abort(b'put binary is not specified')
     get_args = ui.configlist(b'infinitepush', b'get_args', [])
     get_binary = ui.config(b'infinitepush', b'get_binary')
     if not get_binary:
         raise error.Abort(b'get binary is not specified')
     from . import store
 
     return store.externalbundlestore(put_binary, put_args, get_binary, get_args)
 
 
 def _buildsqlindex(ui):
     sqlhost = ui.config(b'infinitepush', b'sqlhost')
     if not sqlhost:
         raise error.Abort(_(b'please set infinitepush.sqlhost'))
     host, port, db, user, password = sqlhost.split(b':')
     reponame = ui.config(b'infinitepush', b'reponame')
     if not reponame:
         raise error.Abort(_(b'please set infinitepush.reponame'))
 
     logfile = ui.config(b'infinitepush', b'logfile', b'')
     waittimeout = ui.configint(b'infinitepush', b'waittimeout', 300)
     locktimeout = ui.configint(b'infinitepush', b'locktimeout', 120)
     from . import sqlindexapi
 
     return sqlindexapi.sqlindexapi(
         reponame,
         host,
         port,
         db,
         user,
         password,
         logfile,
         _getloglevel(ui),
         waittimeout=waittimeout,
         locktimeout=locktimeout,
     )
 
 
 def _getloglevel(ui):
     loglevel = ui.config(b'infinitepush', b'loglevel', b'DEBUG')
     numeric_loglevel = getattr(logging, loglevel.upper(), None)
     if not isinstance(numeric_loglevel, int):
         raise error.Abort(_(b'invalid log level %s') % loglevel)
     return numeric_loglevel
 
 
 def _tryhoist(ui, remotebookmark):
     """returns a bookmarks with hoisted part removed
 
     Remotenames extension has a 'hoist' config that allows to use remote
     bookmarks without specifying remote path. For example, 'hg update master'
     works as well as 'hg update remote/master'. We want to allow the same in
     infinitepush.
     """
 
     if common.isremotebooksenabled(ui):
         hoist = ui.config(b'remotenames', b'hoistedpeer') + b'/'
         if remotebookmark.startswith(hoist):
             return remotebookmark[len(hoist) :]
     return remotebookmark
 
 
 class bundlestore(object):
     def __init__(self, repo):
         self._repo = repo
         storetype = self._repo.ui.config(b'infinitepush', b'storetype')
         if storetype == b'disk':
             from . import store
 
             self.store = store.filebundlestore(self._repo.ui, self._repo)
         elif storetype == b'external':
             self.store = _buildexternalbundlestore(self._repo.ui)
         else:
             raise error.Abort(
                 _(b'unknown infinitepush store type specified %s') % storetype
             )
 
         indextype = self._repo.ui.config(b'infinitepush', b'indextype')
         if indextype == b'disk':
             from . import fileindexapi
 
             self.index = fileindexapi.fileindexapi(self._repo)
         elif indextype == b'sql':
             self.index = _buildsqlindex(self._repo.ui)
         else:
             raise error.Abort(
                 _(b'unknown infinitepush index type specified %s') % indextype
             )
 
 
 def _isserver(ui):
     return ui.configbool(b'infinitepush', b'server')
 
 
 def reposetup(ui, repo):
     if _isserver(ui) and repo.local():
         repo.bundlestore = bundlestore(repo)
 
 
 def extsetup(ui):
     commonsetup(ui)
     if _isserver(ui):
         serverextsetup(ui)
     else:
         clientextsetup(ui)
 
 
 def commonsetup(ui):
     wireprotov1server.commands[b'listkeyspatterns'] = (
         wireprotolistkeyspatterns,
         b'namespace patterns',
     )
     scratchbranchpat = ui.config(b'infinitepush', b'branchpattern')
     if scratchbranchpat:
         global _scratchbranchmatcher
         kind, pat, _scratchbranchmatcher = stringutil.stringmatcher(
             scratchbranchpat
         )
 
 
 def serverextsetup(ui):
     origpushkeyhandler = bundle2.parthandlermapping[b'pushkey']
 
     def newpushkeyhandler(*args, **kwargs):
         bundle2pushkey(origpushkeyhandler, *args, **kwargs)
 
     newpushkeyhandler.params = origpushkeyhandler.params
     bundle2.parthandlermapping[b'pushkey'] = newpushkeyhandler
 
     orighandlephasehandler = bundle2.parthandlermapping[b'phase-heads']
     newphaseheadshandler = lambda *args, **kwargs: bundle2handlephases(
         orighandlephasehandler, *args, **kwargs
     )
     newphaseheadshandler.params = orighandlephasehandler.params
     bundle2.parthandlermapping[b'phase-heads'] = newphaseheadshandler
 
     extensions.wrapfunction(
         localrepo.localrepository, b'listkeys', localrepolistkeys
     )
     wireprotov1server.commands[b'lookup'] = (
         _lookupwrap(wireprotov1server.commands[b'lookup'][0]),
         b'key',
     )
     extensions.wrapfunction(exchange, b'getbundlechunks', getbundlechunks)
 
     extensions.wrapfunction(bundle2, b'processparts', processparts)
 
 
 def clientextsetup(ui):
     entry = extensions.wrapcommand(commands.table, b'push', _push)
 
     entry[1].append(
         (
             b'',
             b'bundle-store',
             None,
             _(b'force push to go to bundle store (EXPERIMENTAL)'),
         )
     )
 
     extensions.wrapcommand(commands.table, b'pull', _pull)
 
     extensions.wrapfunction(discovery, b'checkheads', _checkheads)
 
     wireprotov1peer.wirepeer.listkeyspatterns = listkeyspatterns
 
     partorder = exchange.b2partsgenorder
     index = partorder.index(b'changeset')
     partorder.insert(
         index, partorder.pop(partorder.index(scratchbranchparttype))
     )
 
 
 def _checkheads(orig, pushop):
     if pushop.ui.configbool(experimental, configscratchpush, False):
         return
     return orig(pushop)
 
 
 def wireprotolistkeyspatterns(repo, proto, namespace, patterns):
     patterns = wireprototypes.decodelist(patterns)
     d = pycompat.iteritems(repo.listkeys(encoding.tolocal(namespace), patterns))
     return pushkey.encodekeys(d)
 
 
 def localrepolistkeys(orig, self, namespace, patterns=None):
     if namespace == b'bookmarks' and patterns:
         index = self.bundlestore.index
         results = {}
         bookmarks = orig(self, namespace)
         for pattern in patterns:
             results.update(index.getbookmarks(pattern))
             if pattern.endswith(b'*'):
                 pattern = b're:^' + pattern[:-1] + b'.*'
             kind, pat, matcher = stringutil.stringmatcher(pattern)
             for bookmark, node in pycompat.iteritems(bookmarks):
                 if matcher(bookmark):
                     results[bookmark] = node
         return results
     else:
         return orig(self, namespace)
 
 
 @wireprotov1peer.batchable
 def listkeyspatterns(self, namespace, patterns):
     if not self.capable(b'pushkey'):
-        yield {}, None
-    f = wireprotov1peer.future()
+        return {}, None
     self.ui.debug(b'preparing listkeys for "%s"\n' % namespace)
-    yield {
-        b'namespace': encoding.fromlocal(namespace),
-        b'patterns': wireprototypes.encodelist(patterns),
-    }, f
-    d = f.value
+
+    def decode(d):
     self.ui.debug(
         b'received listkey for "%s": %i bytes\n' % (namespace, len(d))
     )
-    yield pushkey.decodekeys(d)
+        return pushkey.decodekeys(d)
+
+    return {
+        b'namespace': encoding.fromlocal(namespace),
+        b'patterns': wireprototypes.encodelist(patterns),
+    }, decode
 
 
 def _readbundlerevs(bundlerepo):
     return list(bundlerepo.revs(b'bundle()'))
 
 
 def _includefilelogstobundle(bundlecaps, bundlerepo, bundlerevs, ui):
     """Tells remotefilelog to include all changed files to the changegroup
 
     By default remotefilelog doesn't include file content to the changegroup.
     But we need to include it if we are fetching from bundlestore.
     """
     changedfiles = set()
     cl = bundlerepo.changelog
     for r in bundlerevs:
         # [3] means changed files
         changedfiles.update(cl.read(r)[3])
     if not changedfiles:
         return bundlecaps
 
     changedfiles = b'\0'.join(changedfiles)
     newcaps = []
     appended = False
     for cap in bundlecaps or []:
         if cap.startswith(b'excludepattern='):
             newcaps.append(b'\0'.join((cap, changedfiles)))
             appended = True
         else:
             newcaps.append(cap)
     if not appended:
         # Not found excludepattern cap. Just append it
         newcaps.append(b'excludepattern=' + changedfiles)
 
     return newcaps
 
 
 def _rebundle(bundlerepo, bundleroots, unknownhead):
     """
     Bundle may include more revision then user requested. For example,
     if user asks for revision but bundle also consists its descendants.
     This function will filter out all revision that user is not requested.
     """
     parts = []
 
     version = b'02'
     outgoing = discovery.outgoing(
         bundlerepo, commonheads=bundleroots, ancestorsof=[unknownhead]
     )
     cgstream = changegroup.makestream(bundlerepo, outgoing, version, b'pull')
     cgstream = util.chunkbuffer(cgstream).read()
     cgpart = bundle2.bundlepart(b'changegroup', data=cgstream)
     cgpart.addparam(b'version', version)
     parts.append(cgpart)
 
     return parts
 
 
 def _getbundleroots(oldrepo, bundlerepo, bundlerevs):
     cl = bundlerepo.changelog
     bundleroots = []
     for rev in bundlerevs:
         node = cl.node(rev)
         parents = cl.parents(node)
         for parent in parents:
             # include all revs that exist in the main repo
             # to make sure that bundle may apply client-side
             if parent in oldrepo:
                 bundleroots.append(parent)
     return bundleroots
 
 
 def _needsrebundling(head, bundlerepo):
     bundleheads = list(bundlerepo.revs(b'heads(bundle())'))
     return not (
         len(bundleheads) == 1 and bundlerepo[bundleheads[0]].node() == head
     )
 
 
 def _generateoutputparts(head, bundlerepo, bundleroots, bundlefile):
     """generates bundle that will be send to the user
 
     returns tuple with raw bundle string and bundle type
     """
     parts = []
     if not _needsrebundling(head, bundlerepo):
         with util.posixfile(bundlefile, b"rb") as f:
             unbundler = exchange.readbundle(bundlerepo.ui, f, bundlefile)
             if isinstance(unbundler, changegroup.cg1unpacker):
                 part = bundle2.bundlepart(
                     b'changegroup', data=unbundler._stream.read()
                 )
                 part.addparam(b'version', b'01')
                 parts.append(part)
             elif isinstance(unbundler, bundle2.unbundle20):
                 haschangegroup = False
                 for part in unbundler.iterparts():
                     if part.type == b'changegroup':
                         haschangegroup = True
                     newpart = bundle2.bundlepart(part.type, data=part.read())
                     for key, value in pycompat.iteritems(part.params):
                         newpart.addparam(key, value)
                     parts.append(newpart)
 
                 if not haschangegroup:
                     raise error.Abort(
                         b'unexpected bundle without changegroup part, '
                         + b'head: %s' % hex(head),
                         hint=b'report to administrator',
                     )
             else:
                 raise error.Abort(b'unknown bundle type')
     else:
         parts = _rebundle(bundlerepo, bundleroots, head)
 
     return parts
 
 
 def getbundlechunks(orig, repo, source, heads=None, bundlecaps=None, **kwargs):
     heads = heads or []
     # newheads are parents of roots of scratch bundles that were requested
     newphases = {}
     scratchbundles = []
     newheads = []
     scratchheads = []
     nodestobundle = {}
     allbundlestocleanup = []
     try:
         for head in heads:
             if not repo.changelog.index.has_node(head):
                 if head not in nodestobundle:
                     newbundlefile = common.downloadbundle(repo, head)
                     bundlepath = b"bundle:%s+%s" % (repo.root, newbundlefile)
                     bundlerepo = hg.repository(repo.ui, bundlepath)
 
                     allbundlestocleanup.append((bundlerepo, newbundlefile))
                     bundlerevs = set(_readbundlerevs(bundlerepo))
                     bundlecaps = _includefilelogstobundle(
                         bundlecaps, bundlerepo, bundlerevs, repo.ui
                     )
                     cl = bundlerepo.changelog
                     bundleroots = _getbundleroots(repo, bundlerepo, bundlerevs)
                     for rev in bundlerevs:
                         node = cl.node(rev)
                         newphases[hex(node)] = str(phases.draft)
                         nodestobundle[node] = (
                             bundlerepo,
                             bundleroots,
                             newbundlefile,
                         )
 
                 scratchbundles.append(
                     _generateoutputparts(head, *nodestobundle[head])
                 )
                 newheads.extend(bundleroots)
                 scratchheads.append(head)
     finally:
         for bundlerepo, bundlefile in allbundlestocleanup:
             bundlerepo.close()
             try:
                 os.unlink(bundlefile)
             except (IOError, OSError):
                 # if we can't cleanup the file then just ignore the error,
                 # no need to fail
                 pass
 
     pullfrombundlestore = bool(scratchbundles)
     wrappedchangegrouppart = False
     wrappedlistkeys = False
     oldchangegrouppart = exchange.getbundle2partsmapping[b'changegroup']
     try:
 
         def _changegrouppart(bundler, *args, **kwargs):
             # Order is important here. First add non-scratch part
             # and only then add parts with scratch bundles because
             # non-scratch part contains parents of roots of scratch bundles.
             result = oldchangegrouppart(bundler, *args, **kwargs)
             for bundle in scratchbundles:
                 for part in bundle:
                     bundler.addpart(part)
             return result
 
         exchange.getbundle2partsmapping[b'changegroup'] = _changegrouppart
         wrappedchangegrouppart = True
 
         def _listkeys(orig, self, namespace):
             origvalues = orig(self, namespace)
             if namespace == b'phases' and pullfrombundlestore:
                 if origvalues.get(b'publishing') == b'True':
                     # Make repo non-publishing to preserve draft phase
                     del origvalues[b'publishing']
                 origvalues.update(newphases)
             return origvalues
 
         extensions.wrapfunction(
             localrepo.localrepository, b'listkeys', _listkeys
         )
         wrappedlistkeys = True
         heads = list((set(newheads) | set(heads)) - set(scratchheads))
         result = orig(
             repo, source, heads=heads, bundlecaps=bundlecaps, **kwargs
         )
     finally:
         if wrappedchangegrouppart:
             exchange.getbundle2partsmapping[b'changegroup'] = oldchangegrouppart
         if wrappedlistkeys:
             extensions.unwrapfunction(
                 localrepo.localrepository, b'listkeys', _listkeys
             )
     return result
 
 
 def _lookupwrap(orig):
     def _lookup(repo, proto, key):
         localkey = encoding.tolocal(key)
 
         if isinstance(localkey, str) and _scratchbranchmatcher(localkey):
             scratchnode = repo.bundlestore.index.getnode(localkey)
             if scratchnode:
                 return b"%d %s\n" % (1, scratchnode)
             else:
                 return b"%d %s\n" % (
                     0,
                     b'scratch branch %s not found' % localkey,
                 )
         else:
             try:
                 r = hex(repo.lookup(localkey))
                 return b"%d %s\n" % (1, r)
             except Exception as inst:
                 if repo.bundlestore.index.getbundle(localkey):
                     return b"%d %s\n" % (1, localkey)
                 else:
                     r = stringutil.forcebytestr(inst)
                     return b"%d %s\n" % (0, r)
 
     return _lookup
 
 
 def _pull(orig, ui, repo, source=b"default", **opts):
     opts = pycompat.byteskwargs(opts)
     # Copy paste from `pull` command
     source, branches = urlutil.get_unique_pull_path(
         b"infinite-push's pull",
         repo,
         ui,
         source,
         default_branches=opts.get(b'branch'),
     )
 
     scratchbookmarks = {}
     unfi = repo.unfiltered()
     unknownnodes = []
     for rev in opts.get(b'rev', []):
         if rev not in unfi:
             unknownnodes.append(rev)
     if opts.get(b'bookmark'):
         bookmarks = []
         revs = opts.get(b'rev') or []
         for bookmark in opts.get(b'bookmark'):
             if _scratchbranchmatcher(bookmark):
                 # rev is not known yet
                 # it will be fetched with listkeyspatterns next
                 scratchbookmarks[bookmark] = b'REVTOFETCH'
             else:
                 bookmarks.append(bookmark)
 
         if scratchbookmarks:
             other = hg.peer(repo, opts, source)
             try:
                 fetchedbookmarks = other.listkeyspatterns(
                     b'bookmarks', patterns=scratchbookmarks
                 )
                 for bookmark in scratchbookmarks:
                     if bookmark not in fetchedbookmarks:
                         raise error.Abort(
                             b'remote bookmark %s not found!' % bookmark
                         )
                     scratchbookmarks[bookmark] = fetchedbookmarks[bookmark]
                     revs.append(fetchedbookmarks[bookmark])
             finally:
                 other.close()
         opts[b'bookmark'] = bookmarks
         opts[b'rev'] = revs
 
     if scratchbookmarks or unknownnodes:
         # Set anyincoming to True
         extensions.wrapfunction(
             discovery, b'findcommonincoming', _findcommonincoming
         )
     try:
         # Remote scratch bookmarks will be deleted because remotenames doesn't
         # know about them. Let's save it before pull and restore after
         remotescratchbookmarks = _readscratchremotebookmarks(ui, repo, source)
         result = orig(ui, repo, source, **pycompat.strkwargs(opts))
         # TODO(stash): race condition is possible
         # if scratch bookmarks was updated right after orig.
         # But that's unlikely and shouldn't be harmful.
         if common.isremotebooksenabled(ui):
             remotescratchbookmarks.update(scratchbookmarks)
             _saveremotebookmarks(repo, remotescratchbookmarks, source)
         else:
             _savelocalbookmarks(repo, scratchbookmarks)
         return result
     finally:
         if scratchbookmarks:
             extensions.unwrapfunction(discovery, b'findcommonincoming')
 
 
 def _readscratchremotebookmarks(ui, repo, other):
     if common.isremotebooksenabled(ui):
         remotenamesext = extensions.find(b'remotenames')
         remotepath = remotenamesext.activepath(repo.ui, other)
         result = {}
         # Let's refresh remotenames to make sure we have it up to date
         # Seems that `repo.names['remotebookmarks']` may return stale bookmarks
         # and it results in deleting scratch bookmarks. Our best guess how to
         # fix it is to use `clearnames()`
         repo._remotenames.clearnames()
         for remotebookmark in repo.names[b'remotebookmarks'].listnames(repo):
             path, bookname = remotenamesext.splitremotename(remotebookmark)
             if path == remotepath and _scratchbranchmatcher(bookname):
                 nodes = repo.names[b'remotebookmarks'].nodes(
                     repo, remotebookmark
                 )
                 if nodes:
                     result[bookname] = hex(nodes[0])
         return result
     else:
         return {}
 
 
 def _saveremotebookmarks(repo, newbookmarks, remote):
     remotenamesext = extensions.find(b'remotenames')
     remotepath = remotenamesext.activepath(repo.ui, remote)
     branches = collections.defaultdict(list)
     bookmarks = {}
     remotenames = remotenamesext.readremotenames(repo)
     for hexnode, nametype, remote, rname in remotenames:
         if remote != remotepath:
             continue
         if nametype == b'bookmarks':
             if rname in newbookmarks:
                 # It's possible if we have a normal bookmark that matches
                 # scratch branch pattern. In this case just use the current
                 # bookmark node
                 del newbookmarks[rname]
             bookmarks[rname] = hexnode
         elif nametype == b'branches':
             # saveremotenames expects 20 byte binary nodes for branches
             branches[rname].append(bin(hexnode))
 
     for bookmark, hexnode in pycompat.iteritems(newbookmarks):
         bookmarks[bookmark] = hexnode
     remotenamesext.saveremotenames(repo, remotepath, branches, bookmarks)
 
 
 def _savelocalbookmarks(repo, bookmarks):
     if not bookmarks:
         return
     with repo.wlock(), repo.lock(), repo.transaction(b'bookmark') as tr:
         changes = []
         for scratchbook, node in pycompat.iteritems(bookmarks):
             changectx = repo[node]
             changes.append((scratchbook, changectx.node()))
         repo._bookmarks.applychanges(repo, tr, changes)
 
 
 def _findcommonincoming(orig, *args, **kwargs):
     common, inc, remoteheads = orig(*args, **kwargs)
     return common, True, remoteheads
 
 
 def _push(orig, ui, repo, *dests, **opts):
     opts = pycompat.byteskwargs(opts)
     bookmark = opts.get(b'bookmark')
     # we only support pushing one infinitepush bookmark at once
     if len(bookmark) == 1:
         bookmark = bookmark[0]
     else:
         bookmark = b''
 
     oldphasemove = None
     overrides = {(experimental, configbookmark): bookmark}
 
     with ui.configoverride(overrides, b'infinitepush'):
         scratchpush = opts.get(b'bundle_store')
         if _scratchbranchmatcher(bookmark):
             scratchpush = True
             # bundle2 can be sent back after push (for example, bundle2
             # containing `pushkey` part to update bookmarks)
             ui.setconfig(experimental, b'bundle2.pushback', True)
 
         if scratchpush:
             # this is an infinitepush, we don't want the bookmark to be applied
             # rather that should be stored in the bundlestore
             opts[b'bookmark'] = []
             ui.setconfig(experimental, configscratchpush, True)
             oldphasemove = extensions.wrapfunction(
                 exchange, b'_localphasemove', _phasemove
             )
 
         paths = list(urlutil.get_push_paths(repo, ui, dests))
         if len(paths) > 1:
             msg = _(b'cannot push to multiple path with infinitepush')
             raise error.Abort(msg)
 
         path = paths[0]
         destpath = path.pushloc or path.loc
         # Remote scratch bookmarks will be deleted because remotenames doesn't
         # know about them. Let's save it before push and restore after
         remotescratchbookmarks = _readscratchremotebookmarks(ui, repo, destpath)
         result = orig(ui, repo, *dests, **pycompat.strkwargs(opts))
         if common.isremotebooksenabled(ui):
             if bookmark and scratchpush:
                 other = hg.peer(repo, opts, destpath)
                 try:
                     fetchedbookmarks = other.listkeyspatterns(
                         b'bookmarks', patterns=[bookmark]
                     )
                     remotescratchbookmarks.update(fetchedbookmarks)
                 finally:
                     other.close()
             _saveremotebookmarks(repo, remotescratchbookmarks, destpath)
     if oldphasemove:
         exchange._localphasemove = oldphasemove
     return result
 
 
 def _deleteinfinitepushbookmarks(ui, repo, path, names):
     """Prune remote names by removing the bookmarks we don't want anymore,
     then writing the result back to disk
     """
     remotenamesext = extensions.find(b'remotenames')
 
     # remotename format is:
     # (node, nametype ("branches" or "bookmarks"), remote, name)
     nametype_idx = 1
     remote_idx = 2
     name_idx = 3
     remotenames = [
         remotename
         for remotename in remotenamesext.readremotenames(repo)
         if remotename[remote_idx] == path
     ]
     remote_bm_names = [
         remotename[name_idx]
         for remotename in remotenames
         if remotename[nametype_idx] == b"bookmarks"
     ]
 
     for name in names:
         if name not in remote_bm_names:
             raise error.Abort(
                 _(
                     b"infinitepush bookmark '{}' does not exist "
                     b"in path '{}'"
                 ).format(name, path)
             )
 
     bookmarks = {}
     branches = collections.defaultdict(list)
     for node, nametype, remote, name in remotenames:
         if nametype == b"bookmarks" and name not in names:
             bookmarks[name] = node
         elif nametype == b"branches":
             # saveremotenames wants binary nodes for branches
             branches[name].append(bin(node))
 
     remotenamesext.saveremotenames(repo, path, branches, bookmarks)
 
 
 def _phasemove(orig, pushop, nodes, phase=phases.public):
     """prevent commits from being marked public
 
     Since these are going to a scratch branch, they aren't really being
     published."""
 
     if phase != phases.public:
         orig(pushop, nodes, phase)
 
 
 @exchange.b2partsgenerator(scratchbranchparttype)
 def partgen(pushop, bundler):
     bookmark = pushop.ui.config(experimental, configbookmark)
     scratchpush = pushop.ui.configbool(experimental, configscratchpush)
     if b'changesets' in pushop.stepsdone or not scratchpush:
         return
 
     if scratchbranchparttype not in bundle2.bundle2caps(pushop.remote):
         return
 
     pushop.stepsdone.add(b'changesets')
     if not pushop.outgoing.missing:
         pushop.ui.status(_(b'no changes found\n'))
         pushop.cgresult = 0
         return
 
     # This parameter tells the server that the following bundle is an
     # infinitepush. This let's it switch the part processing to our infinitepush
     # code path.
     bundler.addparam(b"infinitepush", b"True")
 
     scratchparts = bundleparts.getscratchbranchparts(
         pushop.repo, pushop.remote, pushop.outgoing, pushop.ui, bookmark
     )
 
     for scratchpart in scratchparts:
         bundler.addpart(scratchpart)
 
     def handlereply(op):
         # server either succeeds or aborts; no code to read
         pushop.cgresult = 1
 
     return handlereply
 
 
 bundle2.capabilities[bundleparts.scratchbranchparttype] = ()
 
 
 def _getrevs(bundle, oldnode, force, bookmark):
     b'extracts and validates the revs to be imported'
     revs = [bundle[r] for r in bundle.revs(b'sort(bundle())')]
 
     # new bookmark
     if oldnode is None:
         return revs
 
     # Fast forward update
     if oldnode in bundle and list(bundle.set(b'bundle() & %s::', oldnode)):
         return revs
 
     return revs
 
 
 @contextlib.contextmanager
 def logservicecall(logger, service, **kwargs):
     start = time.time()
     logger(service, eventtype=b'start', **kwargs)
     try:
         yield
         logger(
             service,
             eventtype=b'success',
             elapsedms=(time.time() - start) * 1000,
             **kwargs
         )
     except Exception as e:
         logger(
             service,
             eventtype=b'failure',
             elapsedms=(time.time() - start) * 1000,
             errormsg=stringutil.forcebytestr(e),
             **kwargs
         )
         raise
 
 
 def _getorcreateinfinitepushlogger(op):
     logger = op.records[b'infinitepushlogger']
     if not logger:
         ui = op.repo.ui
         try:
             username = procutil.getuser()
         except Exception:
             username = b'unknown'
         # Generate random request id to be able to find all logged entries
         # for the same request. Since requestid is pseudo-generated it may
         # not be unique, but we assume that (hostname, username, requestid)
         # is unique.
         random.seed()
         requestid = random.randint(0, 2000000000)
         hostname = socket.gethostname()
         logger = functools.partial(
             ui.log,
             b'infinitepush',
             user=username,
             requestid=requestid,
             hostname=hostname,
             reponame=ui.config(b'infinitepush', b'reponame'),
         )
         op.records.add(b'infinitepushlogger', logger)
     else:
         logger = logger[0]
     return logger
 
 
 def storetobundlestore(orig, repo, op, unbundler):
     """stores the incoming bundle coming from push command to the bundlestore
     instead of applying on the revlogs"""
 
     repo.ui.status(_(b"storing changesets on the bundlestore\n"))
     bundler = bundle2.bundle20(repo.ui)
 
     # processing each part and storing it in bundler
     with bundle2.partiterator(repo, op, unbundler) as parts:
         for part in parts:
             bundlepart = None
             if part.type == b'replycaps':
                 # This configures the current operation to allow reply parts.
                 bundle2._processpart(op, part)
             else:
                 bundlepart = bundle2.bundlepart(part.type, data=part.read())
                 for key, value in pycompat.iteritems(part.params):
                     bundlepart.addparam(key, value)
 
                 # Certain parts require a response
                 if part.type in (b'pushkey', b'changegroup'):
                     if op.reply is not None:
                         rpart = op.reply.newpart(b'reply:%s' % part.type)
                         rpart.addparam(
                             b'in-reply-to', b'%d' % part.id, mandatory=False
                         )
                         rpart.addparam(b'return', b'1', mandatory=False)
 
             op.records.add(
                 part.type,
                 {
                     b'return': 1,
                 },
             )
             if bundlepart:
                 bundler.addpart(bundlepart)
 
     # storing the bundle in the bundlestore
     buf = util.chunkbuffer(bundler.getchunks())
     fd, bundlefile = pycompat.mkstemp()
     try:
         try:
             fp = os.fdopen(fd, 'wb')
             fp.write(buf.read())
         finally:
             fp.close()
         storebundle(op, {}, bundlefile)
     finally:
         try:
             os.unlink(bundlefile)
         except Exception:
             # we would rather see the original exception
             pass
 
 
 def processparts(orig, repo, op, unbundler):
 
     # make sure we don't wrap processparts in case of `hg unbundle`
     if op.source == b'unbundle':
         return orig(repo, op, unbundler)
 
     # this server routes each push to bundle store
     if repo.ui.configbool(b'infinitepush', b'pushtobundlestore'):
         return storetobundlestore(orig, repo, op, unbundler)
 
     if unbundler.params.get(b'infinitepush') != b'True':
         return orig(repo, op, unbundler)
 
     handleallparts = repo.ui.configbool(b'infinitepush', b'storeallparts')
 
     bundler = bundle2.bundle20(repo.ui)
     cgparams = None
     with bundle2.partiterator(repo, op, unbundler) as parts:
         for part in parts:
             bundlepart = None
             if part.type == b'replycaps':
                 # This configures the current operation to allow reply parts.
                 bundle2._processpart(op, part)
             elif part.type == bundleparts.scratchbranchparttype:
                 # Scratch branch parts need to be converted to normal
                 # changegroup parts, and the extra parameters stored for later
                 # when we upload to the store. Eventually those parameters will
                 # be put on the actual bundle instead of this part, then we can
                 # send a vanilla changegroup instead of the scratchbranch part.
                 cgversion = part.params.get(b'cgversion', b'01')
                 bundlepart = bundle2.bundlepart(
                     b'changegroup', data=part.read()
                 )
                 bundlepart.addparam(b'version', cgversion)
                 cgparams = part.params
 
                 # If we're not dumping all parts into the new bundle, we need to
                 # alert the future pushkey and phase-heads handler to skip
                 # the part.
                 if not handleallparts:
                     op.records.add(
                         scratchbranchparttype + b'_skippushkey', True
                     )
                     op.records.add(
                         scratchbranchparttype + b'_skipphaseheads', True
                     )
             else:
                 if handleallparts:
                     # Ideally we would not process any parts, and instead just
                     # forward them to the bundle for storage, but since this
                     # differs from previous behavior, we need to put it behind a
                     # config flag for incremental rollout.
                     bundlepart = bundle2.bundlepart(part.type, data=part.read())
                     for key, value in pycompat.iteritems(part.params):
                         bundlepart.addparam(key, value)
 
                     # Certain parts require a response
                     if part.type == b'pushkey':
                         if op.reply is not None:
                             rpart = op.reply.newpart(b'reply:pushkey')
                             rpart.addparam(
                                 b'in-reply-to', str(part.id), mandatory=False
                             )
                             rpart.addparam(b'return', b'1', mandatory=False)
                 else:
                     bundle2._processpart(op, part)
 
             if handleallparts:
                 op.records.add(
                     part.type,
                     {
                         b'return': 1,
                     },
                 )
             if bundlepart:
                 bundler.addpart(bundlepart)
 
     # If commits were sent, store them
     if cgparams:
         buf = util.chunkbuffer(bundler.getchunks())
         fd, bundlefile = pycompat.mkstemp()
         try:
             try:
                 fp = os.fdopen(fd, 'wb')
                 fp.write(buf.read())
             finally:
                 fp.close()
             storebundle(op, cgparams, bundlefile)
         finally:
             try:
                 os.unlink(bundlefile)
             except Exception:
                 # we would rather see the original exception
                 pass
 
 
 def storebundle(op, params, bundlefile):
     log = _getorcreateinfinitepushlogger(op)
     parthandlerstart = time.time()
     log(scratchbranchparttype, eventtype=b'start')
     index = op.repo.bundlestore.index
     store = op.repo.bundlestore.store
     op.records.add(scratchbranchparttype + b'_skippushkey', True)
 
     bundle = None
     try:  # guards bundle
         bundlepath = b"bundle:%s+%s" % (op.repo.root, bundlefile)
         bundle = hg.repository(op.repo.ui, bundlepath)
 
         bookmark = params.get(b'bookmark')
         bookprevnode = params.get(b'bookprevnode', b'')
         force = params.get(b'force')
 
         if bookmark:
             oldnode = index.getnode(bookmark)
         else:
             oldnode = None
         bundleheads = bundle.revs(b'heads(bundle())')
         if bookmark and len(bundleheads) > 1:
             raise error.Abort(
                 _(b'cannot push more than one head to a scratch branch')
             )
 
         revs = _getrevs(bundle, oldnode, force, bookmark)
 
         # Notify the user of what is being pushed
         plural = b's' if len(revs) > 1 else b''
         op.repo.ui.warn(_(b"pushing %d commit%s:\n") % (len(revs), plural))
         maxoutput = 10
         for i in range(0, min(len(revs), maxoutput)):
             firstline = bundle[revs[i]].description().split(b'\n')[0][:50]
             op.repo.ui.warn(b"    %s  %s\n" % (revs[i], firstline))
 
         if len(revs) > maxoutput + 1:
             op.repo.ui.warn(b"    ...\n")
             firstline = bundle[revs[-1]].description().split(b'\n')[0][:50]
             op.repo.ui.warn(b"    %s  %s\n" % (revs[-1], firstline))
 
         nodesctx = [bundle[rev] for rev in revs]
         inindex = lambda rev: bool(index.getbundle(bundle[rev].hex()))
         if bundleheads:
             newheadscount = sum(not inindex(rev) for rev in bundleheads)
         else:
             newheadscount = 0
         # If there's a bookmark specified, there should be only one head,
         # so we choose the last node, which will be that head.
         # If a bug or malicious client allows there to be a bookmark
         # with multiple heads, we will place the bookmark on the last head.
         bookmarknode = nodesctx[-1].hex() if nodesctx else None
         key = None
         if newheadscount:
             with open(bundlefile, b'rb') as f:
                 bundledata = f.read()
                 with logservicecall(
                     log, b'bundlestore', bundlesize=len(bundledata)
                 ):
                     bundlesizelimit = 100 * 1024 * 1024  # 100 MB
                     if len(bundledata) > bundlesizelimit:
                         error_msg = (
                             b'bundle is too big: %d bytes. '
                             + b'max allowed size is 100 MB'
                         )
                         raise error.Abort(error_msg % (len(bundledata),))
                     key = store.write(bundledata)
 
         with logservicecall(log, b'index', newheadscount=newheadscount), index:
             if key:
                 index.addbundle(key, nodesctx)
             if bookmark:
                 index.addbookmark(bookmark, bookmarknode)
                 _maybeaddpushbackpart(
                     op, bookmark, bookmarknode, bookprevnode, params
                 )
         log(
             scratchbranchparttype,
             eventtype=b'success',
             elapsedms=(time.time() - parthandlerstart) * 1000,
         )
 
     except Exception as e:
         log(
             scratchbranchparttype,
             eventtype=b'failure',
             elapsedms=(time.time() - parthandlerstart) * 1000,
             errormsg=stringutil.forcebytestr(e),
         )
         raise
     finally:
         if bundle:
             bundle.close()
 
 
 @bundle2.parthandler(
     scratchbranchparttype,
     (
         b'bookmark',
         b'bookprevnode',
         b'force',
         b'pushbackbookmarks',
         b'cgversion',
     ),
 )
 def bundle2scratchbranch(op, part):
     '''unbundle a bundle2 part containing a changegroup to store'''
 
     bundler = bundle2.bundle20(op.repo.ui)
     cgversion = part.params.get(b'cgversion', b'01')
     cgpart = bundle2.bundlepart(b'changegroup', data=part.read())
     cgpart.addparam(b'version', cgversion)
     bundler.addpart(cgpart)
     buf = util.chunkbuffer(bundler.getchunks())
 
     fd, bundlefile = pycompat.mkstemp()
     try:
         try:
             fp = os.fdopen(fd, 'wb')
             fp.write(buf.read())
         finally:
             fp.close()
         storebundle(op, part.params, bundlefile)
     finally:
         try:
             os.unlink(bundlefile)
         except OSError as e:
             if e.errno != errno.ENOENT:
                 raise
 
     return 1
 
 
 def _maybeaddpushbackpart(op, bookmark, newnode, oldnode, params):
     if params.get(b'pushbackbookmarks'):
         if op.reply and b'pushback' in op.reply.capabilities:
             params = {
                 b'namespace': b'bookmarks',
                 b'key': bookmark,
                 b'new': newnode,
                 b'old': oldnode,
             }
             op.reply.newpart(
                 b'pushkey', mandatoryparams=pycompat.iteritems(params)
             )
 
 
 def bundle2pushkey(orig, op, part):
     """Wrapper of bundle2.handlepushkey()
 
     The only goal is to skip calling the original function if flag is set.
     It's set if infinitepush push is happening.
     """
     if op.records[scratchbranchparttype + b'_skippushkey']:
         if op.reply is not None:
             rpart = op.reply.newpart(b'reply:pushkey')
             rpart.addparam(b'in-reply-to', str(part.id), mandatory=False)
             rpart.addparam(b'return', b'1', mandatory=False)
         return 1
 
     return orig(op, part)
 
 
 def bundle2handlephases(orig, op, part):
     """Wrapper of bundle2.handlephases()
 
     The only goal is to skip calling the original function if flag is set.
     It's set if infinitepush push is happening.
     """
 
     if op.records[scratchbranchparttype + b'_skipphaseheads']:
         return
 
     return orig(op, part)
 
 
 def _asyncsavemetadata(root, nodes):
     """starts a separate process that fills metadata for the nodes
 
     This function creates a separate process and doesn't wait for it's
     completion. This was done to avoid slowing down pushes
     """
 
     maxnodes = 50
     if len(nodes) > maxnodes:
         return
     nodesargs = []
     for node in nodes:
         nodesargs.append(b'--node')
         nodesargs.append(node)
     with open(os.devnull, b'w+b') as devnull:
         cmdline = [
             util.hgexecutable(),
             b'debugfillinfinitepushmetadata',
             b'-R',
             root,
         ] + nodesargs
         # Process will run in background. We don't care about the return code
         subprocess.Popen(
             pycompat.rapply(procutil.tonativestr, cmdline),
             close_fds=True,
             shell=False,
             stdin=devnull,
             stdout=devnull,
             stderr=devnull,
         )

hgext/largefiles/proto.py

@@ -1,217 +1,218 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.
 from __future__ import absolute_import
 
 import os
 
 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()
 
 
 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
+            def decode(d):
             try:
-                yield int(f.value)
+                    return int(d)
             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
+                    return 2
+
+            result = {b'sha': sha}
+            return result, decode
 
     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)

hgext/remotefilelog/fileserverclient.py

@@ -1,670 +1,673 b''
 # fileserverclient.py - client for communicating with the cache process
 #
 # Copyright 2013 Facebook, 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 io
 import os
 import threading
 import time
 import zlib
 
 from mercurial.i18n import _
 from mercurial.node import bin, hex
 from mercurial import (
     error,
     pycompat,
     revlog,
     sshpeer,
     util,
     wireprotov1peer,
 )
 from mercurial.utils import (
     hashutil,
     procutil,
 )
 
 from . import (
     constants,
     contentstore,
     metadatastore,
 )
 
 _sshv1peer = sshpeer.sshv1peer
 
 # Statistics for debugging
 fetchcost = 0
 fetches = 0
 fetched = 0
 fetchmisses = 0
 
 _lfsmod = None
 
 
 def getcachekey(reponame, file, id):
     pathhash = hex(hashutil.sha1(file).digest())
     return os.path.join(reponame, pathhash[:2], pathhash[2:], id)
 
 
 def getlocalkey(file, id):
     pathhash = hex(hashutil.sha1(file).digest())
     return os.path.join(pathhash, id)
 
 
 def peersetup(ui, peer):
     class remotefilepeer(peer.__class__):
         @wireprotov1peer.batchable
         def x_rfl_getfile(self, file, node):
             if not self.capable(b'x_rfl_getfile'):
                 raise error.Abort(
                     b'configured remotefile server does not support getfile'
                 )
-            f = wireprotov1peer.future()
-            yield {b'file': file, b'node': node}, f
-            code, data = f.value.split(b'\0', 1)
+
+            def decode(d):
+                code, data = d.split(b'\0', 1)
             if int(code):
                 raise error.LookupError(file, node, data)
-            yield data
+                return data
+
+            return {b'file': file, b'node': node}, decode
 
         @wireprotov1peer.batchable
         def x_rfl_getflogheads(self, path):
             if not self.capable(b'x_rfl_getflogheads'):
                 raise error.Abort(
                     b'configured remotefile server does not '
                     b'support getflogheads'
                 )
-            f = wireprotov1peer.future()
-            yield {b'path': path}, f
-            heads = f.value.split(b'\n') if f.value else []
-            yield heads
+
+            def decode(d):
+                return d.split(b'\n') if d else []
+
+            return {b'path': path}, decode
 
         def _updatecallstreamopts(self, command, opts):
             if command != b'getbundle':
                 return
             if (
                 constants.NETWORK_CAP_LEGACY_SSH_GETFILES
                 not in self.capabilities()
             ):
                 return
             if not util.safehasattr(self, '_localrepo'):
                 return
             if (
                 constants.SHALLOWREPO_REQUIREMENT
                 not in self._localrepo.requirements
             ):
                 return
 
             bundlecaps = opts.get(b'bundlecaps')
             if bundlecaps:
                 bundlecaps = [bundlecaps]
             else:
                 bundlecaps = []
 
             # shallow, includepattern, and excludepattern are a hacky way of
             # carrying over data from the local repo to this getbundle
             # command. We need to do it this way because bundle1 getbundle
             # doesn't provide any other place we can hook in to manipulate
             # getbundle args before it goes across the wire. Once we get rid
             # of bundle1, we can use bundle2's _pullbundle2extraprepare to
             # do this more cleanly.
             bundlecaps.append(constants.BUNDLE2_CAPABLITY)
             if self._localrepo.includepattern:
                 patterns = b'\0'.join(self._localrepo.includepattern)
                 includecap = b"includepattern=" + patterns
                 bundlecaps.append(includecap)
             if self._localrepo.excludepattern:
                 patterns = b'\0'.join(self._localrepo.excludepattern)
                 excludecap = b"excludepattern=" + patterns
                 bundlecaps.append(excludecap)
             opts[b'bundlecaps'] = b','.join(bundlecaps)
 
         def _sendrequest(self, command, args, **opts):
             self._updatecallstreamopts(command, args)
             return super(remotefilepeer, self)._sendrequest(
                 command, args, **opts
             )
 
         def _callstream(self, command, **opts):
             supertype = super(remotefilepeer, self)
             if not util.safehasattr(supertype, '_sendrequest'):
                 self._updatecallstreamopts(command, pycompat.byteskwargs(opts))
             return super(remotefilepeer, self)._callstream(command, **opts)
 
     peer.__class__ = remotefilepeer
 
 
 class cacheconnection(object):
     """The connection for communicating with the remote cache. Performs
     gets and sets by communicating with an external process that has the
     cache-specific implementation.
     """
 
     def __init__(self):
         self.pipeo = self.pipei = self.pipee = None
         self.subprocess = None
         self.connected = False
 
     def connect(self, cachecommand):
         if self.pipeo:
             raise error.Abort(_(b"cache connection already open"))
         self.pipei, self.pipeo, self.pipee, self.subprocess = procutil.popen4(
             cachecommand
         )
         self.connected = True
 
     def close(self):
         def tryclose(pipe):
             try:
                 pipe.close()
             except Exception:
                 pass
 
         if self.connected:
             try:
                 self.pipei.write(b"exit\n")
             except Exception:
                 pass
             tryclose(self.pipei)
             self.pipei = None
             tryclose(self.pipeo)
             self.pipeo = None
             tryclose(self.pipee)
             self.pipee = None
             try:
                 # Wait for process to terminate, making sure to avoid deadlock.
                 # See https://docs.python.org/2/library/subprocess.html for
                 # warnings about wait() and deadlocking.
                 self.subprocess.communicate()
             except Exception:
                 pass
             self.subprocess = None
         self.connected = False
 
     def request(self, request, flush=True):
         if self.connected:
             try:
                 self.pipei.write(request)
                 if flush:
                     self.pipei.flush()
             except IOError:
                 self.close()
 
     def receiveline(self):
         if not self.connected:
             return None
         try:
             result = self.pipeo.readline()[:-1]
             if not result:
                 self.close()
         except IOError:
             self.close()
 
         return result
 
 
 def _getfilesbatch(
     remote, receivemissing, progresstick, missed, idmap, batchsize
 ):
     # Over http(s), iterbatch is a streamy method and we can start
     # looking at results early. This means we send one (potentially
     # large) request, but then we show nice progress as we process
     # file results, rather than showing chunks of $batchsize in
     # progress.
     #
     # Over ssh, iterbatch isn't streamy because batch() wasn't
     # explicitly designed as a streaming method. In the future we
     # should probably introduce a streambatch() method upstream and
     # use that for this.
     with remote.commandexecutor() as e:
         futures = []
         for m in missed:
             futures.append(
                 e.callcommand(
                     b'x_rfl_getfile', {b'file': idmap[m], b'node': m[-40:]}
                 )
             )
 
         for i, m in enumerate(missed):
             r = futures[i].result()
             futures[i] = None  # release memory
             file_ = idmap[m]
             node = m[-40:]
             receivemissing(io.BytesIO(b'%d\n%s' % (len(r), r)), file_, node)
             progresstick()
 
 
 def _getfiles_optimistic(
     remote, receivemissing, progresstick, missed, idmap, step
 ):
     remote._callstream(b"x_rfl_getfiles")
     i = 0
     pipeo = remote._pipeo
     pipei = remote._pipei
     while i < len(missed):
         # issue a batch of requests
         start = i
         end = min(len(missed), start + step)
         i = end
         for missingid in missed[start:end]:
             # issue new request
             versionid = missingid[-40:]
             file = idmap[missingid]
             sshrequest = b"%s%s\n" % (versionid, file)
             pipeo.write(sshrequest)
         pipeo.flush()
 
         # receive batch results
         for missingid in missed[start:end]:
             versionid = missingid[-40:]
             file = idmap[missingid]
             receivemissing(pipei, file, versionid)
             progresstick()
 
     # End the command
     pipeo.write(b'\n')
     pipeo.flush()
 
 
 def _getfiles_threaded(
     remote, receivemissing, progresstick, missed, idmap, step
 ):
     remote._callstream(b"x_rfl_getfiles")
     pipeo = remote._pipeo
     pipei = remote._pipei
 
     def writer():
         for missingid in missed:
             versionid = missingid[-40:]
             file = idmap[missingid]
             sshrequest = b"%s%s\n" % (versionid, file)
             pipeo.write(sshrequest)
         pipeo.flush()
 
     writerthread = threading.Thread(target=writer)
     writerthread.daemon = True
     writerthread.start()
 
     for missingid in missed:
         versionid = missingid[-40:]
         file = idmap[missingid]
         receivemissing(pipei, file, versionid)
         progresstick()
 
     writerthread.join()
     # End the command
     pipeo.write(b'\n')
     pipeo.flush()
 
 
 class fileserverclient(object):
     """A client for requesting files from the remote file server."""
 
     def __init__(self, repo):
         ui = repo.ui
         self.repo = repo
         self.ui = ui
         self.cacheprocess = ui.config(b"remotefilelog", b"cacheprocess")
         if self.cacheprocess:
             self.cacheprocess = util.expandpath(self.cacheprocess)
 
         # This option causes remotefilelog to pass the full file path to the
         # cacheprocess instead of a hashed key.
         self.cacheprocesspasspath = ui.configbool(
             b"remotefilelog", b"cacheprocess.includepath"
         )
 
         self.debugoutput = ui.configbool(b"remotefilelog", b"debug")
 
         self.remotecache = cacheconnection()
 
     def setstore(self, datastore, historystore, writedata, writehistory):
         self.datastore = datastore
         self.historystore = historystore
         self.writedata = writedata
         self.writehistory = writehistory
 
     def _connect(self):
         return self.repo.connectionpool.get(self.repo.fallbackpath)
 
     def request(self, fileids):
         """Takes a list of filename/node pairs and fetches them from the
         server. Files are stored in the local cache.
         A list of nodes that the server couldn't find is returned.
         If the connection fails, an exception is raised.
         """
         if not self.remotecache.connected:
             self.connect()
         cache = self.remotecache
         writedata = self.writedata
 
         repo = self.repo
         total = len(fileids)
         request = b"get\n%d\n" % total
         idmap = {}
         reponame = repo.name
         for file, id in fileids:
             fullid = getcachekey(reponame, file, id)
             if self.cacheprocesspasspath:
                 request += file + b'\0'
             request += fullid + b"\n"
             idmap[fullid] = file
 
         cache.request(request)
 
         progress = self.ui.makeprogress(_(b'downloading'), total=total)
         progress.update(0)
 
         missed = []
         while True:
             missingid = cache.receiveline()
             if not missingid:
                 missedset = set(missed)
                 for missingid in idmap:
                     if not missingid in missedset:
                         missed.append(missingid)
                 self.ui.warn(
                     _(
                         b"warning: cache connection closed early - "
                         + b"falling back to server\n"
                     )
                 )
                 break
             if missingid == b"0":
                 break
             if missingid.startswith(b"_hits_"):
                 # receive progress reports
                 parts = missingid.split(b"_")
                 progress.increment(int(parts[2]))
                 continue
 
             missed.append(missingid)
 
         global fetchmisses
         fetchmisses += len(missed)
 
         fromcache = total - len(missed)
         progress.update(fromcache, total=total)
         self.ui.log(
             b"remotefilelog",
             b"remote cache hit rate is %r of %r\n",
             fromcache,
             total,
             hit=fromcache,
             total=total,
         )
 
         oldumask = os.umask(0o002)
         try:
             # receive cache misses from master
             if missed:
                 # When verbose is true, sshpeer prints 'running ssh...'
                 # to stdout, which can interfere with some command
                 # outputs
                 verbose = self.ui.verbose
                 self.ui.verbose = False
                 try:
                     with self._connect() as conn:
                         remote = conn.peer
                         if remote.capable(
                             constants.NETWORK_CAP_LEGACY_SSH_GETFILES
                         ):
                             if not isinstance(remote, _sshv1peer):
                                 raise error.Abort(
                                     b'remotefilelog requires ssh servers'
                                 )
                             step = self.ui.configint(
                                 b'remotefilelog', b'getfilesstep'
                             )
                             getfilestype = self.ui.config(
                                 b'remotefilelog', b'getfilestype'
                             )
                             if getfilestype == b'threaded':
                                 _getfiles = _getfiles_threaded
                             else:
                                 _getfiles = _getfiles_optimistic
                             _getfiles(
                                 remote,
                                 self.receivemissing,
                                 progress.increment,
                                 missed,
                                 idmap,
                                 step,
                             )
                         elif remote.capable(b"x_rfl_getfile"):
                             if remote.capable(b'batch'):
                                 batchdefault = 100
                             else:
                                 batchdefault = 10
                             batchsize = self.ui.configint(
                                 b'remotefilelog', b'batchsize', batchdefault
                             )
                             self.ui.debug(
                                 b'requesting %d files from '
                                 b'remotefilelog server...\n' % len(missed)
                             )
                             _getfilesbatch(
                                 remote,
                                 self.receivemissing,
                                 progress.increment,
                                 missed,
                                 idmap,
                                 batchsize,
                             )
                         else:
                             raise error.Abort(
                                 b"configured remotefilelog server"
                                 b" does not support remotefilelog"
                             )
 
                     self.ui.log(
                         b"remotefilefetchlog",
                         b"Success\n",
                         fetched_files=progress.pos - fromcache,
                         total_to_fetch=total - fromcache,
                     )
                 except Exception:
                     self.ui.log(
                         b"remotefilefetchlog",
                         b"Fail\n",
                         fetched_files=progress.pos - fromcache,
                         total_to_fetch=total - fromcache,
                     )
                     raise
                 finally:
                     self.ui.verbose = verbose
                 # send to memcache
                 request = b"set\n%d\n%s\n" % (len(missed), b"\n".join(missed))
                 cache.request(request)
 
             progress.complete()
 
             # mark ourselves as a user of this cache
             writedata.markrepo(self.repo.path)
         finally:
             os.umask(oldumask)
 
     def receivemissing(self, pipe, filename, node):
         line = pipe.readline()[:-1]
         if not line:
             raise error.ResponseError(
                 _(b"error downloading file contents:"),
                 _(b"connection closed early"),
             )
         size = int(line)
         data = pipe.read(size)
         if len(data) != size:
             raise error.ResponseError(
                 _(b"error downloading file contents:"),
                 _(b"only received %s of %s bytes") % (len(data), size),
             )
 
         self.writedata.addremotefilelognode(
             filename, bin(node), zlib.decompress(data)
         )
 
     def connect(self):
         if self.cacheprocess:
             cmd = b"%s %s" % (self.cacheprocess, self.writedata._path)
             self.remotecache.connect(cmd)
         else:
             # If no cache process is specified, we fake one that always
             # returns cache misses.  This enables tests to run easily
             # and may eventually allow us to be a drop in replacement
             # for the largefiles extension.
             class simplecache(object):
                 def __init__(self):
                     self.missingids = []
                     self.connected = True
 
                 def close(self):
                     pass
 
                 def request(self, value, flush=True):
                     lines = value.split(b"\n")
                     if lines[0] != b"get":
                         return
                     self.missingids = lines[2:-1]
                     self.missingids.append(b'0')
 
                 def receiveline(self):
                     if len(self.missingids) > 0:
                         return self.missingids.pop(0)
                     return None
 
             self.remotecache = simplecache()
 
     def close(self):
         if fetches:
             msg = (
                 b"%d files fetched over %d fetches - "
                 + b"(%d misses, %0.2f%% hit ratio) over %0.2fs\n"
             ) % (
                 fetched,
                 fetches,
                 fetchmisses,
                 float(fetched - fetchmisses) / float(fetched) * 100.0,
                 fetchcost,
             )
             if self.debugoutput:
                 self.ui.warn(msg)
             self.ui.log(
                 b"remotefilelog.prefetch",
                 msg.replace(b"%", b"%%"),
                 remotefilelogfetched=fetched,
                 remotefilelogfetches=fetches,
                 remotefilelogfetchmisses=fetchmisses,
                 remotefilelogfetchtime=fetchcost * 1000,
             )
 
         if self.remotecache.connected:
             self.remotecache.close()
 
     def prefetch(
         self, fileids, force=False, fetchdata=True, fetchhistory=False
     ):
         """downloads the given file versions to the cache"""
         repo = self.repo
         idstocheck = []
         for file, id in fileids:
             # hack
             # - we don't use .hgtags
             # - workingctx produces ids with length 42,
             #   which we skip since they aren't in any cache
             if (
                 file == b'.hgtags'
                 or len(id) == 42
                 or not repo.shallowmatch(file)
             ):
                 continue
 
             idstocheck.append((file, bin(id)))
 
         datastore = self.datastore
         historystore = self.historystore
         if force:
             datastore = contentstore.unioncontentstore(*repo.shareddatastores)
             historystore = metadatastore.unionmetadatastore(
                 *repo.sharedhistorystores
             )
 
         missingids = set()
         if fetchdata:
             missingids.update(datastore.getmissing(idstocheck))
         if fetchhistory:
             missingids.update(historystore.getmissing(idstocheck))
 
         # partition missing nodes into nullid and not-nullid so we can
         # warn about this filtering potentially shadowing bugs.
         nullids = len(
             [None for unused, id in missingids if id == self.repo.nullid]
         )
         if nullids:
             missingids = [
                 (f, id) for f, id in missingids if id != self.repo.nullid
             ]
             repo.ui.develwarn(
                 (
                     b'remotefilelog not fetching %d null revs'
                     b' - this is likely hiding bugs' % nullids
                 ),
                 config=b'remotefilelog-ext',
             )
         if missingids:
             global fetches, fetched, fetchcost
             fetches += 1
 
             # We want to be able to detect excess individual file downloads, so
             # let's log that information for debugging.
             if fetches >= 15 and fetches < 18:
                 if fetches == 15:
                     fetchwarning = self.ui.config(
                         b'remotefilelog', b'fetchwarning'
                     )
                     if fetchwarning:
                         self.ui.warn(fetchwarning + b'\n')
                 self.logstacktrace()
             missingids = [(file, hex(id)) for file, id in sorted(missingids)]
             fetched += len(missingids)
             start = time.time()
             missingids = self.request(missingids)
             if missingids:
                 raise error.Abort(
                     _(b"unable to download %d files") % len(missingids)
                 )
             fetchcost += time.time() - start
             self._lfsprefetch(fileids)
 
     def _lfsprefetch(self, fileids):
         if not _lfsmod or not util.safehasattr(
             self.repo.svfs, b'lfslocalblobstore'
         ):
             return
         if not _lfsmod.wrapper.candownload(self.repo):
             return
         pointers = []
         store = self.repo.svfs.lfslocalblobstore
         for file, id in fileids:
             node = bin(id)
             rlog = self.repo.file(file)
             if rlog.flags(node) & revlog.REVIDX_EXTSTORED:
                 text = rlog.rawdata(node)
                 p = _lfsmod.pointer.deserialize(text)
                 oid = p.oid()
                 if not store.has(oid):
                     pointers.append(p)
         if len(pointers) > 0:
             self.repo.svfs.lfsremoteblobstore.readbatch(pointers, store)
             assert all(store.has(p.oid()) for p in pointers)
 
     def logstacktrace(self):
         import traceback
 
         self.ui.log(
             b'remotefilelog',
             b'excess remotefilelog fetching:\n%s\n',
             b''.join(pycompat.sysbytes(s) for s in traceback.format_stack()),
         )

mercurial/wireprotov1peer.py

@@ -1,673 +1,649 b''
 # wireprotov1peer.py - Client-side functionality for wire protocol version 1.
 #
 # Copyright 2005-2010 Olivia Mackall <[e-mail]>
 #
 # 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 sys
 import weakref
 
 from .i18n import _
 from .node import bin
 from .pycompat import (
     getattr,
     setattr,
 )
 from . import (
     bundle2,
     changegroup as changegroupmod,
     encoding,
     error,
     pushkey as pushkeymod,
     pycompat,
     util,
     wireprototypes,
 )
 from .interfaces import (
     repository,
     util as interfaceutil,
 )
 from .utils import hashutil
 
 urlreq = util.urlreq
 
 
-def batchable_new_style(f):
+def batchable(f):
     """annotation for batchable methods
 
     Such methods must implement a coroutine as follows:
 
     @batchable
     def sample(self, one, two=None):
         # Build list of encoded arguments suitable for your wire protocol:
         encoded_args = [('one', encode(one),), ('two', encode(two),)]
         # Return it, along with a function that will receive the result
         # from the batched request.
         return encoded_args, decode
 
     The decorator returns a function which wraps this coroutine as a plain
     method, but adds the original method as an attribute called "batchable",
     which is used by remotebatch to split the call into separate encoding and
     decoding phases.
     """
 
     def plain(*args, **opts):
         encoded_args_or_res, decode = f(*args, **opts)
         if not decode:
             return encoded_args_or_res  # a local result in this case
         self = args[0]
         cmd = pycompat.bytesurl(f.__name__)  # ensure cmd is ascii bytestr
         encoded_res = self._submitone(cmd, encoded_args_or_res)
         return decode(encoded_res)
 
     setattr(plain, 'batchable', f)
     setattr(plain, '__name__', f.__name__)
     return plain
 
 
-def batchable(f):
-    def upgraded(*args, **opts):
-        batchable = f(*args, **opts)
-        encoded_args_or_res, encoded_res_future = next(batchable)
-        if not encoded_res_future:
-            decode = None
-        else:
-
-            def decode(d):
-                encoded_res_future.set(d)
-                return next(batchable)
-
-        return encoded_args_or_res, decode
-
-    setattr(upgraded, '__name__', f.__name__)
-    return batchable_new_style(upgraded)
-
-
-class future(object):
-    '''placeholder for a value to be set later'''
-
-    def set(self, value):
-        if util.safehasattr(self, b'value'):
-            raise error.RepoError(b"future is already set")
-        self.value = value
-
-
 def encodebatchcmds(req):
     """Return a ``cmds`` argument value for the ``batch`` command."""
     escapearg = wireprototypes.escapebatcharg
 
     cmds = []
     for op, argsdict in req:
         # Old servers didn't properly unescape argument names. So prevent
         # the sending of argument names that may not be decoded properly by
         # servers.
         assert all(escapearg(k) == k for k in argsdict)
 
         args = b','.join(
             b'%s=%s' % (escapearg(k), escapearg(v))
             for k, v in pycompat.iteritems(argsdict)
         )
         cmds.append(b'%s %s' % (op, args))
 
     return b';'.join(cmds)
 
 
 class unsentfuture(pycompat.futures.Future):
     """A Future variation to represent an unsent command.
 
     Because we buffer commands and don't submit them immediately, calling
     ``result()`` on an unsent future could deadlock. Futures for buffered
     commands are represented by this type, which wraps ``result()`` to
     call ``sendcommands()``.
     """
 
     def result(self, timeout=None):
         if self.done():
             return pycompat.futures.Future.result(self, timeout)
 
         self._peerexecutor.sendcommands()
 
         # This looks like it will infinitely recurse. However,
         # sendcommands() should modify __class__. This call serves as a check
         # on that.
         return self.result(timeout)
 
 
 @interfaceutil.implementer(repository.ipeercommandexecutor)
 class peerexecutor(object):
     def __init__(self, peer):
         self._peer = peer
         self._sent = False
         self._closed = False
         self._calls = []
         self._futures = weakref.WeakSet()
         self._responseexecutor = None
         self._responsef = None
 
     def __enter__(self):
         return self
 
     def __exit__(self, exctype, excvalee, exctb):
         self.close()
 
     def callcommand(self, command, args):
         if self._sent:
             raise error.ProgrammingError(
                 b'callcommand() cannot be used after commands are sent'
             )
 
         if self._closed:
             raise error.ProgrammingError(
                 b'callcommand() cannot be used after close()'
             )
 
         # Commands are dispatched through methods on the peer.
         fn = getattr(self._peer, pycompat.sysstr(command), None)
 
         if not fn:
             raise error.ProgrammingError(
                 b'cannot call command %s: method of same name not available '
                 b'on peer' % command
             )
 
         # Commands are either batchable or they aren't. If a command
         # isn't batchable, we send it immediately because the executor
         # can no longer accept new commands after a non-batchable command.
         # If a command is batchable, we queue it for later. But we have
         # to account for the case of a non-batchable command arriving after
         # a batchable one and refuse to service it.
 
         def addcall():
             f = pycompat.futures.Future()
             self._futures.add(f)
             self._calls.append((command, args, fn, f))
             return f
 
         if getattr(fn, 'batchable', False):
             f = addcall()
 
             # But since we don't issue it immediately, we wrap its result()
             # to trigger sending so we avoid deadlocks.
             f.__class__ = unsentfuture
             f._peerexecutor = self
         else:
             if self._calls:
                 raise error.ProgrammingError(
                     b'%s is not batchable and cannot be called on a command '
                     b'executor along with other commands' % command
                 )
 
             f = addcall()
 
             # Non-batchable commands can never coexist with another command
             # in this executor. So send the command immediately.
             self.sendcommands()
 
         return f
 
     def sendcommands(self):
         if self._sent:
             return
 
         if not self._calls:
             return
 
         self._sent = True
 
         # Unhack any future types so caller seens a clean type and to break
         # cycle between us and futures.
         for f in self._futures:
             if isinstance(f, unsentfuture):
                 f.__class__ = pycompat.futures.Future
                 f._peerexecutor = None
 
         calls = self._calls
         # Mainly to destroy references to futures.
         self._calls = None
 
         # Simple case of a single command. We call it synchronously.
         if len(calls) == 1:
             command, args, fn, f = calls[0]
 
             # Future was cancelled. Ignore it.
             if not f.set_running_or_notify_cancel():
                 return
 
             try:
                 result = fn(**pycompat.strkwargs(args))
             except Exception:
                 pycompat.future_set_exception_info(f, sys.exc_info()[1:])
             else:
                 f.set_result(result)
 
             return
 
         # Batch commands are a bit harder. First, we have to deal with the
         # @batchable coroutine. That's a bit annoying. Furthermore, we also
         # need to preserve streaming. i.e. it should be possible for the
         # futures to resolve as data is coming in off the wire without having
         # to wait for the final byte of the final response. We do this by
         # spinning up a thread to read the responses.
 
         requests = []
         states = []
 
         for command, args, fn, f in calls:
             # Future was cancelled. Ignore it.
             if not f.set_running_or_notify_cancel():
                 continue
 
             try:
                 encoded_args_or_res, decode = fn.batchable(
                     fn.__self__, **pycompat.strkwargs(args)
                 )
             except Exception:
                 pycompat.future_set_exception_info(f, sys.exc_info()[1:])
                 return
 
             if not decode:
                 f.set_result(encoded_args_or_res)
             else:
                 requests.append((command, encoded_args_or_res))
                 states.append((command, f, batchable, decode))
 
         if not requests:
             return
 
         # This will emit responses in order they were executed.
         wireresults = self._peer._submitbatch(requests)
 
         # The use of a thread pool executor here is a bit weird for something
         # that only spins up a single thread. However, thread management is
         # hard and it is easy to encounter race conditions, deadlocks, etc.
         # concurrent.futures already solves these problems and its thread pool
         # executor has minimal overhead. So we use it.
         self._responseexecutor = pycompat.futures.ThreadPoolExecutor(1)
         self._responsef = self._responseexecutor.submit(
             self._readbatchresponse, states, wireresults
         )
 
     def close(self):
         self.sendcommands()
 
         if self._closed:
             return
 
         self._closed = True
 
         if not self._responsef:
             return
 
         # We need to wait on our in-flight response and then shut down the
         # executor once we have a result.
         try:
             self._responsef.result()
         finally:
             self._responseexecutor.shutdown(wait=True)
             self._responsef = None
             self._responseexecutor = None
 
             # If any of our futures are still in progress, mark them as
             # errored. Otherwise a result() could wait indefinitely.
             for f in self._futures:
                 if not f.done():
                     f.set_exception(
                         error.ResponseError(
                             _(b'unfulfilled batch command response'), None
                         )
                     )
 
             self._futures = None
 
     def _readbatchresponse(self, states, wireresults):
         # Executes in a thread to read data off the wire.
 
         for command, f, batchable, decode in states:
             # Grab raw result off the wire and teach the internal future
             # about it.
             try:
                 remoteresult = next(wireresults)
             except StopIteration:
                 # This can happen in particular because next(batchable)
                 # in the previous iteration can call peer._abort, which
                 # may close the peer.
                 f.set_exception(
                     error.ResponseError(
                         _(b'unfulfilled batch command response'), None
                     )
                 )
             else:
                 try:
                     result = decode(remoteresult)
                 except Exception:
                     pycompat.future_set_exception_info(f, sys.exc_info()[1:])
                 else:
                     f.set_result(result)
 
 
 @interfaceutil.implementer(
     repository.ipeercommands, repository.ipeerlegacycommands
 )
 class wirepeer(repository.peer):
     """Client-side interface for communicating with a peer repository.
 
     Methods commonly call wire protocol commands of the same name.
 
     See also httppeer.py and sshpeer.py for protocol-specific
     implementations of this interface.
     """
 
     def commandexecutor(self):
         return peerexecutor(self)
 
     # Begin of ipeercommands interface.
 
     def clonebundles(self):
         self.requirecap(b'clonebundles', _(b'clone bundles'))
         return self._call(b'clonebundles')
 
     @batchable
     def lookup(self, key):
         self.requirecap(b'lookup', _(b'look up remote revision'))
-        f = future()
-        yield {b'key': encoding.fromlocal(key)}, f
-        d = f.value
+
+        def decode(d):
         success, data = d[:-1].split(b" ", 1)
         if int(success):
-            yield bin(data)
+                return bin(data)
         else:
             self._abort(error.RepoError(data))
 
+        return {b'key': encoding.fromlocal(key)}, decode
+
     @batchable
     def heads(self):
-        f = future()
-        yield {}, f
-        d = f.value
+        def decode(d):
         try:
-            yield wireprototypes.decodelist(d[:-1])
+                return wireprototypes.decodelist(d[:-1])
         except ValueError:
             self._abort(error.ResponseError(_(b"unexpected response:"), d))
 
+        return {}, decode
+
     @batchable
     def known(self, nodes):
-        f = future()
-        yield {b'nodes': wireprototypes.encodelist(nodes)}, f
-        d = f.value
+        def decode(d):
         try:
-            yield [bool(int(b)) for b in pycompat.iterbytestr(d)]
+                return [bool(int(b)) for b in pycompat.iterbytestr(d)]
         except ValueError:
             self._abort(error.ResponseError(_(b"unexpected response:"), d))
 
+        return {b'nodes': wireprototypes.encodelist(nodes)}, decode
+
     @batchable
     def branchmap(self):
-        f = future()
-        yield {}, f
-        d = f.value
+        def decode(d):
         try:
             branchmap = {}
             for branchpart in d.splitlines():
                 branchname, branchheads = branchpart.split(b' ', 1)
                 branchname = encoding.tolocal(urlreq.unquote(branchname))
                 branchheads = wireprototypes.decodelist(branchheads)
                 branchmap[branchname] = branchheads
-            yield branchmap
+                return branchmap
         except TypeError:
             self._abort(error.ResponseError(_(b"unexpected response:"), d))
 
+        return {}, decode
+
     @batchable
     def listkeys(self, namespace):
         if not self.capable(b'pushkey'):
-            yield {}, None
-        f = future()
+            return {}, None
         self.ui.debug(b'preparing listkeys for "%s"\n' % namespace)
-        yield {b'namespace': encoding.fromlocal(namespace)}, f
-        d = f.value
+
+        def decode(d):
         self.ui.debug(
             b'received listkey for "%s": %i bytes\n' % (namespace, len(d))
         )
-        yield pushkeymod.decodekeys(d)
+            return pushkeymod.decodekeys(d)
+
+        return {b'namespace': encoding.fromlocal(namespace)}, decode
 
     @batchable
     def pushkey(self, namespace, key, old, new):
         if not self.capable(b'pushkey'):
-            yield False, None
-        f = future()
+            return False, None
         self.ui.debug(b'preparing pushkey for "%s:%s"\n' % (namespace, key))
-        yield {
-            b'namespace': encoding.fromlocal(namespace),
-            b'key': encoding.fromlocal(key),
-            b'old': encoding.fromlocal(old),
-            b'new': encoding.fromlocal(new),
-        }, f
-        d = f.value
+
+        def decode(d):
         d, output = d.split(b'\n', 1)
         try:
             d = bool(int(d))
         except ValueError:
             raise error.ResponseError(
                 _(b'push failed (unexpected response):'), d
             )
         for l in output.splitlines(True):
             self.ui.status(_(b'remote: '), l)
-        yield d
+            return d
+
+        return {
+            b'namespace': encoding.fromlocal(namespace),
+            b'key': encoding.fromlocal(key),
+            b'old': encoding.fromlocal(old),
+            b'new': encoding.fromlocal(new),
+        }, decode
 
     def stream_out(self):
         return self._callstream(b'stream_out')
 
     def getbundle(self, source, **kwargs):
         kwargs = pycompat.byteskwargs(kwargs)
         self.requirecap(b'getbundle', _(b'look up remote changes'))
         opts = {}
         bundlecaps = kwargs.get(b'bundlecaps') or set()
         for key, value in pycompat.iteritems(kwargs):
             if value is None:
                 continue
             keytype = wireprototypes.GETBUNDLE_ARGUMENTS.get(key)
             if keytype is None:
                 raise error.ProgrammingError(
                     b'Unexpectedly None keytype for key %s' % key
                 )
             elif keytype == b'nodes':
                 value = wireprototypes.encodelist(value)
             elif keytype == b'csv':
                 value = b','.join(value)
             elif keytype == b'scsv':
                 value = b','.join(sorted(value))
             elif keytype == b'boolean':
                 value = b'%i' % bool(value)
             elif keytype != b'plain':
                 raise KeyError(b'unknown getbundle option type %s' % keytype)
             opts[key] = value
         f = self._callcompressable(b"getbundle", **pycompat.strkwargs(opts))
         if any((cap.startswith(b'HG2') for cap in bundlecaps)):
             return bundle2.getunbundler(self.ui, f)
         else:
             return changegroupmod.cg1unpacker(f, b'UN')
 
     def unbundle(self, bundle, heads, url):
         """Send cg (a readable file-like object representing the
         changegroup to push, typically a chunkbuffer object) to the
         remote server as a bundle.
 
         When pushing a bundle10 stream, return an integer indicating the
         result of the push (see changegroup.apply()).
 
         When pushing a bundle20 stream, return a bundle20 stream.
 
         `url` is the url the client thinks it's pushing to, which is
         visible to hooks.
         """
 
         if heads != [b'force'] and self.capable(b'unbundlehash'):
             heads = wireprototypes.encodelist(
                 [b'hashed', hashutil.sha1(b''.join(sorted(heads))).digest()]
             )
         else:
             heads = wireprototypes.encodelist(heads)
 
         if util.safehasattr(bundle, b'deltaheader'):
             # this a bundle10, do the old style call sequence
             ret, output = self._callpush(b"unbundle", bundle, heads=heads)
             if ret == b"":
                 raise error.ResponseError(_(b'push failed:'), output)
             try:
                 ret = int(ret)
             except ValueError:
                 raise error.ResponseError(
                     _(b'push failed (unexpected response):'), ret
                 )
 
             for l in output.splitlines(True):
                 self.ui.status(_(b'remote: '), l)
         else:
             # bundle2 push. Send a stream, fetch a stream.
             stream = self._calltwowaystream(b'unbundle', bundle, heads=heads)
             ret = bundle2.getunbundler(self.ui, stream)
         return ret
 
     # End of ipeercommands interface.
 
     # Begin of ipeerlegacycommands interface.
 
     def branches(self, nodes):
         n = wireprototypes.encodelist(nodes)
         d = self._call(b"branches", nodes=n)
         try:
             br = [tuple(wireprototypes.decodelist(b)) for b in d.splitlines()]
             return br
         except ValueError:
             self._abort(error.ResponseError(_(b"unexpected response:"), d))
 
     def between(self, pairs):
         batch = 8  # avoid giant requests
         r = []
         for i in pycompat.xrange(0, len(pairs), batch):
             n = b" ".join(
                 [
                     wireprototypes.encodelist(p, b'-')
                     for p in pairs[i : i + batch]
                 ]
             )
             d = self._call(b"between", pairs=n)
             try:
                 r.extend(
                     l and wireprototypes.decodelist(l) or []
                     for l in d.splitlines()
                 )
             except ValueError:
                 self._abort(error.ResponseError(_(b"unexpected response:"), d))
         return r
 
     def changegroup(self, nodes, source):
         n = wireprototypes.encodelist(nodes)
         f = self._callcompressable(b"changegroup", roots=n)
         return changegroupmod.cg1unpacker(f, b'UN')
 
     def changegroupsubset(self, bases, heads, source):
         self.requirecap(b'changegroupsubset', _(b'look up remote changes'))
         bases = wireprototypes.encodelist(bases)
         heads = wireprototypes.encodelist(heads)
         f = self._callcompressable(
             b"changegroupsubset", bases=bases, heads=heads
         )
         return changegroupmod.cg1unpacker(f, b'UN')
 
     # End of ipeerlegacycommands interface.
 
     def _submitbatch(self, req):
         """run batch request <req> on the server
 
         Returns an iterator of the raw responses from the server.
         """
         ui = self.ui
         if ui.debugflag and ui.configbool(b'devel', b'debug.peer-request'):
             ui.debug(b'devel-peer-request: batched-content\n')
             for op, args in req:
                 msg = b'devel-peer-request:    - %s (%d arguments)\n'
                 ui.debug(msg % (op, len(args)))
 
         unescapearg = wireprototypes.unescapebatcharg
 
         rsp = self._callstream(b"batch", cmds=encodebatchcmds(req))
         chunk = rsp.read(1024)
         work = [chunk]
         while chunk:
             while b';' not in chunk and chunk:
                 chunk = rsp.read(1024)
                 work.append(chunk)
             merged = b''.join(work)
             while b';' in merged:
                 one, merged = merged.split(b';', 1)
                 yield unescapearg(one)
             chunk = rsp.read(1024)
             work = [merged, chunk]
         yield unescapearg(b''.join(work))
 
     def _submitone(self, op, args):
         return self._call(op, **pycompat.strkwargs(args))
 
     def debugwireargs(self, one, two, three=None, four=None, five=None):
         # don't pass optional arguments left at their default value
         opts = {}
         if three is not None:
             opts['three'] = three
         if four is not None:
             opts['four'] = four
         return self._call(b'debugwireargs', one=one, two=two, **opts)
 
     def _call(self, cmd, **args):
         """execute <cmd> on the server
 
         The command is expected to return a simple string.
 
         returns the server reply as a string."""
         raise NotImplementedError()
 
     def _callstream(self, cmd, **args):
         """execute <cmd> on the server
 
         The command is expected to return a stream. Note that if the
         command doesn't return a stream, _callstream behaves
         differently for ssh and http peers.
 
         returns the server reply as a file like object.
         """
         raise NotImplementedError()
 
     def _callcompressable(self, cmd, **args):
         """execute <cmd> on the server
 
         The command is expected to return a stream.
 
         The stream may have been compressed in some implementations. This
         function takes care of the decompression. This is the only difference
         with _callstream.
 
         returns the server reply as a file like object.
         """
         raise NotImplementedError()
 
     def _callpush(self, cmd, fp, **args):
         """execute a <cmd> on server
 
         The command is expected to be related to a push. Push has a special
         return method.
 
         returns the server reply as a (ret, output) tuple. ret is either
         empty (error) or a stringified int.
         """
         raise NotImplementedError()
 
     def _calltwowaystream(self, cmd, fp, **args):
         """execute <cmd> on server
 
         The command will send a stream to the server and get a stream in reply.
         """
         raise NotImplementedError()
 
     def _abort(self, exception):
         """clearly abort the wire protocol connection and raise the exception"""
         raise NotImplementedError()

tests/test-batching.py

@@ -1,258 +1,254 b''
 # test-batching.py - tests for transparent command batching
 #
 # Copyright 2011 Peter Arrenbrecht <[e-mail]>
 #
 # 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, print_function
 
 import contextlib
 
 from mercurial import (
     localrepo,
     pycompat,
     wireprotov1peer,
 )
 
 
 def bprint(*bs):
     print(*[pycompat.sysstr(b) for b in bs])
 
 
 # equivalent of repo.repository
 class thing(object):
     def hello(self):
         return b"Ready."
 
 
 # equivalent of localrepo.localrepository
 class localthing(thing):
     def foo(self, one, two=None):
         if one:
             return b"%s and %s" % (
                 one,
                 two,
             )
         return b"Nope"
 
     def bar(self, b, a):
         return b"%s und %s" % (
             b,
             a,
         )
 
     def greet(self, name=None):
         return b"Hello, %s" % name
 
     @contextlib.contextmanager
     def commandexecutor(self):
         e = localrepo.localcommandexecutor(self)
         try:
             yield e
         finally:
             e.close()
 
 
 # usage of "thing" interface
 def use(it):
 
     # Direct call to base method shared between client and server.
     bprint(it.hello())
 
     # Direct calls to proxied methods. They cause individual roundtrips.
     bprint(it.foo(b"Un", two=b"Deux"))
     bprint(it.bar(b"Eins", b"Zwei"))
 
     # Batched call to a couple of proxied methods.
 
     with it.commandexecutor() as e:
         ffoo = e.callcommand(b'foo', {b'one': b'One', b'two': b'Two'})
         fbar = e.callcommand(b'bar', {b'b': b'Eins', b'a': b'Zwei'})
         fbar2 = e.callcommand(b'bar', {b'b': b'Uno', b'a': b'Due'})
 
     bprint(ffoo.result())
     bprint(fbar.result())
     bprint(fbar2.result())
 
 
 # local usage
 mylocal = localthing()
 print()
 bprint(b"== Local")
 use(mylocal)
 
 # demo remoting; mimicks what wireproto and HTTP/SSH do
 
 # shared
 
 
 def escapearg(plain):
     return (
         plain.replace(b':', b'::')
         .replace(b',', b':,')
         .replace(b';', b':;')
         .replace(b'=', b':=')
     )
 
 
 def unescapearg(escaped):
     return (
         escaped.replace(b':=', b'=')
         .replace(b':;', b';')
         .replace(b':,', b',')
         .replace(b'::', b':')
     )
 
 
 # server side
 
 # equivalent of wireproto's global functions
 class server(object):
     def __init__(self, local):
         self.local = local
 
     def _call(self, name, args):
         args = dict(arg.split(b'=', 1) for arg in args)
         return getattr(self, name)(**args)
 
     def perform(self, req):
         bprint(b"REQ:", req)
         name, args = req.split(b'?', 1)
         args = args.split(b'&')
         vals = dict(arg.split(b'=', 1) for arg in args)
         res = getattr(self, pycompat.sysstr(name))(**pycompat.strkwargs(vals))
         bprint(b"  ->", res)
         return res
 
     def batch(self, cmds):
         res = []
         for pair in cmds.split(b';'):
             name, args = pair.split(b':', 1)
             vals = {}
             for a in args.split(b','):
                 if a:
                     n, v = a.split(b'=')
                     vals[n] = unescapearg(v)
             res.append(
                 escapearg(
                     getattr(self, pycompat.sysstr(name))(
                         **pycompat.strkwargs(vals)
                     )
                 )
             )
         return b';'.join(res)
 
     def foo(self, one, two):
         return mangle(self.local.foo(unmangle(one), unmangle(two)))
 
     def bar(self, b, a):
         return mangle(self.local.bar(unmangle(b), unmangle(a)))
 
     def greet(self, name):
         return mangle(self.local.greet(unmangle(name)))
 
 
 myserver = server(mylocal)
 
 # local side
 
 # equivalent of wireproto.encode/decodelist, that is, type-specific marshalling
 # here we just transform the strings a bit to check we're properly en-/decoding
 def mangle(s):
     return b''.join(pycompat.bytechr(ord(c) + 1) for c in pycompat.bytestr(s))
 
 
 def unmangle(s):
     return b''.join(pycompat.bytechr(ord(c) - 1) for c in pycompat.bytestr(s))
 
 
 # equivalent of wireproto.wirerepository and something like http's wire format
 class remotething(thing):
     def __init__(self, server):
         self.server = server
 
     def _submitone(self, name, args):
         req = name + b'?' + b'&'.join([b'%s=%s' % (n, v) for n, v in args])
         return self.server.perform(req)
 
     def _submitbatch(self, cmds):
         req = []
         for name, args in cmds:
             args = b','.join(n + b'=' + escapearg(v) for n, v in args)
             req.append(name + b':' + args)
         req = b';'.join(req)
         res = self._submitone(
             b'batch',
             [
                 (
                     b'cmds',
                     req,
                 )
             ],
         )
         for r in res.split(b';'):
             yield r
 
     @contextlib.contextmanager
     def commandexecutor(self):
         e = wireprotov1peer.peerexecutor(self)
         try:
             yield e
         finally:
             e.close()
 
     @wireprotov1peer.batchable
     def foo(self, one, two=None):
         encoded_args = [
             (
                 b'one',
                 mangle(one),
             ),
             (
                 b'two',
                 mangle(two),
             ),
         ]
-        encoded_res_future = wireprotov1peer.future()
-        yield encoded_args, encoded_res_future
-        yield unmangle(encoded_res_future.value)
+        return encoded_args, unmangle
 
     @wireprotov1peer.batchable
     def bar(self, b, a):
-        encresref = wireprotov1peer.future()
-        yield [
+        return [
             (
                 b'b',
                 mangle(b),
             ),
             (
                 b'a',
                 mangle(a),
             ),
-        ], encresref
-        yield unmangle(encresref.value)
+        ], unmangle
 
     # greet is coded directly. It therefore does not support batching. If it
     # does appear in a batch, the batch is split around greet, and the call to
     # greet is done in its own roundtrip.
     def greet(self, name=None):
         return unmangle(
             self._submitone(
                 b'greet',
                 [
                     (
                         b'name',
                         mangle(name),
                     )
                 ],
             )
         )
 
 
 # demo remote usage
 
 myproxy = remotething(myserver)
 print()
 bprint(b"== Remote")
 use(myproxy)

tests/test-wireproto.py

@@ -1,126 +1,124 b''
 from __future__ import absolute_import, print_function
 
 import sys
 
 from mercurial import (
     error,
     pycompat,
     ui as uimod,
     util,
     wireprototypes,
     wireprotov1peer,
     wireprotov1server,
 )
 from mercurial.utils import stringutil
 
 stringio = util.stringio
 
 
 class proto(object):
     def __init__(self, args):
         self.args = args
         self.name = 'dummyproto'
 
     def getargs(self, spec):
         args = self.args
         args.setdefault(b'*', {})
         names = spec.split()
         return [args[n] for n in names]
 
     def checkperm(self, perm):
         pass
 
 
 wireprototypes.TRANSPORTS['dummyproto'] = {
     'transport': 'dummy',
     'version': 1,
 }
 
 
 class clientpeer(wireprotov1peer.wirepeer):
     def __init__(self, serverrepo, ui):
         self.serverrepo = serverrepo
         self.ui = ui
 
     def url(self):
         return b'test'
 
     def local(self):
         return None
 
     def peer(self):
         return self
 
     def canpush(self):
         return True
 
     def close(self):
         pass
 
     def capabilities(self):
         return [b'batch']
 
     def _call(self, cmd, **args):
         args = pycompat.byteskwargs(args)
         res = wireprotov1server.dispatch(self.serverrepo, proto(args), cmd)
         if isinstance(res, wireprototypes.bytesresponse):
             return res.data
         elif isinstance(res, bytes):
             return res
         else:
             raise error.Abort('dummy client does not support response type')
 
     def _callstream(self, cmd, **args):
         return stringio(self._call(cmd, **args))
 
     @wireprotov1peer.batchable
     def greet(self, name):
-        f = wireprotov1peer.future()
-        yield {b'name': mangle(name)}, f
-        yield unmangle(f.value)
+        return {b'name': mangle(name)}, unmangle
 
 
 class serverrepo(object):
     def __init__(self, ui):
         self.ui = ui
 
     def greet(self, name):
         return b"Hello, " + name
 
     def filtered(self, name):
         return self
 
 
 def mangle(s):
     return b''.join(pycompat.bytechr(ord(c) + 1) for c in pycompat.bytestr(s))
 
 
 def unmangle(s):
     return b''.join(pycompat.bytechr(ord(c) - 1) for c in pycompat.bytestr(s))
 
 
 def greet(repo, proto, name):
     return mangle(repo.greet(unmangle(name)))
 
 
 wireprotov1server.commands[b'greet'] = (greet, b'name')
 
 srv = serverrepo(uimod.ui())
 clt = clientpeer(srv, uimod.ui())
 
 
 def printb(data, end=b'\n'):
     out = getattr(sys.stdout, 'buffer', sys.stdout)
     out.write(data + end)
     out.flush()
 
 
 printb(clt.greet(b"Foobar"))
 
 with clt.commandexecutor() as e:
     fgreet1 = e.callcommand(b'greet', {b'name': b'Fo, =;:<o'})
     fgreet2 = e.callcommand(b'greet', {b'name': b'Bar'})
 
 printb(
     stringutil.pprint([f.result() for f in (fgreet1, fgreet2)], bprefix=True)
 )