From cff5cb27c04f4db8d115493c2e2704c6a10726df Mon Sep 17 00:00:00 2001 From: Felipe Padula Sanches Date: Wed, 30 Oct 2024 12:15:24 +1300 Subject: [PATCH] Add support for ulimit flag (#291) * Barebones extension * Add expected format to help message * Add tests for ulimit extension --------- Co-authored-by: Tully Foote --- setup.py | 6 +- src/rocker/ulimit_extension.py | 67 ++++++++++++++++++++++ test/test_ulimit.py | 100 +++++++++++++++++++++++++++++++++ 3 files changed, 170 insertions(+), 3 deletions(-) create mode 100644 src/rocker/ulimit_extension.py create mode 100644 test/test_ulimit.py diff --git a/setup.py b/setup.py index 4d66983..2cc1d9d 100644 --- a/setup.py +++ b/setup.py @@ -43,7 +43,7 @@ 'console_scripts': [ 'rocker = rocker.cli:main', 'detect_docker_image_os = rocker.cli:detect_image_os', - ], + ], 'rocker.extensions': [ 'cuda = rocker.nvidia_extension:Cuda', 'devices = rocker.extensions:Devices', @@ -63,11 +63,12 @@ 'pulse = rocker.extensions:PulseAudio', 'rmw = rocker.rmw_extension:RMW', 'ssh = rocker.ssh_extension:Ssh', + 'ulimit = rocker.ulimit_extension:Ulimit', 'user = rocker.extensions:User', 'volume = rocker.volume_extension:Volume', 'x11 = rocker.nvidia_extension:X11', ] - }, + }, 'author': 'Tully Foote', 'author_email': 'tfoote@osrfoundation.org', 'keywords': ['Docker'], @@ -91,4 +92,3 @@ } setup(**kwargs) - diff --git a/src/rocker/ulimit_extension.py b/src/rocker/ulimit_extension.py new file mode 100644 index 0000000..105650a --- /dev/null +++ b/src/rocker/ulimit_extension.py @@ -0,0 +1,67 @@ +# Copyright 2019 Open Source Robotics Foundation + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from argparse import ArgumentTypeError +import re +from rocker.extensions import RockerExtension, name_to_argument + + +class Ulimit(RockerExtension): + """ + A RockerExtension to handle ulimit settings for Docker containers. + + This extension allows specifying ulimit options in the format TYPE=SOFT_LIMIT[:HARD_LIMIT] + and validates the format before passing them as Docker arguments. + """ + EXPECTED_FORMAT = "TYPE=SOFT_LIMIT[:HARD_LIMIT]" + + @staticmethod + def get_name(): + return 'ulimit' + + def get_docker_args(self, cliargs): + args = [''] + ulimits = [x for sublist in cliargs[Ulimit.get_name()] for x in sublist] + for ulimit in ulimits: + if self.arg_format_is_valid(ulimit): + args.append(f"--ulimit {ulimit}") + else: + raise ArgumentTypeError( + f"Error processing {Ulimit.get_name()} flag '{ulimit}': expected format" + f" {Ulimit.EXPECTED_FORMAT}") + return ' '.join(args) + + def arg_format_is_valid(self, arg: str): + """ + Validate the format of the ulimit argument. + + Args: + arg (str): The ulimit argument to validate. + + Returns: + bool: True if the format is valid, False otherwise. + """ + ulimit_format = r'(\w+)=(\w+)(:\w+)?$' + match = re.match(ulimit_format, arg) + return match is not None + + @staticmethod + def register_arguments(parser, defaults): + parser.add_argument(name_to_argument(Ulimit.get_name()), + type=str, + nargs='+', + action='append', + metavar=Ulimit.EXPECTED_FORMAT, + default=defaults.get(Ulimit.get_name(), None), + help='ulimit options to add into the container.') diff --git a/test/test_ulimit.py b/test/test_ulimit.py new file mode 100644 index 0000000..4db8a6d --- /dev/null +++ b/test/test_ulimit.py @@ -0,0 +1,100 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import unittest +from argparse import ArgumentTypeError + +from rocker.ulimit_extension import Ulimit + + +class UlimitTest(unittest.TestCase): + """Unit tests for the Ulimit class.""" + + def setUp(self): + self._instance = Ulimit() + + def _is_arg_translation_ok(self, mock_cliargs, expected): + is_ok = False + message_string = "" + try: + docker_args = self._instance.get_docker_args( + {self._instance.get_name(): [mock_cliargs]}) + is_ok = docker_args == expected + message_string = f"Expected: '{expected}', got: '{docker_args}'" + except ArgumentTypeError: + message_string = "Incorrect argument format" + return (is_ok, message_string) + + def test_args_single_soft(self): + """Test single soft limit argument.""" + mock_cliargs = ["rtprio=99"] + expected = " --ulimit rtprio=99" + self.assertTrue(*self._is_arg_translation_ok(mock_cliargs, expected)) + + def test_args_multiple_soft(self): + """Test multiple soft limit arguments.""" + mock_cliargs = ["rtprio=99", "memlock=102400"] + expected = " --ulimit rtprio=99 --ulimit memlock=102400" + self.assertTrue(*self._is_arg_translation_ok(mock_cliargs, expected)) + + def test_args_single_hard(self): + """Test single hard limit argument.""" + mock_cliargs = ["nofile=1024:524288"] + expected = " --ulimit nofile=1024:524288" + self.assertTrue(*self._is_arg_translation_ok(mock_cliargs, expected)) + + def test_args_multiple_hard(self): + """Test multiple hard limit arguments.""" + mock_cliargs = ["nofile=1024:524288", "rtprio=90:99"] + expected = " --ulimit nofile=1024:524288 --ulimit rtprio=90:99" + self.assertTrue(*self._is_arg_translation_ok(mock_cliargs, expected)) + + def test_args_multiple_mix(self): + """Test multiple mixed limit arguments.""" + mock_cliargs = ["rtprio=99", "memlock=102400", "nofile=1024:524288"] + expected = " --ulimit rtprio=99 --ulimit memlock=102400 --ulimit nofile=1024:524288" + self.assertTrue(*self._is_arg_translation_ok(mock_cliargs, expected)) + + def test_args_wrong_single_soft(self): + """Test if single soft limit argument is wrong.""" + mock_cliargs = ["rtprio99"] + expected = " --ulimit rtprio99" + self.assertFalse(*self._is_arg_translation_ok(mock_cliargs, expected)) + + def test_args_wrong_multiple_soft(self): + """Test if multiple soft limit arguments are wrong.""" + mock_cliargs = ["rtprio=99", "memlock102400"] + expected = " --ulimit rtprio=99 --ulimit memlock=102400" + self.assertFalse(*self._is_arg_translation_ok(mock_cliargs, expected)) + + def test_args_wrong_single_hard(self): + """Test if single hard limit arguments are wrong.""" + mock_cliargs = ["nofile=1024:524288:"] + expected = " --ulimit nofile=1024:524288" + self.assertFalse(*self._is_arg_translation_ok(mock_cliargs, expected)) + + def test_args_wrong_multiple_hard(self): + """Test if multiple hard limit arguments are wrong.""" + mock_cliargs = ["nofile1024524288", "rtprio=90:99"] + expected = " --ulimit nofile=1024:524288 --ulimit rtprio=90:99" + self.assertFalse(*self._is_arg_translation_ok(mock_cliargs, expected)) + + def test_args_wrong_multiple_mix(self): + """Test if multiple mixed limit arguments are wrong.""" + mock_cliargs = ["rtprio=:", "memlock102400", "nofile1024:524288:"] + expected = " --ulimit rtprio=99 --ulimit memlock=102400 --ulimit nofile=1024:524288" + self.assertFalse(*self._is_arg_translation_ok(mock_cliargs, expected))