From af26748add1ce23263cdd38d91daeb75c4a99a2f Mon Sep 17 00:00:00 2001
From: "Shawn O. Pearce" <sop@google.com>
Date: Fri, 10 Apr 2009 18:53:46 -0700
Subject: [PATCH] Automatically use SSH control master support during sync

By creating a background ssh "control master" process which lives
for the duration of our sync cycle we can easily cut the time for
a no-op sync of 132 projects from 60s to 18s.

Bug: REPO-11
Signed-off-by: Shawn O. Pearce <sop@google.com>
---
 git_command.py | 10 ++++++++
 git_config.py  | 70 ++++++++++++++++++++++++++++++++++++++++++++++++++
 git_ssh        |  2 ++
 main.py        |  6 ++++-
 project.py     | 10 +++++++-
 5 files changed, 96 insertions(+), 2 deletions(-)
 create mode 100755 git_ssh

diff --git a/git_command.py b/git_command.py
index a3bd9192e..122bbcded 100644
--- a/git_command.py
+++ b/git_command.py
@@ -30,6 +30,13 @@ try:
 except KeyError:
   TRACE = False
 
+_ssh_proxy_path = None
+def _ssh_proxy():
+  global _ssh_proxy_path
+  if _ssh_proxy_path is None:
+    _ssh_proxy_path = os.path.join(os.path.dirname(__file__),'git_ssh')
+  return _ssh_proxy_path
+
 
 class _GitCall(object):
   def version(self):
@@ -56,6 +63,7 @@ class GitCommand(object):
                capture_stdout = False,
                capture_stderr = False,
                disable_editor = False,
+               ssh_proxy = False,
                cwd = None,
                gitdir = None):
     env = dict(os.environ)
@@ -72,6 +80,8 @@ class GitCommand(object):
 
     if disable_editor:
       env['GIT_EDITOR'] = ':'
+    if ssh_proxy:
+      env['GIT_SSH'] = _ssh_proxy()
 
     if project:
       if not cwd:
diff --git a/git_config.py b/git_config.py
index 1d45d92f8..1384b1144 100644
--- a/git_config.py
+++ b/git_config.py
@@ -15,7 +15,10 @@
 
 import os
 import re
+import subprocess
 import sys
+import time
+from signal import SIGTERM
 from urllib2 import urlopen, HTTPError
 from error import GitError, UploadError
 from git_command import GitCommand
@@ -253,6 +256,70 @@ class RefSpec(object):
     return s
 
 
+_ssh_cache = {}
+_ssh_master = True
+
+def _open_ssh(host, port=None):
+  global _ssh_master
+
+  if port is None:
+    port = 22
+
+  key = '%s:%s' % (host, port)
+  if key in _ssh_cache:
+    return True
+
+  if not _ssh_master \
+  or 'GIT_SSH' in os.environ \
+  or sys.platform == 'win32':
+    # failed earlier, or cygwin ssh can't do this
+    #
+    return False
+
+  command = ['ssh',
+             '-o','ControlPath ~/.ssh/master-%r@%h:%p',
+             '-p',port,
+             '-M',
+             '-N',
+             host]
+  try:
+    p = subprocess.Popen(command)
+  except Exception, e:
+    _ssh_master = False
+    print >>sys.stderr, \
+      '\nwarn: cannot enable ssh control master for %s:%s' \
+      % (host,port)
+    return False
+
+  _ssh_cache[key] = p
+  time.sleep(1)
+  return True
+
+def close_ssh():
+  for key,p in _ssh_cache.iteritems():
+    os.kill(p.pid, SIGTERM)
+  _ssh_cache.clear()
+
+URI_SCP = re.compile(r'^([^@:]*@?[^:/]{1,}):')
+URI_ALL = re.compile(r'^([a-z][a-z+]*)://([^@/]*@?[^/])/')
+
+def _preconnect(url):
+  m = URI_ALL.match(url)
+  if m:
+    scheme = m.group(1)
+    host = m.group(2)
+    if ':' in host:
+      host, port = host.split(':')
+    if scheme in ('ssh', 'git+ssh', 'ssh+git'):
+      return _open_ssh(host, port)
+    return False
+
+  m = URI_SCP.match(url)
+  if m:
+    host = m.group(1)
+    return _open_ssh(host)
+
+
 class Remote(object):
   """Configuration options related to a remote.
   """
@@ -266,6 +333,9 @@ class Remote(object):
                      self._Get('fetch', all=True))
     self._review_protocol = None
 
+  def PreConnectFetch(self):
+    return _preconnect(self.url)
+
   @property
   def ReviewProtocol(self):
     if self._review_protocol is None:
diff --git a/git_ssh b/git_ssh
new file mode 100755
index 000000000..5f5c55fd4
--- /dev/null
+++ b/git_ssh
@@ -0,0 +1,2 @@
+#!/bin/sh
+exec ssh -o 'ControlPath ~/.ssh/master-%r@%h:%p' "$@"
diff --git a/main.py b/main.py
index df1818356..834db95a6 100755
--- a/main.py
+++ b/main.py
@@ -28,6 +28,7 @@ import re
 import sys
 
 import git_command
+from git_config import close_ssh
 from command import InteractiveCommand
 from command import MirrorSafeCommand
 from command import PagedCommand
@@ -211,7 +212,10 @@ def _Main(argv):
 
   repo = _Repo(opt.repodir)
   try:
-    repo._Run(argv)
+    try:
+      repo._Run(argv)
+    finally:
+      close_ssh()
   except KeyboardInterrupt:
     sys.exit(1)
   except RepoChangedException, rce:
diff --git a/project.py b/project.py
index bd35c65e6..b86e08e3e 100644
--- a/project.py
+++ b/project.py
@@ -868,11 +868,19 @@ class Project(object):
   def _RemoteFetch(self, name=None):
     if not name:
       name = self.remote.name
+
+    ssh_proxy = False
+    if self.GetRemote(name).PreConnectFetch():
+      ssh_proxy = True
+
     cmd = ['fetch']
     if not self.worktree:
       cmd.append('--update-head-ok')
     cmd.append(name)
-    return GitCommand(self, cmd, bare = True).Wait() == 0
+    return GitCommand(self,
+                      cmd,
+                      bare = True,
+                      ssh_proxy = ssh_proxy).Wait() == 0
 
   def _Checkout(self, rev, quiet=False):
     cmd = ['checkout']