Add timeout on HTTP request calls.android-q-preview-1android-o-mr1-iot-release-1.0.10

Test: go/vtslab-schedule-dev
Bug: 80503411
Change-Id: Ib9cb194403bc6d824d98d630491dc9faed7d1788

3 files changed, 108 insertions, 38 deletions

diff --git a/harnesses/host_controller/build/build_provider_pab.py b/harnesses/host_controller/build/build_provider_pab.py
index 7ec472c..9d932b6 100644
--- a/harnesses/host_controller/build/build_provider_pab.py
+++ b/harnesses/host_controller/build/build_provider_pab.py
@@ -44,6 +44,9 @@ from host_controller.build import build_provider
 GET = 'GET'
 POST = 'POST'
 
+# timeout seconds for requests
+REQUESTS_TIMEOUT_SECONDS = 60
+
 
 class BuildProviderPAB(build_provider.BuildProvider):
     """Client that manages Partner Android Build downloading.
@@ -233,7 +236,12 @@ class BuildProviderPAB(build_provider.BuildProvider):
         headers['Content-Type'] = 'application/json'
         headers['x-alkali-account'] = account_id
 
-        response = requests.post(self.SVC_URL, data=data, headers=headers)
+        try:
+            response = requests.post(self.SVC_URL, data=data, headers=headers,
+                                     timeout=REQUESTS_TIMEOUT_SECONDS)
+        except requests.exceptions.Timeout as e:
+            logging.exception(e)
+            raise ValueError("Request timeout.")
 
         responseJSON = {}
 
@@ -310,11 +318,14 @@ class BuildProviderPAB(build_provider.BuildProvider):
             url = path_urljoin(self.BASE_URL, 'build', 'builds', action,
                                branch, target, dummy,
                                dummy) + '?a=' + str(account_id)
-
-            response = requests.get(url, headers=headers)
             try:
+                response = requests.get(url, headers=headers,
+                                        timeout=REQUESTS_TIMEOUT_SECONDS)
                 responseJSON = response.json()
                 builds = responseJSON['build']
+            except requests.exceptions.Timeout as e:
+                logging.exception(e)
+                raise ValueError("Request timeout.")
             except ValueError as e:
                 logging.exception(e)
                 raise ValueError("Backend error -- check your account ID")
@@ -341,6 +352,10 @@ class BuildProviderPAB(build_provider.BuildProvider):
                     except requests.HTTPError:
                         logging.debug("The build is not signed.")
                         build["signed"] = False
+                    except requests.exceptions.Timeout as e:
+                        logging.debug("Server is not responding.")
+                        logging.exception(e)
+                        build["signed"] = False
             return builds
 
     def GetLatestBuildId(self, account_id, branch, target, method=GET):
@@ -466,10 +481,14 @@ class BuildProviderPAB(build_provider.BuildProvider):
                                    branch, target, build_id,
                                    artifact_name) + '?a=' + str(account_id)
 
-            response = requests.get(get_url, headers=headers)
             try:
+                response = requests.get(get_url, headers=headers,
+                                        timeout=REQUESTS_TIMEOUT_SECONDS)
                 responseJSON = response.json()
                 return responseJSON['url']
+            except requests.exceptions.Timeout as e:
+                logging.exception(e)
+                raise ValueError("Request timeout.")
             except ValueError:
                 raise ValueError("Backend error -- check your account ID")
 
@@ -485,10 +504,9 @@ class BuildProviderPAB(build_provider.BuildProvider):
         """
         try:
             response = self.GetResponseWithURL(download_url)
-        except requests.HTTPError as httpError:
-            logging.exception(httpError)
+        except (requests.HTTPError, requests.exceptions.Timeout) as error:
+            logging.exception(error)
             return False
-
         logging.info('%s now downloading...', download_url)
         with open(filename, 'wb') as handle:
             for block in response.iter_content(self.DEFAULT_CHUNK_SIZE):
@@ -661,11 +679,15 @@ class BuildProviderPAB(build_provider.BuildProvider):
 
         Returns:
             A Response object received from the server.
