localstack · joe4dev · Mar 20, 2023 · Mar 10, 2023 · Mar 14, 2023 · Mar 14, 2023
diff --git a/localstack/utils/container_utils/container_client.py b/localstack/utils/container_utils/container_client.py
diff --git a/localstack/utils/container_utils/docker_cmd_client.py b/localstack/utils/container_utils/docker_cmd_client.py
@@ -22,6 +22,7 @@
     PortMappings,
     RegistryConnectionError,
     SimpleVolumeBind,
+    Ulimit,
     Util,
     VolumeBind,
 )
@@ -629,6 +630,7 @@ def _build_run_create_cmd(
         privileged: Optional[bool] = None,
         labels: Optional[Dict[str, str]] = None,
         platform: Optional[DockerPlatform] = None,
+        ulimits: Optional[List[Ulimit]] = None,
     ) -> Tuple[List[str], str]:
         env_file = None
         cmd = self._docker_cmd() + [action]
@@ -678,6 +680,10 @@ def _build_run_create_cmd(
                 cmd += ["--label", f"{key}={value}"]
         if platform:
             cmd += ["--platform", platform]
+        if ulimits:
+            cmd += list(
+                itertools.chain.from_iterable(["--ulimits", str(ulimit)] for ulimit in ulimits)
+            )
 
         if additional_flags:
             cmd += shlex.split(additional_flags)

diff --git a/localstack/utils/container_utils/docker_sdk_client.py b/localstack/utils/container_utils/docker_sdk_client.py
@@ -24,6 +24,7 @@
     PortMappings,
     RegistryConnectionError,
     SimpleVolumeBind,
+    Ulimit,
     Util,
 )
 from localstack.utils.strings import to_bytes, to_str
@@ -516,29 +517,32 @@ def create_container(
         privileged: Optional[bool] = None,
         labels: Optional[Dict[str, str]] = None,
         platform: Optional[DockerPlatform] = None,
+        ulimits: Optional[List[Ulimit]] = None,
     ) -> str:
         LOG.debug("Creating container with attributes: %s", locals())
         extra_hosts = None
         if additional_flags:
             parsed_flags = Util.parse_additional_flags(
                 additional_flags,
-                env_vars,
-                ports,
-                mount_volumes,
-                network,
-                user,
-                platform,
-                privileged,
+                env_vars=env_vars,
+                mounts=mount_volumes,
+                network=network,
+                platform=platform,
+                privileged=privileged,
+                ports=ports,
+                ulimits=ulimits,
+                user=user,
             )
             env_vars = parsed_flags.env_vars
-            ports = parsed_flags.ports
-            mount_volumes = parsed_flags.mounts
             extra_hosts = parsed_flags.extra_hosts
-            network = parsed_flags.network
+            mount_volumes = parsed_flags.mounts
             labels = parsed_flags.labels
-            user = parsed_flags.user
+            network = parsed_flags.network
             platform = parsed_flags.platform
             privileged = parsed_flags.privileged
+            ports = parsed_flags.ports
+            ulimits = parsed_flags.ulimits
+            user = parsed_flags.user
 
         try:
             kwargs = {}
