perf(storages): ISO chunks to an SSH compute reuse the SSH connection

Each 32 MiB chunk is its own request and opened a new SSH connection and
SFTP session (~200 ms). A worker now keeps its connection to the host for
the next chunk; each request still opens its own SFTP session (~50 ms),
so an SFTP error never reaches the next one. A connection is not reused
after an error, when the host dropped it or once it is 10 minutes old; an
idle one is closed after a minute, and a forked worker leaves the parent's
to the parent. The SFTP start, which paramiko waits for without a deadline,
is ended after 10 s. Live: 256 MiB in 8 chunks 4.5-4.9 s -> 3.7 s.

storages/test_ssh_pool.py

@@ -0,0 +1,141 @@
+"""Chunks to an SSH compute reuse this process's SSH connection: each request
+takes it for itself and opens its own SFTP session. A connection is not
+reused after an error, when the host dropped it, when it is too old, or by
+another login; an idle one is closed, and a forked worker leaves the
+parent's connections to the parent."""
+
+import os
+import threading
+import time
+import unittest
+from unittest.mock import MagicMock, patch
+
+from storages import upload
+from storages.upload import _ssh_pool, _UploadTarget
+from vrtManager.connection import CONN_SSH
+
+
+def client():
+    return MagicMock(**{"get_transport.return_value.is_active.return_value": True})
+
+
+KEY = ("192.0.2.1", 22, "root", "secret")
+
+
+def target(host="192.0.2.1", login="root", passwd="secret"):
+    return _UploadTarget(MagicMock(conn=CONN_SSH, host=host, login=login, passwd=passwd), "/pool", "x.iso", "id")
+
+
+class SSHPoolTestCase(unittest.TestCase):
+    def setUp(self):
+        self.clients = []
+        patcher = patch("paramiko.SSHClient", side_effect=lambda: self.clients.append(client()) or self.clients[-1])
+        patcher.start()
+        self.addCleanup(patcher.stop)
+        self.addCleanup(_ssh_pool.close_all)
+
+    def request(self, **kwargs):
+        with target(**kwargs) as t:
+            return t.ssh
+
+    def test_chunks_and_uploads_share_one_connection(self):
+        first = self.request()
+        self.assertIs(self.request(), first)
+        self.assertEqual(len(self.clients), 1)
+        self.assertEqual(first.open_sftp.call_count, 2)
+        first.close.assert_not_called()
+
+    def test_another_login_port_or_password_has_its_own_connection(self):
+        for kwargs in ({}, {"login": "admin"}, {"host": "192.0.2.1:2222"}, {"passwd": "changed"}):
+            self.request(**kwargs)
+        self.assertEqual(len(self.clients), 4)
+        self.assertEqual(self.clients[2].connect.call_args.kwargs["port"], 2222)
+
+    def test_two_requests_at_once_do_not_share_and_the_surplus_is_closed(self):
+        with target() as a, target() as b:
+            self.assertIsNot(a.ssh, b.ssh)
+        first, second = self.clients  # b is given back first and kept
+        first.close.assert_called_once()
+        second.close.assert_not_called()
+        self.assertIs(self.request(), second)
+
+    def test_a_connection_is_not_reused_after_an_error(self):
+        with self.assertRaises(OSError):
+            with target():
+                raise OSError("write failed")
+        self.clients[0].close.assert_called_once()
+        self.request()
+        self.assertEqual(len(self.clients), 2)
+
+    def test_a_dropped_connection_is_replaced(self):
+        stale = self.request()
+        stale.get_transport.return_value.is_active.return_value = False
+        self.assertIsNot(self.request(), stale)
+        stale.close.assert_called_once()
+
+    def test_a_connection_whose_sftp_fails_is_replaced(self):
+        stale = self.request()
+        stale.open_sftp.side_effect = EOFError("connection reset")
+        self.assertIsNot(self.request(), stale)
+        stale.close.assert_called_once()
+
+    def test_an_sftp_start_that_never_answers_is_ended(self):
+        stale = self.request()
+        closed = threading.Event()
+
+        def hang():  # paramiko's wait for the subsystem, woken by a close
+            closed.wait(5)
+            raise EOFError("closed")
+
+        stale.open_sftp.side_effect = hang
+        stale.close.side_effect = closed.set
+        start = time.monotonic()
+        with patch("storages.upload.SSH_CONNECT_TIMEOUT_SECONDS", 0.05):
+            self.assertIsNot(self.request(), stale)
+        self.assertLess(time.monotonic() - start, 1)
+
+    def test_a_connection_without_an_idle_timer_is_not_kept(self):
+        idle = client()
+        with patch("threading.Timer.start", side_effect=RuntimeError("can't start new thread")):
+            _ssh_pool.give(KEY, idle, time.monotonic())
+        idle.close.assert_called_once()
+        self.assertEqual(_ssh_pool._idle, {})
+
+    def test_a_connection_is_not_reused_after_its_lifetime(self):
+        stale = self.request()
+        with patch("storages.upload.time.monotonic", return_value=time.monotonic() + upload.SSH_MAX_AGE_SECONDS):
+            self.assertIsNot(self.request(), stale)
+        stale.close.assert_called_once()
+
+    def test_an_idle_connection_is_closed(self):
+        with patch("storages.upload.SSH_IDLE_SECONDS", 0.01):
+            idle = self.request()
+            time.sleep(0.2)
+        idle.close.assert_called_once()
+        self.assertIsNot(self.request(), idle)
+
+    def test_an_old_timer_does_not_close_a_connection_given_back_since(self):
+        first = self.request()
+        old_timer = _ssh_pool._idle[KEY][2]
+        self.request()  # taken and given back: a new timer
+        _ssh_pool._expire(KEY, old_timer)
+        first.close.assert_not_called()
+        self.assertIs(self.request(), first)
+
+    def test_a_forked_child_drops_the_parents_connections(self):
+        parent = self.request()
+        with _ssh_pool._lock:  # held while forking
+            pid = os.fork()
+        if pid == 0:
+            ok = _ssh_pool._idle == {} and parent.get_transport.return_value.atfork.called
+            with _ssh_pool._lock:
+                pass
+            os._exit(0 if ok and not parent.close.called else 1)
+        _, status = os.waitpid(pid, 0)
+        self.assertEqual(os.waitstatus_to_exitcode(status), 0)
+        parent.get_transport.return_value.atfork.assert_not_called()
+        self.assertIs(self.request(), parent)
+
+
+if __name__ == "__main__":
+    unittest.main()

