Roll forward a fix for #862. (#876)

Support retry limits in metadata retrieval without changing guest
behaviors.

4 files changed, 23 insertions, 21 deletions

diff --git a/packages/python-google-compute-engine/google_compute_engine/metadata_scripts/script_retriever.py b/packages/python-google-compute-engine/google_compute_engine/metadata_scripts/script_retriever.py
index 3de425b..7acd26f 100644
--- a/packages/python-google-compute-engine/google_compute_engine/metadata_scripts/script_retriever.py
+++ b/packages/python-google-compute-engine/google_compute_engine/metadata_scripts/script_retriever.py
@@ -97,7 +97,7 @@ class ScriptRetriever(object):
 
     if not self.token:
       response = self.watcher.GetMetadata(
-          self.token_metadata_key, recursive=False, retry=False)
+          self.token_metadata_key, recursive=False, retry_limit=3)
 
       if not response:
         self.logger.info(
diff --git a/packages/python-google-compute-engine/google_compute_engine/metadata_scripts/tests/script_retriever_test.py b/packages/python-google-compute-engine/google_compute_engine/metadata_scripts/tests/script_retriever_test.py
index 0df8830..558b0b0 100644
--- a/packages/python-google-compute-engine/google_compute_engine/metadata_scripts/tests/script_retriever_test.py
+++ b/packages/python-google-compute-engine/google_compute_engine/metadata_scripts/tests/script_retriever_test.py
@@ -89,7 +89,7 @@ class ScriptRetrieverTest(unittest.TestCase):
     # GetMetadata includes a prefix, so remove it.
     stripped_url = token_url.replace(metadata_prefix, '')
     mock_get_metadata.assert_called_once_with(
-        stripped_url, recursive=False, retry=False)
+        stripped_url, recursive=False, retry_limit=3)
 
     self.assertEqual(self.retriever.token, 'foo bar')
 
@@ -119,7 +119,7 @@ class ScriptRetrieverTest(unittest.TestCase):
     prefix = 'http://metadata.google.internal/computeMetadata/v1/'
     stripped_url = token_url.replace(prefix, '')
     mock_get_metadata.assert_called_once_with(
-        stripped_url, recursive=False, retry=False)
+        stripped_url, recursive=False, retry_limit=3)
     mock_download_url.assert_called_once_with(auth_url, self.dest_dir)
 
     self.assertIsNone(self.retriever.token)
diff --git a/packages/python-google-compute-engine/google_compute_engine/metadata_watcher.py b/packages/python-google-compute-engine/google_compute_engine/metadata_watcher.py
index 72a12f9..e399839 100644
--- a/packages/python-google-compute-engine/google_compute_engine/metadata_watcher.py
+++ b/packages/python-google-compute-engine/google_compute_engine/metadata_watcher.py
@@ -154,7 +154,7 @@ class MetadataWatcher(object):
 
   def _HandleMetadataUpdate(
       self, metadata_key='', recursive=True, wait=True, timeout=None,
-      retry=True):
+      retry_limit=None):
     """Wait for a successful metadata response.
 
     Args:
@@ -162,25 +162,25 @@ class MetadataWatcher(object):
       recursive: bool, True if we should recursively watch for metadata changes.
       wait: bool, True if we should wait for a metadata change.
       timeout: int, timeout in seconds for returning metadata output.
-      retry: bool, True if we should retry on failure.
+      retry_limit: int or None, limit for number of times to retry on failure.
+          Retry indefinitely when set to None. Do not retry when set to zero.
 
     Returns:
       json, the deserialized contents of the metadata server.
     """
     exception = None
-    while True:
+    while retry_limit is None or retry_limit >= 0:
       try:
         return self._GetMetadataUpdate(
             metadata_key=metadata_key, recursive=recursive, wait=wait,
             timeout=timeout)
       except (httpclient.HTTPException, socket.error, urlerror.URLError) as e:
+        if retry_limit is not None:
+          retry_limit -= 1
         if not isinstance(e, type(exception)):
           exception = e
           self.logger.error('GET request error retrieving metadata. %s.', e)
