Refactor server-side sleep into parent class

Author: dsotirho-ucsc

src/azul/chalice.py

@@ -15,6 +15,7 @@
 import mimetypes
 import os
 import pathlib
+import time
 from typing import (
     Any,
     Iterator,
@@ -617,3 +618,27 @@ def lambda_context(self) -> LambdaContext:
     @property
     def current_request(self) -> AzulRequest:
         return self.app.current_request
+
+    def server_side_sleep(self, max_seconds: float) -> float:
+        """
+        Sleep in the current Lambda function.
+
+        :param max_seconds: The requested amount of time to sleep in seconds.
+                            The actual time slept will be less if the requested
+                            amount would cause the Lambda function to exceed its
+                            execution timeout.
+
+        :return: The actual amount of time slept in seconds
+        """
+        assert isinstance(max_seconds, float), max_seconds
+        assert 0 <= max_seconds <= config.api_gateway_lambda_timeout, max_seconds
+        remaining_time = self.lambda_context.get_remaining_time_in_millis() / 1000
+        # A small buffer is subtracted from the Lambda's remaining time to
+        # ensure that it wakes up before it runs out of execution time (and
+        # before API Gateway times out) so it gets a chance to return a response
+        # to the client.
+        actual_seconds = min(max_seconds,
+                             remaining_time - config.api_gateway_timeout_padding - 3)
+        log.debug('Sleeping for %.3f seconds', actual_seconds)
+        time.sleep(actual_seconds)
+        return actual_seconds

src/azul/service/drs_controller.py

@@ -12,7 +12,6 @@
 from datetime import (
     datetime,
 )
-import time
 from typing import (
     Optional,
 )
@@ -184,10 +183,7 @@ def _dos_gs_url(self, file_uuid, version) -> mutable_furl:
                     return url
                 elif dss_response.status_code == 301:
                     url = dss_response.next.url
-                    remaining_lambda_seconds = self.lambda_context.get_remaining_time_in_millis() / 1000
-                    server_side_sleep = min(1,
-                                            max(remaining_lambda_seconds - config.api_gateway_timeout_padding - 3, 0))
-                    time.sleep(server_side_sleep)
+                    self.server_side_sleep(1.0)
                 else:
                     raise ChaliceViewError({
                         'msg': f'Received {dss_response.status_code} from DSS. Could not get file'

src/azul/service/repository_controller.py

@@ -4,7 +4,6 @@
 )
 import json
 import logging
-import time
 from typing import (
     Optional,
     TYPE_CHECKING,
@@ -21,7 +20,6 @@
     CatalogName,
     cache,
     cached_property,
-    config,
     reject,
 )
 from azul.auth import (
@@ -264,14 +262,8 @@ def download_file(self,
                 if wait == '0':
                     pass
                 elif wait == '1':
-                    # Sleep in the lambda but ensure that we wake up before it
-                    # runs out of execution time (and before API Gateway times
-                    # out) so we get a chance to return a response to the client
-                    remaining_time = self.lambda_context.get_remaining_time_in_millis() / 1000
-                    server_side_sleep = min(float(retry_after),
-                                            remaining_time - config.api_gateway_timeout_padding - 3)
-                    time.sleep(server_side_sleep)
-                    retry_after = round(retry_after - server_side_sleep)
+                    time_slept = self.server_side_sleep(float(retry_after))
+                    retry_after = round(retry_after - time_slept)
                 else:
                     assert False, wait
                 query_params['wait'] = wait