localstack · dominikschubert · Jun 13, 2023 · Jun 14, 2023 · Jun 14, 2023 · Jun 14, 2023
diff --git a/localstack/services/lambda_/event_source_listeners/adapters.py b/localstack/services/lambda_/event_source_listeners/adapters.py
@@ -3,7 +3,6 @@
 import logging
 import threading
 from abc import ABC
-from concurrent.futures import Future
 from functools import lru_cache
 from typing import Callable, Optional
 
@@ -13,7 +12,7 @@
 from localstack.aws.protocol.serializer import gen_amzn_requestid
 from localstack.services.lambda_ import api_utils
 from localstack.services.lambda_.api_utils import function_locators_from_arn, qualifier_is_version
-from localstack.services.lambda_.invocation.lambda_models import InvocationError, InvocationResult
+from localstack.services.lambda_.invocation.lambda_models import InvocationResult
 from localstack.services.lambda_.invocation.lambda_service import LambdaService
 from localstack.services.lambda_.invocation.models import lambda_stores
 from localstack.services.lambda_.lambda_executors import (
@@ -23,6 +22,7 @@
 from localstack.utils.aws.client_types import ServicePrincipal
 from localstack.utils.json import BytesEncoder
 from localstack.utils.strings import to_bytes, to_str
+from localstack.utils.threads import FuncThread
 
 LOG = logging.getLogger(__name__)
 
@@ -143,29 +143,26 @@ def __init__(self, lambda_service: LambdaService):
         self.lambda_service = lambda_service
 
     def invoke(self, function_arn, context, payload, invocation_type, callback=None):
+        def _invoke(*args, **kwargs):
+            # split ARN ( a bit unnecessary since we build an ARN again in the service)
+            fn_parts = api_utils.FULL_FN_ARN_PATTERN.search(function_arn).groupdict()
 
-        # split ARN ( a bit unnecessary since we build an ARN again in the service)
-        fn_parts = api_utils.FULL_FN_ARN_PATTERN.search(function_arn).groupdict()
-
-        ft = self.lambda_service.invoke(
-            # basically function ARN
-            function_name=fn_parts["function_name"],
-            qualifier=fn_parts["qualifier"],
-            region=fn_parts["region_name"],
-            account_id=fn_parts["account_id"],
-            invocation_type=invocation_type,
-            client_context=json.dumps(context or {}),
-            payload=to_bytes(json.dumps(payload or {}, cls=BytesEncoder)),
-            request_id=gen_amzn_requestid(),
-        )
-
-        if callback:
+            result = self.lambda_service.invoke(
+                # basically function ARN
+                function_name=fn_parts["function_name"],
+                qualifier=fn_parts["qualifier"],
+                region=fn_parts["region_name"],
+                account_id=fn_parts["account_id"],
+                invocation_type=invocation_type,
+                client_context=json.dumps(context or {}),
+                payload=to_bytes(json.dumps(payload or {}, cls=BytesEncoder)),
+                request_id=gen_amzn_requestid(),
+            )
 
-            def mapped_callback(ft_result: Future[InvocationResult]) -> None:
+            if callback:
                 try:
-                    result = ft_result.result(timeout=10)
                     error = None
-                    if isinstance(result, InvocationError):
+                    if result.is_error:
                         error = "?"
                     callback(
                         result=LegacyInvocationResult(
@@ -187,7 +184,8 @@ def mapped_callback(ft_result: Future[InvocationResult]) -> None:
                         error=e,
                     )
 
-            ft.add_done_callback(mapped_callback)
+        thread = FuncThread(_invoke)
+        thread.start()
 
     def invoke_with_statuscode(
         self,
@@ -204,7 +202,7 @@ def invoke_with_statuscode(
         fn_parts = api_utils.FULL_FN_ARN_PATTERN.search(function_arn).groupdict()
 
         try:
-            ft = self.lambda_service.invoke(
+            result = self.lambda_service.invoke(
                 # basically function ARN
                 function_name=fn_parts["function_name"],
                 qualifier=fn_parts["qualifier"],
@@ -218,11 +216,10 @@ def invoke_with_statuscode(
 
             if callback:
 
-                def mapped_callback(ft_result: Future[InvocationResult]) -> None:
+                def mapped_callback(result: InvocationResult) -> None:
                     try:
-                        result = ft_result.result(timeout=10)
                         error = None
-                        if isinstance(result, InvocationError):
+                        if result.is_error:
                             error = "?"
                         callback(
                             result=LegacyInvocationResult(
@@ -243,11 +240,10 @@ def mapped_callback(ft_result: Future[InvocationResult]) -> None:
                             error=e,
                         )
 
-                ft.add_done_callback(mapped_callback)
+                mapped_callback(result)
 
             # they're always synchronous in the ASF provider
-            result = ft.result(timeout=900)
-            if isinstance(result, InvocationError):
+            if result.is_error:
                 return 500
             else:
                 return 200

diff --git a/localstack/services/lambda_/invocation/assignment.py b/localstack/services/lambda_/invocation/assignment.py
@@ -0,0 +1,148 @@
+import contextlib
+import logging
+from collections import defaultdict
+from concurrent.futures import Future, ThreadPoolExecutor
+from typing import ContextManager
+
+from localstack.services.lambda_.invocation.execution_environment import (
+    ExecutionEnvironment,
+    InvalidStatusException,
+)
+from localstack.services.lambda_.invocation.lambda_models import (
+    FunctionVersion,
+    InitializationType,
+    OtherServiceEndpoint,
+)
+
+LOG = logging.getLogger(__name__)
+
+
+class AssignmentException(Exception):
+    pass
+
+
+class AssignmentService(OtherServiceEndpoint):
+    """
+    scope: LocalStack global
+    """
+
+    # function_version (fully qualified function ARN) => runtime_environment_id => runtime_environment
+    environments: dict[str, dict[str, ExecutionEnvironment]]
+
+    # Global pool for spawning and killing provisioned Lambda runtime environments
+    provisioning_pool: ThreadPoolExecutor
+
+    def __init__(self):
+        self.environments = defaultdict(dict)
+        self.provisioning_pool = ThreadPoolExecutor(thread_name_prefix="lambda-provisioning-pool")
+
+    @contextlib.contextmanager
+    def get_environment(
+        self, function_version: FunctionVersion, provisioning_type: InitializationType
+    ) -> ContextManager[ExecutionEnvironment]:
+        version_arn = function_version.qualified_arn
+        applicable_envs = (
+            env
+            for env in self.environments[version_arn].values()
+            if env.initialization_type == provisioning_type
+        )
+        for environment in applicable_envs:
+            try:
+                environment.reserve()
+                execution_environment = environment
+                break
+            except InvalidStatusException:
+                pass
+        else:
+            if provisioning_type == "provisioned-concurrency":
+                raise AssignmentException(
+                    "No provisioned concurrency environment available despite lease."
+                )
+            elif provisioning_type == "on-demand":
+                execution_environment = self.start_environment(function_version)
+                self.environments[version_arn][execution_environment.id] = execution_environment
+                execution_environment.reserve()
+            else:
+                raise ValueError(f"Invalid provisioning type {provisioning_type}")
+
+        try:
+            yield execution_environment
+            execution_environment.release()
+        except InvalidStatusException as invalid_e:
+            LOG.error("Should not happen: %s", invalid_e)
+        except Exception as e:
+            LOG.error("Failed invocation %s", e)
+            self.stop_environment(execution_environment)
+            raise e
+
+    def start_environment(self, function_version: FunctionVersion) -> ExecutionEnvironment:
+        LOG.debug("Starting new environment")
+        execution_environment = ExecutionEnvironment(
+            function_version=function_version,
+            initialization_type="on-demand",
+            on_timeout=self.on_timeout,
+        )
+        try:
+            execution_environment.start()
+        except Exception as e:
+            message = f"Could not start new environment: {e}"
+            LOG.error(message, exc_info=LOG.isEnabledFor(logging.DEBUG))
+            raise AssignmentException(message) from e
+        return execution_environment
+
+    def on_timeout(self, version_arn: str, environment_id: str) -> None:
+        """Callback for deleting environment after function times out"""
+        del self.environments[version_arn][environment_id]
+
+    def stop_environment(self, environment: ExecutionEnvironment) -> None:
+        version_arn = environment.function_version.qualified_arn
+        try:
+            environment.stop()
+            self.environments.get(version_arn).pop(environment.id)
+        except Exception as e:
+            LOG.debug(
+                "Error while stopping environment for lambda %s, environment: %s, error: %s",
+                version_arn,
+                environment.id,
+                e,
+            )
+
+    def stop_environments_for_version(self, function_version: FunctionVersion):
+        # We have to materialize the list before iterating due to concurrency
+        environments_to_stop = list(
+            self.environments.get(function_version.qualified_arn, {}).values()
+        )
+        for env in environments_to_stop:
+            self.stop_environment(env)
+
+    def scale_provisioned_concurrency(
+        self, function_version: FunctionVersion, target_provisioned_environments: int
+    ) -> list[Future[None]]:
+        version_arn = function_version.qualified_arn
+        current_provisioned_environments = [
+            e
+            for e in self.environments[version_arn].values()
+            if e.initialization_type == "provisioned-concurrency"
+        ]
+        # TODO: refine scaling loop to re-use existing environments instead of re-creating all
+        # current_provisioned_environments_count = len(current_provisioned_environments)
+        # diff = target_provisioned_environments - current_provisioned_environments_count
+
+        # TODO: handle case where no provisioned environment is available during scaling
+        # Most simple scaling implementation for now:
+        futures = []
+        # 1) Re-create new target
+        for _ in range(target_provisioned_environments):
+            execution_environment = ExecutionEnvironment(
+                function_version=function_version,
+                initialization_type="provisioned-concurrency",
+                on_timeout=self.on_timeout,
+            )
+            self.environments[version_arn][execution_environment.id] = execution_environment
+            futures.append(self.provisioning_pool.submit(execution_environment.start))
+        # 2) Kill all existing
+        for env in current_provisioned_environments:
+            # TODO: think about concurrent updates while deleting a function
+            futures.append(self.provisioning_pool.submit(self.stop_environment, env))
+
+        return futures