@@ -558,6 +562,13 @@ def create_container(
                 kwargs["privileged"] = True
             if labels:
                 kwargs["labels"] = labels
+            if ulimits:
+                kwargs["ulimits"] = [
+                    docker.types.Ulimit(
+                        name=ulimit.name, soft=ulimit.soft_limit, hard=ulimit.hard_limit
+                    )
+                    for ulimit in ulimits
+                ]
             mounts = None
             if mount_volumes:
                 mounts = Util.convert_mount_list_to_dict(mount_volumes)
@@ -615,11 +626,21 @@ def run_container(
         dns: Optional[str] = None,
         additional_flags: Optional[str] = None,
         workdir: Optional[str] = None,
+        platform: Optional[DockerPlatform] = None,
         privileged: Optional[bool] = None,
+        ulimits: Optional[List[Ulimit]] = None,
     ) -> Tuple[bytes, bytes]:
         LOG.debug("Running container with image: %s", image_name)
         container = None
         try:
+            kwargs = {}
+            if ulimits:
+                kwargs["ulimits"] = [
+                    docker.types.Ulimit(
+                        name=ulimit.name, soft=ulimit.soft_limit, hard=ulimit.hard_limit
+                    )
+                    for ulimit in ulimits
+                ]
             container = self.create_container(
                 image_name,
                 name=name,
@@ -641,6 +662,8 @@ def run_container(
                 additional_flags=additional_flags,
                 workdir=workdir,
                 privileged=privileged,
+                platform=platform,
+                **kwargs,
             )
             result = self.start_container(
                 container_name_or_id=container,

diff --git a/localstack/utils/no_exit_argument_parser.py b/localstack/utils/no_exit_argument_parser.py
@@ -0,0 +1,19 @@
+import argparse
+import logging
+from typing import NoReturn, Optional
+
+LOG = logging.getLogger(__name__)
+
+
+class NoExitArgumentParser(argparse.ArgumentParser):
+    """Implements the `exit_on_error=False` behavior introduced in Python 3.9 to support older Python versions
+    and prevents further SystemExit for other error categories.
+    * Limitations of error categories: https://stackoverflow.com/a/67891066/6875981
+    * ArgumentParser subclassing example: https://stackoverflow.com/a/59072378/6875981
+    """
+
+    def exit(self, status: int = ..., message: Optional[str] = ...) -> NoReturn:
+        LOG.warning(f"Error in argument parser but preventing exit: {message}")
+
+    def error(self, message: str) -> NoReturn:
+        raise NotImplementedError(f"Unsupported flag by this Docker client: {message}")
diff --git a/tests/integration/awslambda/functions/lambda_ulimits.py b/tests/integration/awslambda/functions/lambda_ulimits.py
@@ -0,0 +1,18 @@
+import resource
+
+
+def handler(event, context):
+    # https://docs.python.org/3/library/resource.html
+    ulimit_names = {
+        "RLIMIT_AS": resource.RLIMIT_AS,
+        "RLIMIT_CORE": resource.RLIMIT_CORE,
+        "RLIMIT_CPU": resource.RLIMIT_CPU,
+        "RLIMIT_DATA": resource.RLIMIT_DATA,
+        "RLIMIT_FSIZE": resource.RLIMIT_FSIZE,
+        "RLIMIT_MEMLOCK": resource.RLIMIT_MEMLOCK,
+        "RLIMIT_NOFILE": resource.RLIMIT_NOFILE,
+        "RLIMIT_NPROC": resource.RLIMIT_NPROC,
+        "RLIMIT_RSS": resource.RLIMIT_RSS,
+        "RLIMIT_STACK": resource.RLIMIT_STACK,
+    }
+    return {label: resource.getrlimit(res) for label, res in ulimit_names.items()}
diff --git a/tests/integration/awslambda/test_lambda.py b/tests/integration/awslambda/test_lambda.py
@@ -83,6 +83,7 @@
 TEST_LAMBDA_TIMEOUT_ENV_PYTHON = os.path.join(THIS_FOLDER, "functions/lambda_timeout_env.py")
 TEST_LAMBDA_SLEEP_ENVIRONMENT = os.path.join(THIS_FOLDER, "functions/lambda_sleep_environment.py")
 TEST_LAMBDA_INTROSPECT_PYTHON = os.path.join(THIS_FOLDER, "functions/lambda_introspect.py")
+TEST_LAMBDA_ULIMITS = os.path.join(THIS_FOLDER, "functions/lambda_ulimits.py")
 TEST_LAMBDA_VERSION = os.path.join(THIS_FOLDER, "functions/lambda_version.py")
 
 TEST_GOLANG_LAMBDA_URL_TEMPLATE = "https://github.com/localstack/awslamba-go-runtime/releases/download/v{version}/example-handler-{os}-{arch}.tar.gz"
@@ -379,6 +380,31 @@ def test_runtime_introspection_arm(self, lambda_client, create_lambda_function,
         invoke_result = lambda_client.invoke(FunctionName=func_name)
         snapshot.match("invoke_runtime_arm_introspection", invoke_result)
 
+    @pytest.mark.skipif(
+        is_old_local_executor(),
+        reason="Monkey-patching of Docker flags is not applicable because no new container is spawned",
+    )
+    @pytest.mark.skip_snapshot_verify(condition=is_old_provider, paths=["$..LogResult"])
+    @pytest.mark.aws_validated
+    def test_runtime_ulimits(self, lambda_client, create_lambda_function, snapshot, monkeypatch):
+        """We consider ulimits parity as opt-in because development environments could hit these limits unlike in
+        optimized production deployments."""
+        monkeypatch.setattr(
+            config,
+            "LAMBDA_DOCKER_FLAGS",
+            "--ulimit nofile=1024:1024 --ulimit nproc=735:735 --ulimit core=-1:-1 --ulimit stack=8388608:-1",
+        )
+
+        func_name = f"test_lambda_ulimits_{short_uid()}"
+        create_lambda_function(
+            func_name=func_name,
+            handler_file=TEST_LAMBDA_ULIMITS,
+            runtime=Runtime.python3_9,
+        )
+
+        invoke_result = lambda_client.invoke(FunctionName=func_name)
+        snapshot.match("invoke_runtime_ulimits", invoke_result)
+
     @pytest.mark.skipif(is_old_provider(), reason="unsupported in old provider")
     @pytest.mark.skipif(
         is_old_local_executor(),

diff --git a/tests/integration/awslambda/test_lambda.snapshot.json b/tests/integration/awslambda/test_lambda.snapshot.json
@@ -2658,5 +2658,60 @@
         }
       }
     }
+  },
+  "tests/integration/awslambda/test_lambda.py::TestLambdaBehavior::test_runtime_ulimits": {
+    "recorded-date": "15-03-2023, 00:11:15",
+    "recorded-content": {
+      "invoke_runtime_ulimits": {
+        "ExecutedVersion": "$LATEST",
+        "Payload": {
+          "RLIMIT_AS": [
+            -1,
+            -1
+          ],
+          "RLIMIT_CORE": [
+            -1,
+            -1
+          ],
+          "RLIMIT_CPU": [
+            -1,
+            -1
+          ],
+          "RLIMIT_DATA": [
+            -1,
+            -1
+          ],
+          "RLIMIT_FSIZE": [
+            -1,
+            -1
+          ],
+          "RLIMIT_MEMLOCK": [
+            65536,
+            65536
+          ],
+          "RLIMIT_NOFILE": [
+            1024,
+            1024
+          ],
+          "RLIMIT_NPROC": [
+            735,
+            735
+          ],
+          "RLIMIT_RSS": [
+            -1,
+            -1
+          ],
+          "RLIMIT_STACK": [
+            8388608,
+            -1
+          ]
+        },
+        "StatusCode": 200,
+        "ResponseMetadata": {
+          "HTTPHeaders": {},
+          "HTTPStatusCode": 200
+        }
+      }
+    }
   }
 }