From 47620b5a3233ca2644d4a02073801c03ac846bac Mon Sep 17 00:00:00 2001
From: velorums <velorums@gmail.com>
Date: Wed, 23 Mar 2022 17:40:46 +0100
Subject: [PATCH] Reuse ssh connection

---
 docker/transport/sshconn.py | 42 ++++++++++++++++---------------------
 1 file changed, 18 insertions(+), 24 deletions(-)

diff --git a/docker/transport/sshconn.py b/docker/transport/sshconn.py
index 8e6beb254..070b69377 100644
--- a/docker/transport/sshconn.py
+++ b/docker/transport/sshconn.py
@@ -18,7 +18,18 @@
 except ImportError:
     import urllib3
 
-RecentlyUsedContainer = urllib3._collections.RecentlyUsedContainer
+
+class DontCloseStreamWrapper:
+    """Make close() noop for the wrapped object."""
+
+    def __init__(self, obj):
+        self.obj = obj
+
+    def __getattr__(self, name):
+        def wrapper(*args, **kwargs):
+            if name != "close":
+                return getattr(self.obj, name)(*args, **kwargs)
+        return wrapper
 
 
 class SSHSocket(socket.socket):
@@ -90,7 +101,7 @@ def makefile(self, mode):
             self.connect()
         self.proc.stdout.channel = self
 
-        return self.proc.stdout
+        return DontCloseStreamWrapper(self.proc.stdout)
 
     def close(self):
         if not self.proc or self.proc.stdin.closed:
@@ -164,7 +175,7 @@ def _get_conn(self, timeout):
 class SSHHTTPAdapter(BaseHTTPAdapter):
 
     __attrs__ = requests.adapters.HTTPAdapter.__attrs__ + [
-        'pools', 'timeout', 'ssh_client', 'ssh_params', 'max_pool_size'
+        'pool', 'timeout', 'ssh_client', 'ssh_params', 'max_pool_size'
     ]
 
     def __init__(self, base_url, timeout=60,
@@ -181,10 +192,8 @@ def __init__(self, base_url, timeout=60,
             self.ssh_host = base_url[len('ssh://'):]
 
         self.timeout = timeout
+        self.pool = None
         self.max_pool_size = max_pool_size
-        self.pools = RecentlyUsedContainer(
-            pool_connections, dispose_func=lambda p: p.close()
-        )
         super().__init__()
 
     def _create_paramiko_client(self, base_url):
@@ -223,31 +232,16 @@ def _connect(self):
             self.ssh_client.connect(**self.ssh_params)
 
     def get_connection(self, url, proxies=None):
-        if not self.ssh_client:
-            return SSHConnectionPool(
-                ssh_client=self.ssh_client,
-                timeout=self.timeout,
-                maxsize=self.max_pool_size,
-                host=self.ssh_host
-            )
-        with self.pools.lock:
-            pool = self.pools.get(url)
-            if pool:
-                return pool
-
-            # Connection is closed try a reconnect
+        if self.pool is None:
             if self.ssh_client and not self.ssh_client.get_transport():
                 self._connect()
-
-            pool = SSHConnectionPool(
+            self.pool = SSHConnectionPool(
                 ssh_client=self.ssh_client,
                 timeout=self.timeout,
                 maxsize=self.max_pool_size,
                 host=self.ssh_host
             )
-            self.pools[url] = pool
-
-        return pool
+        return self.pool
 
     def close(self):
         super().close()