Merge pull request #2560 from damz/pr/global-timeout

Use urllib timeout param instead of hacking socket global timeout. Fixes #2560, #1935.

2 files changed, 28 insertions, 27 deletions

diff --git a/boto/utils.py b/boto/utils.py
index cfd1a758..dca332fe 100644
--- a/boto/utils.py
+++ b/boto/utils.py
@@ -39,8 +39,6 @@
 Some handy utility functions used by several classes.
 """
 
-import socket
-import imp
 import subprocess
 import time
 import logging.handlers
@@ -197,7 +195,7 @@ def get_aws_metadata(headers, provider=None):
     return metadata
 
 
-def retry_url(url, retry_on_404=True, num_retries=10):
+def retry_url(url, retry_on_404=True, num_retries=10, timeout=None):
     """
     Retry a url.  This is specifically used for accessing the metadata
     service on an instance.  Since this address should never be proxied
@@ -209,7 +207,7 @@ def retry_url(url, retry_on_404=True, num_retries=10):
             proxy_handler = urllib.request.ProxyHandler({})
             opener = urllib.request.build_opener(proxy_handler)
             req = urllib.request.Request(url)
-            r = opener.open(req)
+            r = opener.open(req, timeout=timeout)
             result = r.read()
 
             if(not isinstance(result, six.string_types) and
@@ -232,17 +230,18 @@ def retry_url(url, retry_on_404=True, num_retries=10):
     return ''
 
 
-def _get_instance_metadata(url, num_retries):
-    return LazyLoadMetadata(url, num_retries)
+def _get_instance_metadata(url, num_retries, timeout=None):
+    return LazyLoadMetadata(url, num_retries, timeout)
 
 
 class LazyLoadMetadata(dict):
-    def __init__(self, url, num_retries):
+    def __init__(self, url, num_retries, timeout=None):
         self._url = url
         self._num_retries = num_retries
         self._leaves = {}
         self._dicts = []
-        data = boto.utils.retry_url(self._url, num_retries=self._num_retries)
+        self._timeout = timeout
+        data = boto.utils.retry_url(self._url, num_retries=self._num_retries, timeout=self._timeout)
         if data:
             fields = data.split('\n')
             for field in fields:
@@ -282,7 +281,8 @@ class LazyLoadMetadata(dict):
                     val = boto.utils.retry_url(
                         self._url + urllib.parse.quote(resource,
                                                        safe="/:"),
-                        num_retries=self._num_retries)
+                        num_retries=self._num_retries,
+                        timeout=self._timeout)
                     if val and val[0] == '{':
                         val = json.loads(val)
                         break
@@ -389,17 +389,11 @@ def get_instance_metadata(version='latest', url='http://169.254.169.254',
     will time out after the specified number of seconds.
 
     """
-    if timeout is not None:
-        original = socket.getdefaulttimeout()
-        socket.setdefaulttimeout(timeout)
     try:
         metadata_url = _build_instance_metadata_url(url, version, data)
-        return _get_instance_metadata(metadata_url, num_retries=num_retries)
+        return _get_instance_metadata(metadata_url, num_retries=num_retries, timeout=timeout)
     except urllib.error.URLError as e:
         return None
-    finally:
-        if timeout is not None:
-            socket.setdefaulttimeout(original)
 
 
 def get_instance_identity(version='latest', url='http://169.254.169.254',
@@ -410,14 +404,11 @@ def get_instance_identity(version='latest', url='http://169.254.169.254',
     iid = {}
     base_url = _build_instance_metadata_url(url, version,
                                             'dynamic/instance-identity/')
-    if timeout is not None:
-        original = socket.getdefaulttimeout()
-        socket.setdefaulttimeout(timeout)
     try:
-        data = retry_url(base_url, num_retries=num_retries)
+        data = retry_url(base_url, num_retries=num_retries, timeout=timeout)
         fields = data.split('\n')
         for field in fields:
-            val = retry_url(base_url + '/' + field + '/')
+            val = retry_url(base_url + '/' + field + '/', num_retries=num_retries, timeout=timeout)
             if val[0] == '{':
                 val = json.loads(val)
             if field:
@@ -425,15 +416,12 @@ def get_instance_identity(version='latest', url='http://169.254.169.254',
         return iid
     except urllib.error.URLError as e:
         return None
-    finally:
-        if timeout is not None:
-            socket.setdefaulttimeout(original)
 
 
 def get_instance_userdata(version='latest', sep=None,
-                          url='http://169.254.169.254'):
+                          url='http://169.254.169.254', timeout=None, num_retries=5):
     ud_url = _build_instance_metadata_url(url, version, 'user-data')
-    user_data = retry_url(ud_url, retry_on_404=False)
+    user_data = retry_url(ud_url, retry_on_404=False, num_retries=num_retries, timeout=timeout)
     if user_data:
         if sep:
             l = user_data.split(sep)
diff --git a/tests/unit/utils/test_utils.py b/tests/unit/utils/test_utils.py
index c170506a..d96978c1 100644
--- a/tests/unit/utils/test_utils.py
+++ b/tests/unit/utils/test_utils.py
@@ -272,7 +272,20 @@ class TestLazyLoadMetadata(unittest.TestCase):
 
         boto.utils.retry_url.assert_called_with(
             'http://169.254.169.254/latest/user-data',
-            retry_on_404=False)
+            retry_on_404=False,
+            num_retries=5, timeout=None)
+
+    def test_user_data_timeout(self):
+        self.set_normal_response(['foo'])
+
+        userdata = get_instance_userdata(timeout=1, num_retries=2)
+
+        self.assertEqual('foo', userdata)
+
+        boto.utils.retry_url.assert_called_with(
+            'http://169.254.169.254/latest/user-data',
+            retry_on_404=False,
+            num_retries=2, timeout=1)
 
 
 class TestStringToDatetimeParsing(unittest.TestCase):