Lambda DevX: Bypass Concurrency Settings (#11418)

MEPalma · web-flow · commit 13d85d190a05 · 2024-10-14T09:52:12.000+02:00
diff --git a/localstack-core/localstack/services/lambda_/invocation/assignment.py b/localstack-core/localstack/services/lambda_/invocation/assignment.py
@@ -15,6 +15,7 @@
     InitializationType,
     OtherServiceEndpoint,
 )
+from localstack.utils.lambda_debug_mode.lambda_debug_mode import is_lambda_debug_enabled_for
 
 LOG = logging.getLogger(__name__)
 
@@ -134,6 +135,21 @@ def scale_provisioned_concurrency(
         function_version: FunctionVersion,
         target_provisioned_environments: int,
     ) -> list[Future[None]]:
+        # Enforce a single environment per lambda version if this is a target
+        # of an active Lambda Debug Mode.
+        qualified_lambda_version_arn = function_version.qualified_arn
+        if (
+            is_lambda_debug_enabled_for(qualified_lambda_version_arn)
+            and target_provisioned_environments > 0
+        ):
+            LOG.warning(
+                "Environments for '%s' enforced to '1' by Lambda Debug Mode, "
+                "configurations will continue to report the set value '%s'",
+                qualified_lambda_version_arn,
+                target_provisioned_environments,
+            )
+            target_provisioned_environments = 1
+
         current_provisioned_environments = [
             e
             for e in self.environments[version_manager_id].values()
diff --git a/localstack-core/localstack/services/lambda_/invocation/counting_service.py b/localstack-core/localstack/services/lambda_/invocation/counting_service.py
@@ -11,6 +11,9 @@
     InitializationType,
 )
 from localstack.services.lambda_.invocation.models import lambda_stores
+from localstack.utils.lambda_debug_mode.lambda_debug_mode import (
+    is_lambda_debug_enabled_for,
+)
 
 LOG = logging.getLogger(__name__)
 
@@ -125,6 +128,24 @@ def get_invocation_lease(
         unqualified_function_arn = function_version.id.unqualified_arn()
         qualified_arn = function_version.id.qualified_arn()
 
+        # Enforce one lease per ARN if the global flag is set
+        if is_lambda_debug_enabled_for(qualified_arn):
+            with provisioned_tracker.lock, on_demand_tracker.lock:
+                on_demand_executions: int = on_demand_tracker.concurrent_executions[
+                    unqualified_function_arn
+                ]
+                provisioned_executions = provisioned_tracker.concurrent_executions[qualified_arn]
+                if on_demand_executions or provisioned_executions:
+                    LOG.warning(
+                        "Concurrent lambda invocations disabled for '%s' by Lambda Debug Mode",
+                        qualified_arn,
+                    )
+                    raise TooManyRequestsException(
+                        "Rate Exceeded.",
+                        Reason="SingleLeaseEnforcement",
+                        Type="User",
+                    )
+
         lease_type = None
         with provisioned_tracker.lock:
             # 1) Check for free provisioned concurrency
diff --git a/localstack-core/localstack/utils/lambda_debug_mode/lambda_debug_mode.py b/localstack-core/localstack/utils/lambda_debug_mode/lambda_debug_mode.py
@@ -1,6 +1,7 @@
 from typing import Optional
 
 from localstack.aws.api.lambda_ import Arn
+from localstack.utils.lambda_debug_mode.lambda_debug_mode_config import LambdaDebugModeConfig
 from localstack.utils.lambda_debug_mode.lambda_debug_mode_session import LambdaDebugModeSession
 
 # Specifies the fault timeout value in seconds to be used by time restricted workflows when
@@ -13,19 +14,28 @@ def is_lambda_debug_mode() -> bool:
     return LambdaDebugModeSession.get().is_lambda_debug_mode()
 
 
-def lambda_debug_port_for(lambda_arn: Arn) -> Optional[int]:
+def _lambda_debug_config_for(lambda_arn: Arn) -> Optional[LambdaDebugModeConfig]:
     if not is_lambda_debug_mode():
         return None
     debug_configuration = LambdaDebugModeSession.get().debug_config_for(lambda_arn=lambda_arn)
+    return debug_configuration
+
+
+def is_lambda_debug_enabled_for(lambda_arn: Arn) -> bool:
+    """Returns True if the given lambda arn is subject of an active debugging configuration; False otherwise."""
+    debug_configuration = _lambda_debug_config_for(lambda_arn=lambda_arn)
+    return debug_configuration is not None
+
+
+def lambda_debug_port_for(lambda_arn: Arn) -> Optional[int]:
+    debug_configuration = _lambda_debug_config_for(lambda_arn=lambda_arn)
     if debug_configuration is None:
         return None
     return debug_configuration.debug_port
 
 
 def is_lambda_debug_timeout_enabled_for(lambda_arn: Arn) -> bool:
-    if not is_lambda_debug_mode():
-        return False
-    debug_configuration = LambdaDebugModeSession.get().debug_config_for(lambda_arn=lambda_arn)
+    debug_configuration = _lambda_debug_config_for(lambda_arn=lambda_arn)
     if debug_configuration is None:
         return False
     return not debug_configuration.enforce_timeouts
