From a74bf43a9ca5fda7a188d564fa9d1d86a7ddcb4b Mon Sep 17 00:00:00 2001
From: Jason Wang <jawang35@gmail.com>
Date: Mon, 5 Nov 2018 15:20:09 -0800
Subject: [PATCH] Timeout SSH wait after 5 minutes (#3)

---
 environment.yml |  1 -
 lib/ssh.py      | 15 ++++++++-------
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/environment.yml b/environment.yml
index 390cda2..519a9a3 100644
--- a/environment.yml
+++ b/environment.yml
@@ -6,4 +6,3 @@ dependencies:
   - boto3=1.9.31
   - python=3.7.0
   - fabric=2.4.0
-
diff --git a/lib/ssh.py b/lib/ssh.py
index 1d58b7e..9939236 100644
--- a/lib/ssh.py
+++ b/lib/ssh.py
@@ -1,5 +1,6 @@
 from io import StringIO
 from time import sleep
+from timeit import default_timer
 from fabric import Connection
 from paramiko import RSAKey
 from paramiko.ssh_exception import NoValidConnectionsError
@@ -10,7 +11,7 @@ class SSH():
     SSH context manager for creating an SSH connection.
 
     On enter an SSH connection is attempted every 5 seconds until successful.
-    An exception is raised after 60 attempts.
+    An exception is raised after 5 minutes.
 
     On exit the connection is closed.
 
@@ -26,7 +27,6 @@ def __init__(self, host, user, private_key):
         self.host = host
         self.user = user
         self.private_key = RSAKey.from_private_key(StringIO(private_key))
-        self.wait_count = 0
 
     def __enter__(self):
         self.connection = Connection(
@@ -35,20 +35,21 @@ def __enter__(self):
             connect_kwargs={'pkey': self.private_key},
         )
         print(f'Waiting for SSH to become available on {self.host}...')
-        self.wait_for_ssh()
+        self.wait_for_ssh(default_timer())
         return self.connection
 
     def __exit__(self, type, value, traceback):
         print(f'Closing SSH connection to {self.host}...')
         self.connection.close()
 
-    def wait_for_ssh(self):
-        self.wait_count += 1
+    def wait_for_ssh(self, start):
         try:
             self.connection.open()
         except NoValidConnectionsError:
-            if self.wait_count >= 60:
+            # Error after 5 minutes. Otherwise retry.
+            now = default_timer()
+            if now - start > 300:
                 raise
 
             sleep(5)
-            self.wait_for_ssh()
+            self.wait_for_ssh(start)