Implement Lambda function versions (#5927)

Author: dsotirho-ucsc

scripts/sell_unused_slots.py

@@ -71,7 +71,7 @@ def _list_contribution_lambda_functions(cls) -> list[Lambda]:
         """
         return [
             lambda_
-            for lambda_ in Lambdas().list_lambdas()
+            for lambda_ in Lambdas().list_lambdas(deployment='ALL')
             if lambda_.is_contribution_lambda
         ]
 

src/azul/lambdas.py

@@ -13,6 +13,7 @@
 )
 
 from azul import (
+    JSON,
     R,
     cache,
     config,
@@ -37,6 +38,7 @@ class Lambda:
     name: str
     role: str
     slot_location: Optional[str]
+    version: str
 
     @property
     def is_contribution_lambda(self) -> bool:
@@ -82,13 +84,15 @@ def has_notification_queue(handler) -> bool:
     def from_response(cls, response: 'FunctionConfigurationTypeDef') -> Self:
         name = response['FunctionName']
         role = response['Role']
+        version = response['Version']
         try:
             slot_location = response['Environment']['Variables']['AZUL_TDR_SOURCE_LOCATION']
         except KeyError:
             slot_location = None
         return cls(name=name,
                    role=role,
-                   slot_location=slot_location)
+                   slot_location=slot_location,
+                   version=version)
 
     def __attrs_post_init__(self):
         if self.slot_location is None:
@@ -105,25 +109,65 @@ class Lambdas:
     def _lambda(self):
         return aws.lambda_
 
-    def list_lambdas(self) -> list[Lambda]:
+    def list_lambdas(self,
+                     deployment: str
+                     ) -> list[Lambda]:
+        """
+        Return a list of AWS Lambda functions. Only the largest numbered version
+        of each function will be included in the list.
+
+        :param deployment: Limit output to the specified deployment stage. If
+                           'ALL', functions from all deployments will be
+                           returned.
+        """
+        paginator = self._lambda.get_paginator('list_functions')
+        lambda_prefixes = None if deployment == 'ALL' else [
+            config.qualified_resource_name(lambda_name, stage=deployment)
+            for lambda_name in config.lambda_names()
+        ]
+        functions: dict[str, JSON] = dict()
+        for response in paginator.paginate(FunctionVersion='ALL'):
+            for function in response['Functions']:
+                version = function['Version']
+                version = None if version == '$LATEST' else int(version)
+                if version and (lambda_prefixes is None or any(
+                    function['FunctionName'].startswith(prefix)
+                    for prefix in lambda_prefixes
+                )):
+                    name = function['FunctionName']
+                    previous_function = functions.get(name)
+                    if previous_function is None or version > int(previous_function['Version']):
+                        functions[name] = function
         return [
             Lambda.from_response(function)
-            for response in self._lambda.get_paginator('list_functions').paginate()
-            for function in response['Functions']
+            for function in functions.values()
         ]
 
+    def get_function(self, function_name: str) -> JSON:
+        """
+        Return the Lambda client `get_function()` response for the largest
+        numbered version of the Lambda function.
+        """
+        paginator = self._lambda.get_paginator('list_versions_by_function')
+        params = {'FunctionName': function_name}
+        version = max([
+            int(function['Version'])
+            for response in paginator.paginate(**params)
+            for function in response['Versions']
+            if function['Version'] != '$LATEST'
+        ])
+        return self._lambda.get_function(FunctionName=function_name,
+                                         Qualifier=str(version))
+
     def manage_lambdas(self, enabled: bool):
-        paginator = self._lambda.get_paginator('list_functions')
-        lambda_prefixes = [config.qualified_resource_name(lambda_infix) for lambda_infix in config.lambda_names()]
-        assert all(lambda_prefixes)
-        for lambda_page in paginator.paginate(FunctionVersion='ALL', MaxItems=500):
-            for lambda_name in [metadata['FunctionName'] for metadata in lambda_page['Functions']]:
-                if any(lambda_name.startswith(prefix) for prefix in lambda_prefixes):
-                    self.manage_lambda(lambda_name, enabled)
+        for function in self.list_lambdas(deployment=config.deployment_stage):
+            self.manage_lambda(function.name, enabled)
 
     def manage_lambda(self, lambda_name: str, enable: bool):
-        lambda_settings = self._lambda.get_function(FunctionName=lambda_name)
+        lambda_settings = self.get_function(function_name=lambda_name)
         lambda_arn = lambda_settings['Configuration']['FunctionArn']
+        # Lambda does not support adding tags to function aliases or versions
+        lambda_arn, _, _ = lambda_arn.rpartition(':')
         lambda_tags = self._lambda.list_tags(Resource=lambda_arn)['Tags']
         lambda_name = lambda_settings['Configuration']['FunctionName']
         if enable:
@@ -132,13 +176,15 @@ def manage_lambda(self, lambda_name: str, enable: bool):
 
                 if original_concurrency_limit is not None:
                     log.info(f'Setting concurrency limit for {lambda_name} back to {original_concurrency_limit}.')
+                    # Concurrency settings apply to the function as a whole,
+                    # including all published versions and the unpublished
+                    # version
                     self._lambda.put_function_concurrency(FunctionName=lambda_name,
                                                           ReservedConcurrentExecutions=original_concurrency_limit)
                 else:
                     log.info(f'Removed concurrency limit for {lambda_name}.')
                     self._lambda.delete_function_concurrency(FunctionName=lambda_name)
 
-                lambda_arn = lambda_settings['Configuration']['FunctionArn']
                 self._lambda.untag_resource(Resource=lambda_arn, TagKeys=[self.tag_name])
             else:
                 log.warning(f'{lambda_name} is already enabled.')
@@ -156,7 +202,7 @@ def manage_lambda(self, lambda_name: str, enable: bool):
 
                 log.info(f'Setting concurrency limit for {lambda_name} to zero.')
                 new_tag = {self.tag_name: repr(concurrency_limit)}
-                self._lambda.tag_resource(Resource=lambda_settings['Configuration']['FunctionArn'], Tags=new_tag)
+                self._lambda.tag_resource(Resource=lambda_arn, Tags=new_tag)
                 self._lambda.put_function_concurrency(FunctionName=lambda_name, ReservedConcurrentExecutions=0)
             else:
                 log.warning(f'{lambda_name} is already disabled.')
@@ -165,7 +211,7 @@ def reset_lambda_roles(self):
         client = self._lambda
         lambda_names = set(config.lambda_names())
 
-        for lambda_ in self.list_lambdas():
+        for lambda_ in self.list_lambdas(deployment='ALL'):
             for lambda_name in lambda_names:
                 if lambda_.name.startswith(config.qualified_resource_name(lambda_name)):
                     other_lambda_name = one(lambda_names - {lambda_name})
@@ -174,6 +220,8 @@ def reset_lambda_roles(self):
                         config.qualified_resource_name(other_lambda_name)
                     )
                     log.info('Temporarily updating %r to role %r', lambda_.name, temporary_role)
+                    # You can’t modify the configuration of a published version,
+                    # only the unpublished version
                     client.update_function_configuration(FunctionName=lambda_.name,
                                                          Role=temporary_role)
                     log.info('Updating %r to role %r', lambda_.name, lambda_.role)

src/azul/queues.py

@@ -502,9 +502,14 @@ def _wait_for_queue_empty(self, queue: 'Queue'):
             time.sleep(3)
             queue.reload()
 
-    def _manage_sqs_push(self, function_name: str, queue: 'Queue', enable: bool):
+    def _manage_sqs_push(self,
+                         function_name: str,
+                         function_version: str,
+                         queue: 'Queue',
+                         enable: bool):
         lambda_ = aws.lambda_
-        response = lambda_.list_event_source_mappings(FunctionName=function_name,
+        partial_arn = f'{function_name}:{function_version}'
+        response = lambda_.list_event_source_mappings(FunctionName=partial_arn,
                                                       EventSourceArn=queue.attributes['QueueArn'])
         mapping_uuid = one(response['EventSourceMappings'])['UUID']
 
@@ -554,6 +559,11 @@ def manage_lambdas(self, queues: Mapping[str, 'Queue'], enable: bool):
         Enable or disable the readers and writers of the given queues.
         """
         functions_by_queue = self.functions_by_queue()
+        versions_by_function = {
+            f.name: f.version
+            for f in self._lambdas.list_lambdas(deployment=config.deployment_stage)
+        }
+        assert set(functions_by_queue.values()) <= set(versions_by_function)
 
         with ThreadPoolExecutor(max_workers=len(queues)) as tpe:
             futures = []
@@ -567,10 +577,11 @@ def submit(f, *args, **kwargs):
                 except KeyError:
                     assert queue_name in config.fail_queue_names
                 else:
+                    version = versions_by_function[function]
                     if queue_name == config.notifications_queue.name:
                         # Prevent new notifications from being added
                         submit(self._manage_lambda, config.indexer_name, enable)
-                    submit(self._manage_sqs_push, function, queue, enable)
+                    submit(self._manage_sqs_push, function, version, queue, enable)
             self._handle_futures(futures)
             futures = [tpe.submit(self._wait_for_queue_idle, queue) for queue in queues.values()]
             self._handle_futures(futures)

src/azul/terraform.py

@@ -14,6 +14,7 @@
 from pathlib import (
     Path,
 )
+import re
 import subprocess
 from typing import (
     Mapping,
@@ -759,6 +760,10 @@ def tf_config(self, app_name):
         for _, resource in functions:
             assert 'layers' not in resource
             resource['layers'] = ['${aws_lambda_layer_version.dependencies.arn}']
+            # Publishing a new Lambda function version each time lets us perform
+            # an atomic update of the lambda function, avoiding a race condition
+            # between the update of the function's configuration and its code.
+            resource['publish'] = True
             env = config.es_endpoint_env(
                 es_endpoint=(
                     aws.es_endpoint
@@ -921,6 +926,38 @@ def tf_config(self, app_name):
                 sqs_name, _ = config.unqualified_resource_name(resource_name, suffix)
                 resource['event_source_arn'] = f'${{aws_sqs_queue.{sqs_name}.arn}}'
 
+        # Replace the references to unqualified Lambda function ARNs emitted by
+        # Chalice with qualified ARNs. Note: `aws_lambda_permission` resources
+        # doesn't support the qualified ARN syntax, instead it has a `qualifier`
+        # argument.
+        #
+        locals[app_name] = re.sub(r'(aws_lambda_function\.[^\.\s]+)\.invoke_arn',
+                                  r'\1.qualified_invoke_arn',
+                                  json_str(locals[app_name]))
+        if app_name == 'indexer':
+            resource_arguments = [
+                ('aws_cloudwatch_event_target', 'arn'),
+                ('aws_lambda_event_source_mapping', 'function_name'),
+            ]
+        elif app_name == 'service':
+            resource_arguments = [
+                ('aws_cloudwatch_event_target', 'arn'),
+            ]
+        else:
+            assert False, app_name
+        for resource_type, argument in resource_arguments:
+            for _, resource in json_item_dicts(resources[resource_type]):
+                value = json_str(resource[argument])
+                assert value.endswith('.arn}'), resource
+                resource[argument] = value.replace('.arn', '.qualified_arn')
+
+        # Add a qualifier argument to `aws_lambda_permission` resources
+        #
+        for _, resource in json_item_dicts(resources['aws_lambda_permission']):
+            assert 'qualifier' not in resource, resource
+            lambda_arn = resource['function_name']
+            resource['qualifier'] = lambda_arn.replace('.arn', '.version')
+
         return {
             'resource': resources,
             'data': data,