From e5429581fe53d9008e5adc87634b7c29f243d412 Mon Sep 17 00:00:00 2001
From: Sander Vrijders <sander.vrijders@ugent.be>
Date: Thu, 26 Oct 2017 15:30:44 +0200
Subject: rumba: Add persistent SSH connections

This adds persistent SSH connections to Rumba. A new SSH client will
be opened for every node in the experiment, which is closed after all
experimentation is done.
---
 rumba/model.py       |  7 +++++++
 rumba/ssh_support.py | 39 ++++++++++++++++-----------------------
 2 files changed, 23 insertions(+), 23 deletions(-)

diff --git a/rumba/model.py b/rumba/model.py
index 7d7e214..b44da88 100644
--- a/rumba/model.py
+++ b/rumba/model.py
@@ -198,6 +198,8 @@ class SSHConfig:
         self.hostname = hostname
         self.port = port
         self.proxy_server = proxy_server
+        self.client = None
+        self.proxy_client = None
 
     def set_username(self, username):
         self.username = username
@@ -809,6 +811,11 @@ class Experiment:
         self.dump_ssh_info()
 
     def swap_out(self):
+        for node in self.nodes:
+            if node.ssh_config.client is not None:
+                node.ssh_config.client.close()
+            if node.ssh_config.proxy_client is not None:
+                node.ssh_config.proxy_client.close()
         # Undo the testbed (testbed-specific)
         self.testbed.swap_out(self)
 
diff --git a/rumba/ssh_support.py b/rumba/ssh_support.py
index 5f69439..93b6ac2 100644
--- a/rumba/ssh_support.py
+++ b/rumba/ssh_support.py
@@ -155,9 +155,12 @@ def execute_commands(testbed, ssh_config, commands, time_out=3):
     be used when no timeout is needed
     """
 
-    ssh_client, proxy_client = ssh_connect(ssh_config.hostname, ssh_config.port,
+    if ssh_config.client is None:
+        client, proxy_client = ssh_connect(ssh_config.hostname, ssh_config.port,
                                            testbed.username, testbed.password,
                                            time_out, ssh_config.proxy_server)
+        ssh_config.client = client
+        ssh_config.proxy_client = proxy_client
 
     o = ""
     for command in commands:
@@ -167,7 +170,7 @@ def execute_commands(testbed, ssh_config, commands, time_out=3):
                                          command))
         envars = '. /etc/profile;'
         command = envars + ' ' + command
-        chan = ssh_client.get_transport().open_session()
+        chan = ssh_config.client.get_transport().open_session()
         stdout = chan.makefile()
         try:
             chan.exec_command(command)
@@ -175,14 +178,8 @@ def execute_commands(testbed, ssh_config, commands, time_out=3):
             raise SSHException('Failed to execute command')
         o = _print_stream(stdout)
         if (chan.recv_exit_status() != 0):
-            ssh_client.close()
-            if proxy_client is not None:
-                proxy_client.close()
             raise SSHException('A remote command returned an error.\n' + o)
 
-    ssh_client.close()
-    if proxy_client is not None:
-        proxy_client.close()
     return o
 
 
@@ -215,21 +212,23 @@ def write_text_to_file(testbed, ssh_config, text, file_name):
     @param text: string to be written in file
     @param file_name: file name (including full path) on the host
     """
-
-    ssh_client, proxy_client = ssh_connect(ssh_config.hostname, ssh_config.port,
+    if ssh_config.client is None:
+        client, proxy_client = ssh_connect(ssh_config.hostname, ssh_config.port,
                                            testbed.username, testbed.password,
                                            None, ssh_config.proxy_server)
+        ssh_config.client = client
+        ssh_config.proxy_client = proxy_client
 
     cmd = "touch " + file_name + "; chmod a+rwx " + file_name
 
     try:
-        stdin, stdout, stderr = ssh_client.exec_command(cmd)
+        stdin, stdout, stderr = ssh_config.client.exec_command(cmd)
         del stdin, stdout
         err = str(stderr.read()).strip('b\'\"\\n')
         if err != "":
             logger.error(err)
 
-        sftp_client = ssh_client.open_sftp()
+        sftp_client = ssh_config.client.open_sftp()
         remote_file = sftp_client.open(file_name, 'w')
 
         remote_file.write(text)
@@ -237,10 +236,6 @@ def write_text_to_file(testbed, ssh_config, text, file_name):
 
     except SSHException as e:
         raise SSHException('Failed to write text to remote file')
-    finally:
-        ssh_client.close()
-        if proxy_client is not None:
-            proxy_client.close()
 
 def copy_files_to_testbed(testbed, ssh_config, paths, destination):
     """
@@ -254,12 +249,15 @@ def copy_files_to_testbed(testbed, ssh_config, paths, destination):
     if destination is not '' and not destination.endswith('/'):
         destination = destination + '/'
 
-    ssh_client, proxy_client = ssh_connect(ssh_config.hostname, ssh_config.port,
+    if ssh_config.client is None:
+        client, proxy_client = ssh_connect(ssh_config.hostname, ssh_config.port,
                                            testbed.username, testbed.password,
                                            None, ssh_config.proxy_server)
+        ssh_config.client = client
+        ssh_config.proxy_client = proxy_client
 
     try:
-        sftp_client = ssh_client.open_sftp()
+        sftp_client = ssh_config.client.open_sftp()
 
         for path in paths:
             file_name = os.path.basename(path)
@@ -274,11 +272,6 @@ def copy_files_to_testbed(testbed, ssh_config, paths, destination):
 
     except Exception as e:
         raise SSHException('Failed to copy files to testbed')
-    finally:
-        ssh_client.close()
-        if proxy_client is not None:
-            proxy_client.close()
-
 
 def copy_file_to_testbed(testbed, ssh_config, path, destination):
     """
-- 
cgit v1.2.3