Restart manifest generations when manifest has expired (#6441)

src/azul/service/async_manifest_service.py

@@ -3,6 +3,7 @@
 import logging
 from typing import (
     Self,
+    TypedDict,
 )
 from uuid import (
     UUID,
@@ -35,15 +36,30 @@ class InvalidTokenError(Exception):
 @attrs.frozen(kw_only=True)
 class Token:
     """
-    Represents an ongoing manifest generation
+    Represents a Step Function execution to generate a manifest
     """
+    #: A hash of the inputs
     generation_id: UUID = strict_auto()
+
+    #: Number of prior executions for the generation represented by this token
+    iteration: int = strict_auto()
+
+    #: The number of times the service received a request to inspect the
+    #: status of the execution represented by this token
     request_index: int = strict_auto()
+
+    #: How long clients should wait before requesting a status update about the
+    #: execution represented by the token
     retry_after: int = strict_auto()
 
+    @property
+    def execution_id(self) -> tuple[UUID, int]:
+        return self.generation_id, self.iteration
+
     def pack(self) -> bytes:
         return msgpack.packb([
             self.generation_id.bytes,
+            self.iteration,
             self.request_index,
             self.retry_after
         ])
@@ -52,6 +68,7 @@ def pack(self) -> bytes:
     def unpack(cls, pack: bytes) -> Self:
         i = iter(msgpack.unpackb(pack))
         return cls(generation_id=UUID(bytes=next(i)),
+                   iteration=next(i),
                    request_index=next(i),
                    retry_after=next(i))
 
@@ -66,8 +83,9 @@ def decode(cls, token: str) -> Self:
             raise InvalidTokenError(token) from e
 
     @classmethod
-    def first(cls, generation_id: UUID) -> Self:
+    def first(cls, generation_id: UUID, iteration: int) -> Self:
         return cls(generation_id=generation_id,
+                   iteration=iteration,
                    request_index=0,
                    retry_after=cls._next_retry_after(0))
 
@@ -87,11 +105,21 @@ def _next_retry_after(cls, request_index: int) -> int:
             return delays[-1]
 
 
+class ExecutionResult(TypedDict):
+    input: JSON
+    output: JSON
+
+
 @attrs.frozen
 class NoSuchGeneration(Exception):
     token: Token = strict_auto()
 
 
+@attrs.frozen
+class GenerationFinished(Exception):
+    token: Token = strict_auto()
+
+
 @attrs.frozen(kw_only=True)
 class GenerationFailed(Exception):
     status: str = strict_auto()
@@ -112,14 +140,18 @@ class AsyncManifestService:
     def machine_name(self):
         return config.qualified_resource_name(config.manifest_sfn)
 
-    def start_generation(self, generation_id: UUID, input: JSON) -> Token:
-        execution_name = self.execution_name(generation_id)
+    def start_generation(self,
+                         generation_id: UUID,
+                         input: JSON,
+                         iteration: int
+                         ) -> Token:
+        execution_name = self.execution_name(generation_id, iteration)
         execution_arn = self.execution_arn(execution_name)
         # The input contains the verbatim manifest key as JSON while the ARN
         # contains the encoded hash of the key so this log line is useful for
         # associating the hash with the key for diagnostic purposes.
         log.info('Starting execution %r for input %r', execution_arn, input)
-        token = Token.first(generation_id)
+        token = Token.first(generation_id, iteration)
         try:
             # If there already is an execution of the given name, and if that
             # execution is still ongoing and was given the same input as what we
@@ -145,16 +177,16 @@ def start_generation(self, generation_id: UUID, input: JSON) -> Token:
             execution = self._sfn.describe_execution(executionArn=execution_arn)
             if input == json.loads(execution['input']):
                 log.info('A completed execution %r already exists', execution_arn)
-                return token
+                raise GenerationFinished(token)
             else:
                 raise InvalidGeneration(token)
         else:
-            assert execution_arn == execution['executionArn'], execution
+            assert execution_arn == execution['executionArn'], (execution_arn, execution)
             log.info('Started execution %r or it was already running', execution_arn)
             return token
 
-    def inspect_generation(self, token: Token) -> Token | JSON:
-        execution_name = self.execution_name(token.generation_id)
+    def inspect_generation(self, token: Token) -> Token | ExecutionResult:
+        execution_name = self.execution_name(token.generation_id, token.iteration)
         execution_arn = self.execution_arn(execution_name)
         try:
             execution = self._sfn.describe_execution(executionArn=execution_arn)
@@ -169,7 +201,7 @@ def inspect_generation(self, token: Token) -> Token | JSON:
                     return token.next(retry_after=1)
                 else:
                     log.info('Execution %r succeeded with output %r', execution_arn, output)
-                    return json.loads(output)
+                    return {k: json.loads(execution[k]) for k in ['input', 'output']}
             elif status == 'RUNNING':
                 log.info('Execution %r is still running', execution_arn)
                 return token.next()
@@ -186,10 +218,11 @@ def machine_arn(self):
     def execution_arn(self, execution_name):
         return self.arn(f'execution:{self.machine_name}:{execution_name}')
 
-    def execution_name(self, generation_id: UUID) -> str:
+    def execution_name(self, generation_id: UUID, iteration: int) -> str:
         assert isinstance(generation_id, UUID), generation_id
-        execution_name = str(generation_id)
-        assert 0 < len(execution_name) <= 80, (generation_id, execution_name)
+        assert isinstance(iteration, int), iteration
+        execution_name = f'{generation_id}_{iteration}'
+        assert 0 < len(execution_name) <= 80, execution_name
         return execution_name
 
     @property

src/azul/service/manifest_controller.py

@@ -37,6 +37,7 @@
 from azul.service.async_manifest_service import (
     AsyncManifestService,
     GenerationFailed,
+    GenerationFinished,
     InvalidTokenError,
     NoSuchGeneration,
     Token,
@@ -127,6 +128,33 @@ def _unpack_token_or_key(self,
                     # The OpenAPI spec doesn't distinguish key and token
                     raise BadRequestError('Invalid token')
 
+    def _start_execution(self,
+                         filters: Filters,
+                         manifest_key: ManifestKey,
+                         previous_token: Token | None = None,
+                         ) -> Token:
+        partition = ManifestPartition.first()
+        state: ManifestGenerationState = {
+            'filters': filters.to_json(),
+            'manifest_key': manifest_key.to_json(),
+            'partition': partition.to_json()
+        }
+        # Manifest keys for catalogs with long names would be too long to be
+        # used directly as state machine execution names.
+        generation_id = manifest_key.uuid
+        # ManifestGenerationState is also JSON but there is no way to express
+        # that since TypedDict rejects a co-parent class.
+        input = cast(JSON, state)
+        next_iteration = 0 if previous_token is None else previous_token.iteration + 1
+        for i in range(10):
+            try:
+                return self.async_service.start_generation(generation_id,
+                                                           input,
+                                                           iteration=next_iteration + i)
+            except GenerationFinished:
+                pass
+        raise ChaliceViewError('Too many executions of this manifest generation')
+
     def get_manifest_async(self,
                            *,
                            token_or_key: str,
@@ -152,19 +180,8 @@ def get_manifest_async(self,
                     # A cache miss, but the exception tells us the cache key
                     manifest, manifest_key = None, e.manifest_key
                     # Prepare the execution that will generate the manifest
-                    partition = ManifestPartition.first()
-                    state: ManifestGenerationState = {
-                        'filters': filters.to_json(),
-                        'manifest_key': manifest_key.to_json(),
-                        'partition': partition.to_json()
-                    }
-                    # Manifest keys for catalogs with long names would be too
-                    # long to be used directly in state machine execution names.
-                    generation_id = manifest_key.uuid
-                    # ManifestGenerationState is also JSON but there is no way
-                    # to express that since TypedDict rejects a co-parent class.
-                    input = cast(JSON, state)
-                    token = self.async_service.start_generation(generation_id, input)
+                    token = self._start_execution(filters=filters,
+                                                  manifest_key=manifest_key)
                 else:
                     # A cache hit
                     manifest_key = manifest.manifest_key
@@ -180,31 +197,48 @@ def get_manifest_async(self,
                     manifest_key = self.service.verify_manifest_key(manifest_key)
                     manifest = self.service.get_cached_manifest_with_key(manifest_key)
                 except CachedManifestNotFound:
-                    # We can't restart the execution to regenerate the manifest
-                    # because we don't know the parameters that were used to
-                    # create it. So the client will have to do it.
-                    raise GoneError('The requested manifest has expired, please request a new one')
+                    # We could start another execution but that would require
+                    # the client to follow more redirects. We've already sent
+                    # the final 302 so we shouldn't that.
+                    raise GoneError('The manifest has expired, please request a new one')
                 except InvalidManifestKeySignature:
                     raise BadRequestError('Invalid token')
         else:
             # A token for an ongoing execution was given
             assert manifest_key is None, manifest_key
             try:
-                token_or_state = self.async_service.inspect_generation(token)
+                token_or_result = self.async_service.inspect_generation(token)
             except NoSuchGeneration:
                 raise BadRequestError('Invalid token')
             except GenerationFailed as e:
                 raise ChaliceViewError('Failed to generate manifest', e.status, e.output)
-            if isinstance(token_or_state, Token):
+            if isinstance(token_or_result, Token):
                 # Execution is still ongoing, we got an updated token
-                token, manifest, manifest_key = token_or_state, None, None
-            elif isinstance(token_or_state, dict):
-                # Eecution is done, a cached manifest should be ready
-                state = token_or_state
-                manifest = Manifest.from_json(state['manifest'])
+                token, manifest, manifest_key = token_or_result, None, None
+            elif isinstance(token_or_result, dict):
+                # The execution is done, the resulting manifest should be ready
+                result = token_or_result
+                manifest = Manifest.from_json(result['output']['manifest'])
                 manifest_key = manifest.manifest_key
+                try:
+                    manifest = self.service.get_cached_manifest_with_key(manifest_key)
+                except CachedManifestNotFound as e:
+                    assert manifest_key == e.manifest_key
+                    # There are two possible causes for the missing manifest: it
+                    # may have expired, in which case the supplied token must be
+                    # really stale, or the manifest was deleted immediately
+                    # after it was created. We haven't sent a 302 redirect yet,
+                    # so we'll just restart the generation by starting another
+                    # execution for it.
+                    manifest = None
+                    filters = Filters.from_json(result['input']['filters'])
+                    token = self._start_execution(filters=filters,
+                                                  manifest_key=manifest_key,
+                                                  previous_token=token)
+                else:
+                    assert manifest_key == manifest.manifest_key
             else:
-                assert False, token_or_state
+                assert False, token_or_result
 
         body: dict[str, int | str | FlatJSON]
 

test/integration_test.py

@@ -109,6 +109,9 @@
 from azul.csp import (
     CSP,
 )
+from azul.deployment import (
+    aws,
+)
 from azul.drs import (
     AccessMethod,
 )
@@ -532,9 +535,10 @@ def _test_manifest(self, catalog: CatalogName):
         assert supported_formats
         for format in [None, *supported_formats]:
             filters = self._manifest_filters(catalog)
-            first_fetch = bool(self.random.getrandbits(1))
-            for fetch in [first_fetch, not first_fetch]:
-                with self.subTest('manifest', catalog=catalog, format=format, fetch=fetch):
+            execution_ids = set()
+            coin_flip = bool(self.random.getrandbits(1))
+            for i, fetch in enumerate([coin_flip, coin_flip, not coin_flip]):
+                with self.subTest('manifest', catalog=catalog, format=format, i=i, fetch=fetch):
                     args = dict(catalog=catalog, filters=json.dumps(filters))
                     if format is None:
                         format = first(supported_formats)
@@ -565,13 +569,22 @@ def worker(_):
                             results = list(tpe.map(worker, range(num_workers)))
 
                     self.assertEqual([None] * num_workers, results)
-                    execution_ids = self._manifest_execution_ids(responses, fetch=fetch)
-                    # The second iteration of the inner-most loop re-requests
-                    # the manifest with only `fetch` being different. In that
-                    # case, the manifest will already be cached and no step
-                    # function execution is expected to have been started.
-                    expect_execution = fetch == first_fetch
-                    self.assertEqual(1 if expect_execution else 0, len(execution_ids))
+
+                    execution_ids.update(self._manifest_execution_ids(responses))
+                    bucket, key = one(self._manifest_objects(responses))
+                    if i == 0:
+                        aws.s3.delete_object(Bucket=bucket, Key=key)
+                        # One execution to generate the manifest
+                        self.assertEqual(1, len(execution_ids))
+                    elif i == 1:
+                        # One more execution to re-generate the manifest
+                        self.assertEqual(2, len(execution_ids))
+                    elif i == 2:
+                        # Only fetch mode changed, cached manifest will be used,
+                        # and no additional executions are expectect
+                        self.assertEqual(2, len(execution_ids))
+                    else:
+                        assert False
 
     def _manifest_filters(self, catalog: CatalogName) -> JSON:
         # IT catalogs with just one public source are always indexed completely
@@ -642,28 +655,43 @@ def _test_manifest_tagging_race(self, catalog: CatalogName):
                             self._manifest_validators[format](catalog, response.data)
                             break
 
-                execution_ids = self._manifest_execution_ids(responses, fetch=False)
+                execution_ids = self._manifest_execution_ids(responses)
                 self.assertEqual(1, len(execution_ids))
 
     def _manifest_execution_ids(self,
-                                responses: list[urllib3.HTTPResponse],
-                                *,
-                                fetch: bool
-                                ) -> set[uuid.UUID]:
-        urls: list[furl]
-        if fetch:
-            responses = [
-                json.loads(r.data)
-                for r in responses
-                if r.status == 200 and r.headers['Content-Type'] == 'application/json'
-            ]
-            urls = [furl(r['Location']) for r in responses if r['Status'] == 301]
-        else:
-            urls = [furl(r.headers['Location']) for r in responses if r.status == 301]
+                                responses: list[urllib3.HTTPResponse]
+                                ) -> set[tuple[uuid.UUID, int]]:
+        urls = self._manifest_urls(responses, status=301)
         tokens = {Token.decode(url.path.segments[-1]) for url in urls}
-        execution_ids = {token.generation_id for token in tokens}
+        execution_ids = {token.execution_id for token in tokens}
         return execution_ids
 
+    def _manifest_objects(self,
+                          responses: list[urllib3.HTTPResponse]
+                          ) -> set[tuple[str, str]]:
+        urls = self._manifest_urls(responses, status=302)
+        return {
+            (url.path.segments[0], '/'.join(url.path.segments[1:]))
+            for url in urls
+            if url.netloc == 's3.amazonaws.com' and url.scheme == 'https'
+        }
+
+    def _manifest_urls(self,
+                       responses: list[urllib3.HTTPResponse],
+                       *,
+                       status: int
+                       ) -> list[furl]:
+        urls: list[furl] = []
+        for response in responses:
+            if response.status == 200:
+                if response.headers['Content-Type'] == 'application/json':
+                    body = json.loads(response.data)
+                    if body['Status'] == status:
+                        urls.append(furl(body['Location']))
+            elif response.status == status:
+                urls.append(furl(response.headers['Location']))
+        return urls
+
     def _get_one_inner_file(self, catalog: CatalogName) -> tuple[JSON, FileInnerEntity]:
         outer_file = self._get_one_outer_file(catalog)
         inner_files: JSONs = outer_file['files']
@@ -1065,8 +1093,8 @@ def _check_curl_manifest(self, _catalog: CatalogName, response: bytes):
         num_files = 0
         for url, output in grouper(lines, 2):
             num_files += 1
-            self.assertTrue(url.startswith('url='))
-            self.assertTrue(output.startswith('output='))
+            self.assertTrue(url.startswith('url='), url)
+            self.assertTrue(output.startswith('output='), output)
         log.info(f'Manifest contains {num_files} files.')
         self.assertGreater(num_files, 0)