From b9a3c7c038834de7047bec3f1db3cf8050fd88d1 Mon Sep 17 00:00:00 2001
From: Alexander Schepanovski <suor.web@gmail.com>
Date: Mon, 15 Jul 2019 03:30:24 +0700
Subject: [PATCH] remote: use ssh connections pool (#2264)

* remote: use ssh connections pool

* dvc: never use paramiko DEBUG mode

It's ridiculously verbose, makes logs unreadable.
---
 dvc/main.py                  |  1 -
 dvc/remote/ssh/__init__.py   | 12 +++------
 dvc/remote/ssh/connection.py | 17 ++++++------
 dvc/remote/ssh/pool.py       | 50 ++++++++++++++++++++++++++++++++++++
 4 files changed, 62 insertions(+), 18 deletions(-)
 create mode 100644 dvc/remote/ssh/pool.py

diff --git a/dvc/main.py b/dvc/main.py
index cb162fb7d..f00dc0ced 100644
--- a/dvc/main.py
+++ b/dvc/main.py
@@ -34,7 +34,6 @@ def main(argv=None):
 
         elif args.verbose:
             logger.setLevel(logging.DEBUG)
-            logging.getLogger("paramiko").setLevel(logging.DEBUG)
 
         cmd = args.func(args)
         ret = cmd.run_cmd()
diff --git a/dvc/remote/ssh/__init__.py b/dvc/remote/ssh/__init__.py
index 07fdb910f..3e34e9ce9 100644
--- a/dvc/remote/ssh/__init__.py
+++ b/dvc/remote/ssh/__init__.py
@@ -13,12 +13,12 @@ except ImportError:
     paramiko = None
 
 import dvc.prompt as prompt
-from dvc.remote.ssh.connection import SSHConnection
 from dvc.config import Config
 from dvc.utils import to_chunks
 from dvc.utils.compat import urlparse, StringIO
 from dvc.remote.base import RemoteBASE
 from dvc.scheme import Schemes
+from .pool import ssh_connection
 
 
 logger = logging.getLogger(__name__)
@@ -108,13 +108,7 @@ class RemoteSSH(RemoteBASE):
     def ssh(self, path_info):
         host, user, port = path_info.host, path_info.user, path_info.port
 
-        logger.debug(
-            "Establishing ssh connection with '{host}' "
-            "through port '{port}' as user '{user}'".format(
-                host=host, user=user, port=port
-            )
-        )
-
+        # NOTE: we use the same password regardless of the server :(
         if self.ask_password and not self.password:
             self.password = prompt.password(
                 "Enter a private key passphrase or a password for "
@@ -123,7 +117,7 @@ class RemoteSSH(RemoteBASE):
                 )
             )
 
-        return SSHConnection(
+        return ssh_connection(
             host,
             username=user,
             port=port,
diff --git a/dvc/remote/ssh/connection.py b/dvc/remote/ssh/connection.py
index 76d14a90c..3aa6eb598 100644
--- a/dvc/remote/ssh/connection.py
+++ b/dvc/remote/ssh/connection.py
@@ -45,23 +45,24 @@ def create_cb(name):
 
 
 class SSHConnection:
-    def __init__(self, *args, **kwargs):
+    def __init__(self, host, *args, **kwargs):
+        logger.debug(
+            "Establishing ssh connection with '{host}' "
+            "through port '{port}' as user '{username}'".format(
+                host=host, **kwargs
+            )
+        )
         self.timeout = kwargs.get("timeout", 1800)
 
         self._ssh = paramiko.SSHClient()
         self._ssh.load_system_host_keys()
         self._ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
 
-        self._ssh.connect(*args, **kwargs)
+        self._ssh.connect(host, *args, **kwargs)
+        self._ssh.get_transport().set_keepalive(10)
         self._sftp = None
         self._sftp_alive = False
 
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exception_type, exception_value, traceback):
-        self.close()
-
     def _sftp_connect(self):
         if not self._sftp or not self._sftp_alive:
             self._sftp = self._ssh.open_sftp()
diff --git a/dvc/remote/ssh/pool.py b/dvc/remote/ssh/pool.py
new file mode 100644
index 000000000..abba1b53c
--- /dev/null
+++ b/dvc/remote/ssh/pool.py
@@ -0,0 +1,50 @@
+from collections import deque
+from contextlib import contextmanager
+from funcy import memoize
+
+from .connection import SSHConnection
+
+
+@contextmanager
+def ssh_connection(*conn_args, **conn_kwargs):
+    pool = get_ssh_pool(*conn_args, **conn_kwargs)
+    conn = pool.get_connection()
+    try:
+        yield conn
+    except BaseException:
+        conn.close()
+    else:
+        pool.release(conn)
+
+
+@memoize
+def get_ssh_pool(*conn_args, **conn_kwargs):
+    return SSHPool(conn_args, conn_kwargs)
+
+
+class SSHPool(object):
+    def __init__(self, conn_args, conn_kwargs):
+        self._conn_args = conn_args
+        self._conn_kwargs = conn_kwargs
+        self._conns = deque()
+        self._closed = False
+
+    def __del__(self):
+        self.close()
+
+    def close(self):
+        while self._conns:
+            self._conns.pop().close()
+        self._closed = True
+
+    def get_connection(self):
+        try:
+            return self._conns.popleft()
+        except IndexError:
+            return SSHConnection(*self._conn_args, **self._conn_kwargs)
+
+    def release(self, conn):
+        if self._closed:
+            conn.close()
+        else:
+            self._conns.append(conn)