storages/tests.py

@@ -18,6 +18,7 @@ from storages.upload import (
     _get_upload_lock_directory,
     _is_session_completed,
     _sftp_replace_file,
+    _ssh_pool,
 )
 from vrtManager.connection import CONN_SOCKET, CONN_SSH, CONN_TCP, CONN_TLS
 
@@ -117,6 +118,7 @@ class StorageUploadSecurityTests(TestCase):
         self.mock_ssh_conn.host = "127.0.0.1"
         self.mock_ssh_conn.login = "test"
         self.mock_ssh_conn.passwd = "test"
+        self.addCleanup(_ssh_pool.close_all)
 
     def tearDown(self):
         shutil.rmtree(self.temp_dir, ignore_errors=True)
@@ -635,6 +637,9 @@ class StorageUploadSecurityTests(TestCase):
             upload_id=upload_id,
         )
         self.assertEqual(remote_files[final_remote], b"DEBIAN_CHUNK0_DEBIAN_CHUNK1")
+        # one connection for the three requests, an SFTP session for each
+        mock_ssh.connect.assert_called_once()
+        self.assertEqual(mock_ssh.open_sftp.call_count, 3)
 
     def test_handle_uploaded_file_ssh_traversal_check(self):
         dummy_chunk = SimpleUploadedFile("test.iso", b"data")

storages/upload.py

@@ -1,5 +1,6 @@
 """Chunked ISO upload operations shared by the storage view."""
 
+import atexit
 import fcntl
 import json
 import logging
@@ -8,6 +9,7 @@ import posixpath
 import re
 import stat
 import tempfile
+import threading
 import time
 from xml.sax.saxutils import escape
 
@@ -20,6 +22,86 @@ logger = logging.getLogger(__name__)
 LOCK_TIMEOUT_SECONDS = 15
 SSH_CONNECT_TIMEOUT_SECONDS = 10
 SSH_IO_TIMEOUT_SECONDS = 30
+SSH_IDLE_SECONDS = 60
+SSH_MAX_AGE_SECONDS = 600
+
+
+class _SSHPool:
+    """SSH connections to computes, kept for the next chunk of this process
+    (one HTTP request per chunk; a new connection costs ~200 ms). A request
+    takes a connection for itself and opens its own SFTP session on it. An
+    idle connection is closed after SSH_IDLE_SECONDS, and none is reused
+    once it is SSH_MAX_AGE_SECONDS old, so host keys and logins are checked
+    again."""
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._idle = {}  # key -> (client, opened, timer)
+
+    @staticmethod
+    def _usable(client, opened):
+        transport = client.get_transport()
+        return transport is not None and transport.is_active() and time.monotonic() - opened < SSH_MAX_AGE_SECONDS
+
+    def take(self, key):
+        """An idle connection and when it was opened, or None"""
+        with self._lock:
+            entry = self._idle.pop(key, None)
+        if entry is None:
+            return None
+        client, opened, timer = entry
+        timer.cancel()
+        if self._usable(client, opened):
+            return client, opened
+        client.close()
+        return None
+
+    def give(self, key, client, opened):
+        if not self._usable(client, opened):
+            client.close()
+            return
+        timer = threading.Timer(SSH_IDLE_SECONDS, lambda: self._expire(key, timer))
+        timer.daemon = True
+        with self._lock:
+            kept = key not in self._idle
+            if kept:
+                try:
+                    timer.start()
+                except RuntimeError:  # no thread for the timer: not kept
+                    kept = False
+                else:
+                    self._idle[key] = (client, opened, timer)
+        if not kept:
+            client.close()
+
+    def _expire(self, key, timer):
+        with self._lock:
+            entry = self._idle.get(key)
+            if entry is None or entry[2] is not timer:  # taken since, perhaps given back
+                return
+            del self._idle[key]
+        entry[0].close()
+
+    def close_all(self):
+        with self._lock:
+            entries, self._idle = list(self._idle.values()), {}
+        for client, _opened, timer in entries:
+            timer.cancel()
+            client.close()
+
+    def forget(self):
+        """For a forked child: drop the parent's connections without ending
+        their sessions; the parent keeps using them"""
+        entries, self._lock, self._idle = list(self._idle.values()), threading.Lock(), {}
+        for client, _opened, _timer in entries:
+            transport = client.get_transport()
+            if transport is not None:
+                transport.atfork()
+
+
+_ssh_pool = _SSHPool()
+os.register_at_fork(after_in_child=_ssh_pool.forget)
+atexit.register(_ssh_pool.close_all)
 
 
 def _fsync_dir(dir_path):
