mirror of
https://github.com/retspen/webvirtcloud
synced 2026-11-08 11:16:47 +00:00
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.
This commit is contained in:
catborise
parent
ff2f16b226
commit
7b5c2cd057
3 changed files with 285 additions and 18 deletions
141
storages/test_ssh_pool.py
Normal file
141
storages/test_ssh_pool.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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.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)
|
||||
self._open_sftp()
|
||||
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:
|
||||
self.sftp.close()
|
||||
ssh.load_system_host_keys()
|
||||
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:
|
||||
self.ssh.close()
|
||||
watchdog.cancel()
|
||||
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:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue