diff hgext/remotefilelog/connectionpool.py @ 40495:3a333a582d7b

remotefilelog: import pruned-down remotefilelog extension from hg-experimental

This is remotefilelog as of my recent patches for compatibility with
current tip of hg, minus support for old versions of Mercurial and
some FB-specific features like their treemanifest extension and
fetching linkrev data from a patched phabricator. The file extutil.py
moved from hgext3rd to remotefilelog.

This is not yet ready to be landed, consider it a preview for
now. Planned changes include:

 * replace lz4 with zstd
 * rename some capabilities, requirements and wireproto commands to mark
   them as experimental
 * consolidate bits of shallowutil with related functions (eg readfile)

I'm certainly open to other (small) changes, but my rough mission is
to land this largely as-is so we can use it as a model of the
functionality we need going forward for lazy-fetching of file contents
from a server.

# no-check-commit because of a few foo_bar functions

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

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/hgext/remotefilelog/connectionpool.py	Thu Sep 27 13:03:19 2018 -0400
@@ -0,0 +1,84 @@
+# connectionpool.py - class for pooling peer connections for reuse
+#
+# Copyright 2017 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
+
+from mercurial import (
+    extensions,
+    hg,
+    sshpeer,
+    util,
+)
+
+_sshv1peer = sshpeer.sshv1peer
+
+class connectionpool(object):
+    def __init__(self, repo):
+        self._repo = repo
+        self._pool = dict()
+
+    def get(self, path):
+        pathpool = self._pool.get(path)
+        if pathpool is None:
+            pathpool = list()
+            self._pool[path] = pathpool
+
+        conn = None
+        if len(pathpool) > 0:
+            try:
+                conn = pathpool.pop()
+                peer = conn.peer
+                # If the connection has died, drop it
+                if isinstance(peer, _sshv1peer):
+                    if peer._subprocess.poll() is not None:
+                        conn = None
+            except IndexError:
+                pass
+
+        if conn is None:
+            def _cleanup(orig):
+                # close pipee first so peer.cleanup reading it won't deadlock,
+                # if there are other processes with pipeo open (i.e. us).
+                peer = orig.im_self
+                if util.safehasattr(peer, 'pipee'):
+                    peer.pipee.close()
+                return orig()
+
+            peer = hg.peer(self._repo.ui, {}, path)
+            if util.safehasattr(peer, 'cleanup'):
+                extensions.wrapfunction(peer, 'cleanup', _cleanup)
+
+            conn = connection(pathpool, peer)
+
+        return conn
+
+    def close(self):
+        for pathpool in self._pool.itervalues():
+            for conn in pathpool:
+                conn.close()
+            del pathpool[:]
+
+class connection(object):
+    def __init__(self, pool, peer):
+        self._pool = pool
+        self.peer = peer
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        # Only add the connection back to the pool if there was no exception,
+        # since an exception could mean the connection is not in a reusable
+        # state.
+        if type is None:
+            self._pool.append(self)
+        else:
+            self.close()
+
+    def close(self):
+        if util.safehasattr(self.peer, 'cleanup'):
+            self.peer.cleanup()