From 702fc5844bdd56d88b7c4b947198fcc25cd693dd Mon Sep 17 00:00:00 2001
From: Robert Barron <rmbarron@google.com>
Date: Wed, 13 Feb 2019 07:01:29 -0800
Subject: [PATCH] Add retry to _metadata.ping()

---
 google/auth/compute_engine/_metadata.py | 33 +++++++++++++++----------
 1 file changed, 20 insertions(+), 13 deletions(-)

diff --git a/google/auth/compute_engine/_metadata.py b/google/auth/compute_engine/_metadata.py
index c47be3fae..d8004bb24 100644
--- a/google/auth/compute_engine/_metadata.py
+++ b/google/auth/compute_engine/_metadata.py
@@ -51,13 +51,15 @@
     _METADATA_DEFAULT_TIMEOUT = 3
 
 
-def ping(request, timeout=_METADATA_DEFAULT_TIMEOUT):
+def ping(request, timeout=_METADATA_DEFAULT_TIMEOUT, retry_count=3):
     """Checks to see if the metadata server is available.
 
     Args:
         request (google.auth.transport.Request): A callable used to make
             HTTP requests.
         timeout (int): How long to wait for the metadata server to respond.
+        retry_count (int): How many times to attempt connecting to metadata
+            server using above timeout.
 
     Returns:
         bool: True if the metadata server is reachable, False otherwise.
@@ -68,18 +70,23 @@ def ping(request, timeout=_METADATA_DEFAULT_TIMEOUT):
     #       could lead to false negatives in the event that we are on GCE, but
     #       the metadata resolution was particularly slow. The latter case is
     #       "unlikely".
-    try:
-        response = request(
-            url=_METADATA_IP_ROOT, method='GET', headers=_METADATA_HEADERS,
-            timeout=timeout)
-
-        metadata_flavor = response.headers.get(_METADATA_FLAVOR_HEADER)
-        return (response.status == http_client.OK and
-                metadata_flavor == _METADATA_FLAVOR_VALUE)
-
-    except exceptions.TransportError:
-        _LOGGER.info('Compute Engine Metadata server unavailable.')
-        return False
+    retries = 0
+    while retries < retry_count:
+        try:
+            response = request(
+                url=_METADATA_IP_ROOT, method='GET', headers=_METADATA_HEADERS,
+                timeout=timeout)
+
+            metadata_flavor = response.headers.get(_METADATA_FLAVOR_HEADER)
+            return (response.status == http_client.OK and
+                    metadata_flavor == _METADATA_FLAVOR_VALUE)
+
+        except exceptions.TransportError:
+            _LOGGER.info('Compute Engine Metadata server unavailable on'
+                         'attempt %s of %s', retries+1, retry_count)
+            retries += 1
+
+    return False
 
 
 def get(request, path, root=_METADATA_ROOT, recursive=False):