Add ulimit to Docker clients

localstack · joe4dev · Mar 20, 2023 · Mar 10, 2023 · Mar 14, 2023 · Mar 14, 2023
commit ed4ce7c52fd36e075e83e72464d24d96627da58c
diff --git a/localstack/utils/container_utils/container_client.py b/localstack/utils/container_utils/container_client.py
@@ -11,7 +11,7 @@
 from abc import ABCMeta, abstractmethod
 from enum import Enum, unique
 from pathlib import Path
-from typing import Dict, List, NamedTuple, Optional, Tuple, Union
+from typing import Dict, List, NamedTuple, Optional, Tuple, TypedDict, Union
 
 from localstack.utils.no_exit_argument_parser import BooleanOptionalAction, NoExitArgumentParser
 
@@ -105,6 +105,23 @@ class DockerPlatform(str):
     linux_arm64 = "linux/arm64"
 
 
+class Ulimit(TypedDict, total=False):
+    """The ``ulimit`` settings for the container.
+    See https://www.tutorialspoint.com/setting-ulimit-values-on-docker-containers
+    """
+
+    name: str
+    soft_limit: int
+    hard_limit: Optional[int] = None
+
+    def __repr__(self):
+        """Format: <type>=<soft limit>[:<hard limit>]"""
+        ulimit_string = f"{self.name}={self.soft_limit}"
+        if self.hard_limit:
+            ulimit_string += f":{self.hard_limit}"
+        return ulimit_string
+
+
 # defines the type for port mappings (source->target port range)
 PortRange = Union[List, HashableList]
 
@@ -376,6 +393,7 @@ class ContainerConfiguration:
     dns: Optional[str] = None
     workdir: Optional[str] = None
     platform: Optional[str] = None
+    ulimits: Optional[List[Ulimit]] = None
 
 
 @dataclasses.dataclass
@@ -393,6 +411,7 @@ class DockerRunFlags:
     platform: Optional[DockerPlatform]
     privileged: Optional[bool]
     ports: Optional[PortMappings]
+    ulimits: Optional[List[Ulimit]]
     user: Optional[str]
 
 
@@ -771,6 +790,7 @@ def run_container(
         additional_flags: Optional[str] = None,
         workdir: Optional[str] = None,
         privileged: Optional[bool] = None,
+        ulimits: Optional[List[Ulimit]] = None,
     ) -> Tuple[bytes, bytes]:
         """Creates and runs a given docker container
 
@@ -917,6 +937,7 @@ def parse_additional_flags(
         ports: Optional[PortMappings] = None,
         privileged: Optional[bool] = None,
         user: Optional[str] = None,
+        ulimits: Optional[List[Ulimit]] = None,
     ) -> DockerRunFlags:
         """Parses additional CLI-formatted Docker flags, which could overwrite provided defaults.
         :param additional_flags: String which contains the flag definitions inspired by the Docker CLI reference:
@@ -928,6 +949,7 @@ def parse_additional_flags(
         :param platform: Platform to execute container. Warning will be printed if platform is overwritten in flags.
         :param ports: PortMapping object. Will be modified in place.
         :param privileged: Run the container in privileged mode. Warning will be printed if overwritten in flags.
+        :param ulimits: ulimit options in the format <type>=<soft limit>[:<hard limit>]
         :param user: User to run first process. Warning will be printed if user is overwritten in flags.
         :return: A DockerRunFlags object that will return new objects if respective parameters were None and
                 additional flags contained a flag for that object or the same which are passed otherwise.
@@ -956,6 +978,7 @@ def parse_additional_flags(
             "--privileged",
             type=bool,
             help="Give extended privileges to this container",
+            # TODO: could replace with 'store_true'
             action=BooleanOptionalAction,
         )
         parser.add_argument(
@@ -965,6 +988,9 @@ def parse_additional_flags(
             dest="publish_ports",
             action="append",
         )
+        parser.add_argument(
+            "--ulimit", help="Container ulimit settings", dest="ulimits", action="append"
+        )
         parser.add_argument("--user", "-u", help="Username or UID to execute first process")
         parser.add_argument(
             "--volume", "-v", help="Bind mount a volume", dest="volumes", action="append"
@@ -1043,6 +1069,19 @@ def parse_additional_flags(
                 ports = ports if ports is not None else PortMappings()
                 ports.add(host_port, int(container_port), protocol)
 
+        if args.ulimits:
+            ulimits = ulimits if ulimits is not None else []
+            ulimits_dict = {ul["name"]: ul for ul in ulimits}
+            for ulimit in args.ulimits:
+                name, _, rhs = ulimit.partition("=")
+                soft, _, hard = rhs.partition(":")
+                hard_limit = int(hard) if hard else int(soft)
+                new_ulimit = Ulimit(name=name, soft_limit=int(soft), hard_limit=hard_limit)
+                if ulimits_dict.get(name):
+                    LOG.warning(f"Overwriting Docker ulimit {new_ulimit}")
+                ulimits_dict[name] = new_ulimit
+            ulimits = list(ulimits_dict.values())
+
         if args.user:
             LOG.warning(
                 "Overwriting Docker user '%s' with new value '%s'",
@@ -1077,6 +1116,7 @@ def parse_additional_flags(
             network=network,
             platform=platform,
             privileged=privileged,
+            ulimits=ulimits,
             user=user,
         )
 

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,6 +517,7 @@ 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
@@ -528,6 +530,7 @@ def create_container(
                 platform=platform,
                 privileged=privileged,
                 ports=ports,
+                ulimits=ulimits,
                 user=user,
             )
             env_vars = parsed_flags.env_vars
@@ -538,6 +541,7 @@ def create_container(
             platform = parsed_flags.platform
             privileged = parsed_flags.privileged
             ports = parsed_flags.ports
+            ulimits = parsed_flags.ulimits
             user = parsed_flags.user
 
         try:
@@ -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)