added retry logic to tklbam stsagent command

cmd_internals/cmd_stsagent.py

@@ -14,6 +14,11 @@
 import sys
 from registry import hub_backups
 import hub
+from retry import retry
+
+@retry(5, backoff=2)
+def get_credentials(hb):
+    return hb.get_credentials()
 
 def usage(e=None):
     if e:
@@ -31,6 +36,7 @@ def format(creds):
     values = [ creds[k] for k in ('accesskey', 'secretkey', 'sessiontoken', 'expiration') ]
     return " ".join(values)
 
+
 def main():
     args = sys.argv[1:]
     if args:
@@ -41,7 +47,7 @@ def main():
     except hub.Backups.NotInitialized, e:
         print >> sys.stderr, "error: " + str(e)
 
-    creds = hb.get_credentials()
+    creds = get_credentials(hb)
     if creds.type != 'iamrole':
         fatal("STS agent incompatible with '%s' type credentials" % creds.type)
 

retry.py

@@ -0,0 +1,44 @@
+"""
+Usage examples:
+
+    @retry(3)               # will retry 3 times (total 4 attempts), no backoff
+    @retry(3, backoff=1)    # will retry 3 times, increasing delay by 100% each attempt
+    @retry(3, backoff=2)    # will retry 3 times, increasing delay by 200% each attempt
+
+
+"""
+from time import sleep
+
+def retry(retries, delay=1, backoff=0, fatal_exceptions=None):
+    """
+    Argument:
+
+        retries             how many times to retry if exception is raised
+        delay               how many seconds to delay in case of failure
+        backoff             linear backoff factor (e.g., 0 = no backoff, 1 = 100% step increase)
+        fatal_exc           fatal exceptions are unrecoverable, raised immediately
+    """
+    def decorator(func):
+
+        _fatal_exceptions = (SyntaxError, KeyboardInterrupt, SystemExit)
+        if fatal_exceptions:
+
+            if issubclass(fatal_exceptions, Exception):
+                _fatal_exceptions += (fatal_exceptions,)
+            else:
+                _fatal_exceptions += fatal_exceptions
+
+        def wrapper(*args, **kwargs):
+            for attempt in range(retries + 1):
+                try:
+                    return func(*args, **kwargs)
+                except _fatal_exceptions:
+                    raise
+                except:
+                    if attempt < retries and delay:
+                        sleep(delay + delay * attempt * backoff)
+            else:
+                raise
+
+        return wrapper
+    return decorator