mirror of
https://github.com/facebook/sapling.git
synced 2024-10-09 08:18:15 +03:00
32ac5d0647
Summary: This makes hg pull use the connectionpool. This means prefetches can reuse the existing ssh connection when appropriate. This both speeds up prefetches, and also means they will speak to the same server that served the pull. Reviewed By: ryanmce Differential Revision: D7481107 fbshipit-source-id: f9a3670527cb7e8956029c86d50d8e030dd3cc01
109 lines
3.1 KiB
Python
109 lines
3.1 KiB
Python
# 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 . import (
|
|
extensions,
|
|
sshpeer,
|
|
util,
|
|
)
|
|
|
|
class connectionpool(object):
|
|
def __init__(self, repo):
|
|
self._repo = repo
|
|
self._pool = dict()
|
|
|
|
def get(self, path, opts=None):
|
|
# Prevent circular dependency
|
|
from . import hg
|
|
|
|
if opts is None:
|
|
opts = {}
|
|
if opts.get('ssh') or opts.get('remotecmd'):
|
|
self._repo.ui.debug("not using connection pool due to ssh or "
|
|
"remotecmd option being set\n")
|
|
peer = hg.peer(self._repo.ui, opts, path)
|
|
return standaloneconnection(peer)
|
|
|
|
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, sshpeer.sshpeer) and
|
|
util.safehasattr(peer, '_subprocess')):
|
|
proc = peer._subprocess
|
|
if proc.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 standaloneconnection(object):
|
|
def __init__(self, peer):
|
|
self.peer = peer
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, type, value, traceback):
|
|
self.close()
|
|
|
|
def close(self):
|
|
if util.safehasattr(self.peer, 'cleanup'):
|
|
self.peer.cleanup()
|
|
|
|
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()
|