-        if retry:
-          continue
-        else:
-          break
+        time.sleep(1)
 
   def WatchMetadata(
       self, handler, metadata_key='', recursive=True, timeout=None):
@@ -202,18 +202,19 @@ class MetadataWatcher(object):
         self.logger.exception('Exception calling the response handler. %s.', e)
 
   def GetMetadata(
-      self, metadata_key='', recursive=True, timeout=None, retry=True):
+      self, metadata_key='', recursive=True, timeout=None, retry_limit=None):
     """Retrieve the contents of metadata server for a metadata key.
 
     Args:
       metadata_key: string, the metadata key to watch for changes.
       recursive: bool, True if we should recursively watch for metadata changes.
       timeout: int, timeout in seconds for returning metadata output.
-      retry: bool, True if we should retry on failure.
+      retry_limit: int or None, limit for number of times to retry on failure.
+          Retry indefinitely when set to None. Do not retry when set to zero.
 
     Returns:
       json, the deserialized contents of the metadata server or None if error.
     """
     return self._HandleMetadataUpdate(
         metadata_key=metadata_key, recursive=recursive, wait=False,
-        timeout=timeout, retry=retry)
+        timeout=timeout, retry_limit=retry_limit)
diff --git a/packages/python-google-compute-engine/google_compute_engine/tests/metadata_watcher_test.py b/packages/python-google-compute-engine/google_compute_engine/tests/metadata_watcher_test.py
index 1bce509..fe8aad9 100644
--- a/packages/python-google-compute-engine/google_compute_engine/tests/metadata_watcher_test.py
+++ b/packages/python-google-compute-engine/google_compute_engine/tests/metadata_watcher_test.py
@@ -259,12 +259,12 @@ class MetadataWatcherTest(unittest.TestCase):
     metadata_key = 'instance/id'
     recursive = False
     wait = False
-    retry = True
+    retry_limit = None
 
     self.assertEqual(
         self.mock_watcher._HandleMetadataUpdate(
             metadata_key=metadata_key, recursive=recursive, wait=wait,
-            timeout=None, retry=retry),
+            timeout=None, retry_limit=retry_limit),
         {})
     expected_calls = [
         mock.call(
@@ -282,12 +282,12 @@ class MetadataWatcherTest(unittest.TestCase):
     metadata_key = 'instance/id'
     recursive = False
     wait = False
-    retry = False
+    retry_limit = 0
 
     self.assertIsNone(
         self.mock_watcher._HandleMetadataUpdate(
             metadata_key=metadata_key, recursive=recursive, wait=wait,
-            timeout=None, retry=retry))
+            timeout=None, retry_limit=retry_limit))
     expected_calls = [
         mock.call(
             metadata_key=metadata_key, recursive=recursive, wait=wait,
@@ -333,7 +333,8 @@ class MetadataWatcherTest(unittest.TestCase):
 
     self.assertEqual(self.mock_watcher.GetMetadata(), {})
     mock_response.assert_called_once_with(
-        metadata_key='', recursive=True, wait=False, timeout=None, retry=True)
+        metadata_key='', recursive=True, wait=False, timeout=None,
+        retry_limit=None)
     self.mock_watcher.logger.exception.assert_not_called()
 
   def testGetMetadataArgs(self):
@@ -342,15 +343,15 @@ class MetadataWatcherTest(unittest.TestCase):
     self.mock_watcher._HandleMetadataUpdate = mock_response
     metadata_key = 'instance/id'
     recursive = False
-    retry = False
+    retry_limit = None
 
     response = self.mock_watcher.GetMetadata(
         metadata_key=metadata_key, recursive=recursive, timeout=60,
-        retry=retry)
+        retry_limit=retry_limit)
     self.assertEqual(response, {})
     mock_response.assert_called_once_with(
         metadata_key=metadata_key, recursive=False, wait=False, timeout=60,
-        retry=False)
+        retry_limit=None)
     self.mock_watcher.logger.exception.assert_not_called()