+
+        Raises:
+            requests.HTTPError if response.status_code is not 200.
+            requests.exceptions.Timeout if the server does not respond.
         """
         headers = {}
         self._credentials.apply(headers)
-
-        response = requests.get(url, headers=headers, stream=True)
+        response = requests.get(url, headers=headers, stream=True,
+                                timeout=REQUESTS_TIMEOUT_SECONDS)
         response.raise_for_status()
 
         return response
diff --git a/harnesses/host_controller/utils/parser/pb2_utils.py b/harnesses/host_controller/utils/parser/pb2_utils.py
index ef6bc83..af96588 100644
--- a/harnesses/host_controller/utils/parser/pb2_utils.py
+++ b/harnesses/host_controller/utils/parser/pb2_utils.py
@@ -18,6 +18,9 @@ import json
 import logging
 import requests
 
+# timeout seconds for requests
+REQUESTS_TIMEOUT_SECONDS = 60
+
 
 def FillDictAndPost(msg,
                     dict_to_fill,
@@ -64,8 +67,13 @@ def FillDictAndPost(msg,
                                       filters, caller_name)
 
     if terminal:
-        response = requests.post(
-            url, data=json.dumps(dict_to_fill), headers=headers)
+        try:
+            response = requests.post(url, data=json.dumps(dict_to_fill),
+                                     headers=headers,
+                                     timeout=REQUESTS_TIMEOUT_SECONDS)
+        except requests.exceptions.Timeout as e:
+            logging.exception(e)
+            return False
         if response.status_code != requests.codes.ok:
             logging.error("%s error: %s", caller_name, response)
             ret = ret and False
diff --git a/harnesses/host_controller/vti_interface/vti_endpoint_client.py b/harnesses/host_controller/vti_interface/vti_endpoint_client.py
index aafbe91..17fdc9d 100644
--- a/harnesses/host_controller/vti_interface/vti_endpoint_client.py
+++ b/harnesses/host_controller/vti_interface/vti_endpoint_client.py
@@ -43,6 +43,9 @@ SCHEDULE_INFO_PB2_ATTR_FILTERS = {
     "name": "build_target",
 }
 
+# timeout seconds for requests
+REQUESTS_TIMEOUT_SECONDS = 60
+
 
 class VtiEndpointClient(object):
     """VTI (Vendor Test Infrastructure) endpoint client.
@@ -80,10 +83,15 @@ class VtiEndpointClient(object):
         url = self._url + "build/v1/set"
         fail = False
         for build in builds:
-            response = requests.post(url, data=json.dumps(build),
-                                     headers=self._headers)
-            if response.status_code != requests.codes.ok:
-                logging.error("UploadBuildInfo error: %s", response)
+            try:
+                response = requests.post(url, data=json.dumps(build),
+                                         headers=self._headers,
+                                         timeout=REQUESTS_TIMEOUT_SECONDS)
+                if response.status_code != requests.codes.ok:
+                    logging.error("UploadBuildInfo error: %s", response)
+                    fail = True
+            except requests.exceptions.Timeout as e:
+                logging.exception(e)
                 fail = True
         if fail:
             return False
@@ -113,8 +121,10 @@ class VtiEndpointClient(object):
 
         try:
             response = requests.post(url, data=json.dumps(payload),
-                                     headers=self._headers)
-        except requests.exceptions.ConnectionError as e:
+                                     headers=self._headers,
+                                     timeout=REQUESTS_TIMEOUT_SECONDS)
+        except (requests.exceptions.ConnectionError,
+                requests.exceptions.Timeout) as e:
             logging.exception(e)
             return False
         if response.status_code != requests.codes.ok:
@@ -139,9 +149,13 @@ class VtiEndpointClient(object):
         url = self._url + "schedule/v1/clear"
         succ = True
         if clear_schedule:
-            response = requests.post(
-                url, data=json.dumps({"manifest_branch": "na"}),
-                headers=self._headers)
+            try:
+                response = requests.post(
+                    url, data=json.dumps({"manifest_branch": "na"}),
+                    headers=self._headers, timeout=REQUESTS_TIMEOUT_SECONDS)
+            except requests.exceptions.Timeout as e:
+                logging.exception(e)
+                return False
             if response.status_code != requests.codes.ok:
                 logging.error("UploadScheduleInfo error: %s", response)
                 succ = False