@@ -348,32 +430,71 @@ class _UploadTarget:
     def __enter__(self):
         if self.remote:
             hostname, separator, port = self.conn.host.partition(":")
-            self.ssh = paramiko.SSHClient()
+            self.key = (hostname, int(port) if separator else 22, self.conn.login, self.conn.passwd)
+            taken = _ssh_pool.take(self.key)
+            if taken is not None:
+                self.ssh, self.opened = taken
+                try:
+                    self._open_sftp()
+                    return self
+                except Exception:  # a connection the host dropped: open a new one
+                    self.ssh.close()
+            self.ssh, self.opened = self._connect(), time.monotonic()
             try:
-                self.ssh.load_system_host_keys()
+                self._open_sftp()
-                self.ssh.set_missing_host_key_policy(paramiko.RejectPolicy())
-                self.ssh.connect(
-                    hostname=hostname, port=int(port) if separator else 22,
-                    username=self.conn.login, password=self.conn.passwd,
-                    timeout=SSH_CONNECT_TIMEOUT_SECONDS,
-                    banner_timeout=SSH_CONNECT_TIMEOUT_SECONDS,
-                    auth_timeout=SSH_CONNECT_TIMEOUT_SECONDS,
-                    channel_timeout=SSH_CONNECT_TIMEOUT_SECONDS,
-                )
-                self.sftp = self.ssh.open_sftp()
-                self.sftp.get_channel().settimeout(SSH_IO_TIMEOUT_SECONDS)
             except Exception:
                 self.ssh.close()
                 raise
         return self
 
-    def __exit__(self, *_):
+    def _connect(self):
+        hostname, port = self.key[:2]
+        ssh = paramiko.SSHClient()
         try:
-            if self.sftp is not None:
+            ssh.load_system_host_keys()
-                self.sftp.close()
+            ssh.set_missing_host_key_policy(paramiko.RejectPolicy())
+            ssh.connect(
+                hostname=hostname, port=port,
+                username=self.conn.login, password=self.conn.passwd,
+                timeout=SSH_CONNECT_TIMEOUT_SECONDS,
+                banner_timeout=SSH_CONNECT_TIMEOUT_SECONDS,
+                auth_timeout=SSH_CONNECT_TIMEOUT_SECONDS,
+                channel_timeout=SSH_CONNECT_TIMEOUT_SECONDS,
+            )
+        except Exception:
+            ssh.close()
+            raise
+        return ssh
+
+    def _open_sftp(self):
+        # paramiko waits for the SFTP subsystem without a deadline; on a dead
+        # connection closing it ends the wait with an error
+        watchdog = threading.Timer(SSH_CONNECT_TIMEOUT_SECONDS, self.ssh.close)
+        watchdog.daemon = True
+        watchdog.start()
+        try:
+            self.sftp = self.ssh.open_sftp()
         finally:
-            if self.ssh is not None:
+            watchdog.cancel()
-                self.ssh.close()
+        transport = self.ssh.get_transport()
+        if transport is None or not transport.is_active():
+            self.sftp.close()
+            raise paramiko.SSHException(_("The SFTP session did not start in time"))
+        self.sftp.get_channel().settimeout(SSH_IO_TIMEOUT_SECONDS)
+
+    def __exit__(self, exc_type, *_):
+        if self.ssh is None:
+            return
+        try:
+            self.sftp.close()
+        except Exception:
+            self.ssh.close()
+            raise
+        # after an error the connection is not trusted for the next chunk
+        if exc_type is None:
+            _ssh_pool.give(self.key, self.ssh, self.opened)
+        else:
+            self.ssh.close()
 
     def info(self, path):
         try: