treemanifest: use a connectionpool

remotefilelog and fastannotate already use a connection pool to share and reuse
connections. Treemanifest often does ondemand downloading of trees, such as
during hg log -p, and would greatly benefit from reusing connections as well.

This patch makes the connectionpool and attribute of the repo object, instead of
the fileserverclient object, which allows treemanifest to make use of it easily.

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

hgext3rd/fastannotate/protocol.py

@@ -138,7 +138,7 @@ def annotatepeer(repo):
 
     if sharepeer and fileservice:
         ui.debug('fastannotate: using remotefilelog connection pool\n')
-        conn = fileservice.connpool.get(repo.fallbackpath)
+        conn = repo.connectionpool.get(repo.fallbackpath)
         peer = conn.peer
         stolen = True
     else:

remotefilelog/fileserverclient.py

@@ -22,7 +22,6 @@ from mercurial import (
 )
 
 from . import (
-    connectionpool,
     constants,
     shallowutil,
     wirepack,
@@ -258,7 +257,6 @@ class fileserverclient(object):
         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(
@@ -267,7 +265,6 @@ class fileserverclient(object):
         self.debugoutput = ui.configbool("remotefilelog", "debug")
 
         self.remotecache = cacheconnection()
-        self.connpool = connectionpool.connectionpool(repo)
 
     def setstore(self, datastore, historystore, writedata, writehistory):
         self.datastore = datastore
@@ -276,7 +273,7 @@ class fileserverclient(object):
         self.writehistory = writehistory
 
     def _connect(self):
-        return self.connpool.get(self.repo.fallbackpath)
+        return self.repo.connectionpool.get(self.repo.fallbackpath)
 
     def request(self, fileids):
         """Takes a list of filename/node pairs and fetches them from the
@@ -526,8 +523,6 @@ class fileserverclient(object):
         if self.remotecache.connected:
             self.remotecache.close()
 
-        self.connpool.close()
-
     def prefetch(self, fileids, force=False, fetchdata=True,
                  fetchhistory=False):
         """downloads the given file versions to the cache

remotefilelog/shallowrepo.py

@@ -9,7 +9,12 @@ from hgext3rd.extutil import runshellcommand
 from mercurial.i18n import _
 from mercurial.node import hex, nullid, nullrev
 from mercurial import error, localrepo, util, match, scmutil
-from . import remotefilelog, remotefilectx, fileserverclient
+from . import (
+    connectionpool,
+    fileserverclient,
+    remotefilelog,
+    remotefilectx,
+)
 import constants, shallowutil
 from contentstore import remotefilelogcontentstore, unioncontentstore
 from contentstore import remotecontentstore
@@ -271,6 +276,10 @@ def wraprepo(repo):
                 results = [(path, hex(fnode)) for (path, fnode) in files]
                 repo.fileservice.prefetch(results)
 
+        def close(self):
+            super(shallowrepository, self).close()
+            self.connectionpool.close()
+
     repo.__class__ = shallowrepository
 
     repo.shallowmatch = match.always(repo.root, '')
@@ -281,6 +290,9 @@ def wraprepo(repo):
                                              None)
     repo.excludepattern = repo.ui.configlist("remotefilelog", "excludepattern",
                                              None)
+    if not util.safehasattr(repo, 'connectionpool'):
+        repo.connectionpool = connectionpool.connectionpool(repo)
+
     if repo.includepattern or repo.excludepattern:
         repo.shallowmatch = match.match(repo.root, '', None,
             repo.includepattern, repo.excludepattern)

treemanifest/__init__.py

@@ -73,7 +73,6 @@ from mercurial import (
     error,
     exchange,
     extensions,
-    hg,
     localrepo,
     manifest,
     mdiff,
@@ -113,7 +112,12 @@ from remotefilelog.historypack import (
     historypackstore,
     mutablehistorypack,
 )
-from remotefilelog import shallowrepo, shallowutil, wirepack
+from remotefilelog import (
+    connectionpool,
+    shallowrepo,
+    shallowutil,
+    wirepack,
+)
 from remotefilelog.repack import (
     _computeincrementaldatapack,
     _computeincrementalhistorypack,
@@ -256,6 +260,9 @@ def clientreposetup(repo):
         except KeyError:
             raise error.Abort(_("cannot use treemanifest without fastmanifest"))
 
+    if not util.safehasattr(repo, 'connectionpool'):
+        repo.connectionpool = connectionpool.connectionpool(repo)
+
 def setuptreestores(repo, mfl):
     if repo.ui.configbool("treemanifest", "server"):
         packpath = repo.vfs.join('cache/packs/%s' % PACK_CATEGORY)
@@ -1070,7 +1077,12 @@ def _prefetchtrees(repo, rootdir, mfnodes, basemfnodes, directories):
         fallbackpath = repo.ui.config('paths', 'default')
 
     start = time.time()
-    remote = hg.peer(repo.ui, {}, fallbackpath)
+    with repo.connectionpool.get(fallbackpath) as conn:
+        remote = conn.peer
+        _gettrees(repo, remote, rootdir, mfnodes, basemfnodes, directories,
+                  start)
+
+def _gettrees(repo, remote, rootdir, mfnodes, basemfnodes, directories, start):
     if 'gettreepack' not in shallowutil.peercapabilities(remote):
         raise error.Abort(_("missing gettreepack capability on remote"))
     remote.ui.pushbuffer()
@@ -1100,14 +1112,16 @@ def _prefetchtrees(repo, rootdir, mfnodes, basemfnodes, directories):
         nodestr = '\n'.join(hexnodes[:10])
         if len(hexnodes) > 10:
             nodestr += '\n...'
+        # Give stderr some time to reach the client, so we can read it into the
+        # currently pushed ui buffer, instead of it randomly showing up in a
+        # future ui read.
+        time.sleep(0.1)
         raise error.Abort(_('unable to download the following trees from the '
                             'server:\n%s') % nodestr, hint=exc.hint)
     except error.BundleValueError as exc:
         raise error.Abort(_('missing support for %s') % exc)
     finally:
-        # Manually destruct the peer, so we can collect any error output
-        remote._cleanup()
-
+        remote._readerr()
         output = remote.ui.popbuffer()
         if output:
             repo.ui.debug(output)