Merge pull request #852 from iocage/issues/74235

tkt-74235: Add Cpuset support to Iocage

iocage_lib/ioc_json.py

@@ -162,6 +162,85 @@ def __write_file(self):
             )
 
 
+class IOCCpuset(object):
+
+    def __init__(self, name):
+        self.jail_name = f'ioc-{name}'
+
+    def set_cpuset(self, value=None):
+        if not value:
+            value = 'all'
+
+        failed = False
+        try:
+            iocage_lib.ioc_exec.SilentExec(
+                ['cpuset', '-l', value, '-j', self.jail_name],
+                None, unjailed=True, decode=True
+            )
+        except iocage_lib.ioc_exceptions.CommandFailed:
+            failed = True
+        finally:
+            return failed
+
+    @staticmethod
+    def retrieve_cpu_sets():
+        cpu_sets = -2
+        try:
+            output = iocage_lib.ioc_exec.SilentExec(
+                ['cpuset', '-g', '-s', '0'],
+                None, unjailed=True, decode=True
+            )
+        except iocage_lib.ioc_exceptions.CommandFailed:
+            pass
+        else:
+            result = re.findall(
+                r'.*mask:.*(\d+)$',
+                output.stdout.split('\n')[0]
+            )
+            if result:
+                cpu_sets = int(result[0])
+        finally:
+            return cpu_sets
+
+    @staticmethod
+    def validate_cpuset_prop(value, raise_error=True):
+        failed = False
+        cpu_sets = IOCCpuset.retrieve_cpu_sets() + 1
+
+        if not any(
+            cond for cond in (
+                re.findall(
+                    fr'^(?!.*(\b\d+\b).*\b\1\b)'
+                    fr'((?:{"|".join(map(str, range(cpu_sets)))})'
+                    fr'(,(?:{"|".join(map(str, range(cpu_sets)))}))*)?$',
+                    value
+                ),
+                value in ('off', 'all'),
+                re.findall(
+                    fr'^(?:{"|".join(map(str, range(cpu_sets - 1)))})-'
+                    fr'(?:{"|".join(map(str, range(cpu_sets)))})$',
+                    value
+                ) and int(value.split('-')[0]) < int(value.split('-')[1])
+            )
+        ):
+            failed = True
+
+        if failed and raise_error:
+            iocage_lib.ioc_common.logit(
+                {
+                    'level': 'EXCEPTION',
+                    'message': 'Please specify a valid format for cpuset '
+                               'value.\nFollowing 4 formats are supported:\n'
+                               '1) comma delimited string i.e 0,1,2,3\n'
+                               '2) a range of values i.e 0-2\n'
+                               '3) "all" - all would mean using all cores\n'
+                               '4) off'
+                }
+            )
+        else:
+            return failed
+
+
 class IOCRCTL(object):
 
     types = {
@@ -1933,6 +2012,17 @@ def json_set_value(self, prop, _import=False, default=False):
                 else:
                     key = key.replace("_", ".")
 
+                if key == 'cpuset':
+                    iocage_lib.ioc_common.logit(
+                        {
+                            'level': 'INFO',
+                            'message': 'cpuset changes '
+                                       'require a jail restart'
+                        },
+                        _callback=self.callback,
+                        silent=self.silent
+                    )
+
                 # Let's set a rctl rule for the prop if applicable
                 if key in IOCRCTL.types:
                     rctl_jail = IOCRCTL(conf['host_hostuuid'])
@@ -2098,7 +2188,7 @@ def json_check_prop(self, key, value, conf):
             "allow_vmm": truth_variations,
             "vnet_interfaces": ("string", ),
             # RCTL limits
-            "cpuset": ("off", "on"),
+            "cpuset": ('string',),
             "rlimits": ("off", "on"),
             "memoryuse": ('string',),
             "memorylocked": ('string',),
@@ -2311,6 +2401,8 @@ def json_check_prop(self, key, value, conf):
                         )
                 elif key in IOCRCTL.types:
                     IOCRCTL.validate_rctl_props(key, value)
+                elif key == 'cpuset':
+                    IOCCpuset.validate_cpuset_prop(value)
 
                 return value, conf
             else:

iocage_lib/ioc_start.py

@@ -656,6 +656,12 @@ def __start_jail__(self):
                 _callback=self.callback,
                 silent=self.silent)
 
+        self.set(
+            "last_started={}".format(
+                datetime.datetime.utcnow().strftime("%F %T")
+            )
+        )
+
         rctl_keys = set(
             filter(
                 lambda k: self.conf.get(k, 'off') != 'off',
@@ -683,14 +689,28 @@ def __start_jail__(self):
 
             if failed:
                 iocage_lib.ioc_common.logit({
-                    'level': 'INFO',
+                    'level': 'ERROR',
                     'message': f'  + Failed to set {", ".join(failed)} '
                     'RCTL props'
                 })
 
-        self.set(
-            "last_started={}".format(datetime.datetime.utcnow().strftime(
-                "%F %T")))
+        cpuset = self.conf.get('cpuset', 'off')
+        if cpuset != 'off':
+            # Let's set the specified rules
+            iocage_lib.ioc_common.logit({
+                'level': 'INFO',
+                'message': f'  + Setting cpuset to: {cpuset}'
+            })
+
+            cpuset_jail = iocage_lib.ioc_json.IOCCpuset(self.uuid)
+            cpuset_jail.validate_cpuset_prop(cpuset)
+
+            failed = cpuset_jail.set_cpuset(cpuset)
+            if failed:
+                iocage_lib.ioc_common.logit({
+                    'level': 'ERROR',
+                    'message': f'  + Failed to set cpuset to: {cpuset}'
+                })
 
     def check_aliases(self, ip_addrs, mode='4'):
         """

tests/conftest.py

@@ -322,3 +322,15 @@ def _default_jails(resources, **kwargs):
         ]
 
     return _default_jails
+
+
+@pytest.fixture
+def run_console():
+    def _run_console(cmd):
+        proc = subprocess.run(
+            cmd, stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE
+        )
+        return proc
+
+    return _run_console

tests/data_classes.py

@@ -608,6 +608,16 @@ def is_thickjail(self):
             'properties'
         ].get('origin', {}).get('value')
 
+    @property
+    def cpuset(self):
+        output = self.run_command(['cpuset', '-g'])[0].split('\n')[0]
+        return list(
+            map(
+                lambda v: int(v.strip()),
+                output.split(':')[1].strip().split(',')
+            )
+        )
+
     @property
     def is_rcjail(self):
         return self.config.get('boot', 0)

tests/functional_tests/0006_set_test.py

@@ -23,6 +23,7 @@
 # POSSIBILITY OF SUCH DAMAGE.
 
 import pytest
+import re
 
 
 require_root = pytest.mark.require_root
@@ -37,13 +38,13 @@
 # TODO: Plugin test left
 
 
-def _set_and_test_note_prop(invoke_cli, value, jail):
+def _set_and_test_prop(invoke_cli, value, jail, prop='notes'):
     invoke_cli(
-        ['set', f'notes={value}', jail.name]
+        ['set', f'{prop}={value}', jail.name]
     )
 
-    assert jail.config.get('notes') == value, \
-        f'Failed to set note value to {value}'
+    assert jail.config.get(prop) == value, \
+        f'Failed to set {prop} value to {value}'
 
 
 @require_root
@@ -52,7 +53,7 @@ def test_01_set_prop_on_jail(resource_selector, invoke_cli, skip_test):
     jails = resource_selector.jails
     skip_test(not jails)
 
-    _set_and_test_note_prop(
+    _set_and_test_prop(
         invoke_cli, 'foo \"bar\"', jails[0]
     )
 
@@ -65,7 +66,7 @@ def test_02_set_prop_on_thickconfig_jail(
     thickconfig_jails = resource_selector.thickconfig_jails
     skip_test(not thickconfig_jails)
 
-    _set_and_test_note_prop(
+    _set_and_test_prop(
         invoke_cli, 'foo \"bar\"', thickconfig_jails[0]
     )
 
@@ -76,7 +77,7 @@ def test_03_set_prop_on_basejail(resource_selector, invoke_cli, skip_test):
     basejails = resource_selector.basejails
     skip_test(not basejails)
 
-    _set_and_test_note_prop(
+    _set_and_test_prop(
         invoke_cli, 'foo \"bar\"', basejails[0]
     )
 
@@ -87,6 +88,74 @@ def test_04_set_prop_on_template_jail(resource_selector, invoke_cli, skip_test):
     template_jails = resource_selector.template_jails
     skip_test(not template_jails)
 
-    _set_and_test_note_prop(
+    _set_and_test_prop(
         invoke_cli, 'foo \"bar\"', template_jails[0]
     )
+
+
+@require_root
+@require_zpool
+def test_04_set_cpuset_prop_on_jail(
+    resource_selector, invoke_cli, skip_test, run_console
+):
+    jails = resource_selector.startable_jails_and_not_running
+    skip_test(not jails)
+
+    # We need to get the no of cpus first
+    cpuset_output = run_console(['cpuset', '-g', '-s', '0'])
+    skip_test(cpuset_output.returncode)
+
+    cpuset_num = re.findall(
+        r'.*mask:.*(\d+)$',
+        cpuset_output.stdout.decode().split('\n')[0]
+    )
+    skip_test(not cpuset_num)
+
+    cpuset_num = int(cpuset_num[0])
+
+    # We would like to test following formats now
+    # 0,1,2,3
+    # 0-2
+    # all
+    # off
+    #
+    # However if the num of cpu is only one, we can't do ranges or multiple
+    # values in that case
+
+    possible_variations = ['off', 'all']
+    if cpuset_num:
+        possible_variations.extend([
+            f'0-{cpuset_num}',
+            f','.join(
+                map(
+                    str,
+                    range(cpuset_num + 1)
+                )
+            )
+        ])
+
+    jail = jails[0]
+    for variation in possible_variations:
+        _set_and_test_prop(
+            invoke_cli, variation, jail, 'cpuset'
+        )
+
+    if cpuset_num:
+        invoke_cli(
+            ['start', jail.name],
+            f'Jail {jail} failed to start'
+        )
+
+        assert jail.running is True
+
+        jail_cpusets = jail.cpuset
+        assert set(jail_cpusets) == set(
+            map(int, possible_variations[-1].split(','))
+        )
+
+        invoke_cli(
+            ['stop', jail.name],
+            f'Jail {jail} failed to stop'
+        )
+
+        assert jail.running is False