@@ -175,8 +189,13 @@ class VtiEndpointClient(object):
         url = self._url + "lab/v1/clear"
         succ = True
         if clear_labinfo:
-            response = requests.post(url, data=json.dumps({"name": "na"}),
-                                     headers=self._headers)
+            try:
+                response = requests.post(url, data=json.dumps({"name": "na"}),
+                                         headers=self._headers,
+                                         timeout=REQUESTS_TIMEOUT_SECONDS)
+            except requests.exceptions.Timeout as e:
+                logging.exception(e)
+                return False
             if response.status_code != requests.codes.ok:
                 logging.error("UploadLabInfo error: %s", response)
                 succ = False
@@ -212,10 +231,15 @@ class VtiEndpointClient(object):
                                 device.device_equipment)
                         new_host["device"].append(new_device)
                 lab["host"].append(new_host)
-            response = requests.post(url, data=json.dumps(lab),
-                                     headers=self._headers)
-            if response.status_code != requests.codes.ok:
-                logging.error("UploadLabInfo error: %s", response)
+            try:
+                response = requests.post(url, data=json.dumps(lab),
+                                         headers=self._headers,
+                                         timeout=REQUESTS_TIMEOUT_SECONDS)
+                if response.status_code != requests.codes.ok:
+                    logging.error("UploadLabInfo error: %s", response)
+                    succ = False
+            except requests.exceptions.Timeout as e:
+                logging.exception(e)
                 succ = False
         return succ
 
@@ -235,8 +259,14 @@ class VtiEndpointClient(object):
             return None, {}
 
         url = self._url + "job/v1/lease"
-        response = requests.post(url, data=json.dumps({"hostname": hostname}),
-                                 headers=self._headers)
+        try:
+            response = requests.post(url, data=json.dumps({"hostname": hostname}),
+                                     headers=self._headers,
+                                     timeout=REQUESTS_TIMEOUT_SECONDS)
+        except requests.exceptions.Timeout as e:
+            logging.exception(e)
+            return None, {}
+
         if response.status_code != requests.codes.ok:
             logging.error("LeaseJob error: %s", response.status_code)
             return None, {}
@@ -294,10 +324,14 @@ class VtiEndpointClient(object):
 
         thread = threading.currentThread()
         while getattr(thread, 'keep_running', True):
-            response = requests.post(url, data=json.dumps(self._job),
-                                     headers=self._headers)
-            if response.status_code != requests.codes.ok:
-                logging.error("UpdateLeasedJobStatus error: %s", response)
+            try:
+                response = requests.post(url, data=json.dumps(self._job),
+                                         headers=self._headers,
+                                         timeout=REQUESTS_TIMEOUT_SECONDS)
+                if response.status_code != requests.codes.ok:
+                    logging.error("UpdateLeasedJobStatus error: %s", response)
+            except requests.exceptions.Timeout as e:
+                logging.exception(e)
             time.sleep(update_interval)
 
     def StartHeartbeat(self, status="leased", update_interval=60):
@@ -334,10 +368,14 @@ class VtiEndpointClient(object):
         self.SetJobStatusFromLeasedTo(status)
         self._job["infra_log_url"] = infra_log_url
 
-        response = requests.post(
-            url, data=json.dumps(self._job), headers=self._headers)
-        if response.status_code != requests.codes.ok:
-            logging.error("StopHeartbeat error: %s", response)
+        try:
+            response = requests.post(url, data=json.dumps(self._job),
+                                     headers=self._headers,
+                                     timeout=REQUESTS_TIMEOUT_SECONDS)
+            if response.status_code != requests.codes.ok:
+                logging.error("StopHeartbeat error: %s", response)
+        except requests.exceptions.Timeout as e:
+            logging.exception(e)
 
         self._job = None
 
@@ -365,8 +403,10 @@ class VtiEndpointClient(object):
 
         try:
             response = requests.post(url, data=json.dumps(host),
-                                    headers=self._headers)
-        except requests.exceptions.ConnectionError as e:
+                                     headers=self._headers,
+                                     timeout=REQUESTS_TIMEOUT_SECONDS)
+        except (requests.exceptions.ConnectionError,
+                requests.exceptions.Timeout) as e:
             logging.exception(e)
             return
         if response.status_code != requests.codes.ok: