diff --git a/sonic_installer/main.py b/sonic_installer/main.py index 322aac8220e7..79022408e15b 100644 --- a/sonic_installer/main.py +++ b/sonic_installer/main.py @@ -4,6 +4,7 @@ import subprocess import sys import time +import utilities_common.cli as clicommon from urllib.request import urlopen, urlretrieve import click @@ -367,6 +368,102 @@ def migrate_sonic_packages(bootloader, binary_image_version): umount(new_image_mount, raise_exception=False) +class SWAPAllocator(object): + """Context class to allocate SWAP memory.""" + + SWAP_MEM_SIZE = 1024 + DISK_FREESPACE_THRESHOLD = 4 * 1024 + TOTAL_MEM_THRESHOLD = 2048 + AVAILABLE_MEM_THRESHOLD = 1200 + SWAP_FILE_PATH = '/host/swapfile' + KiB_TO_BYTES_FACTOR = 1024 + MiB_TO_BYTES_FACTOR = 1024 * 1024 + + def __init__(self, allocate, swap_mem_size=None, total_mem_threshold=None, available_mem_threshold=None): + """ + Initialize the SWAP memory allocator. + The allocator will try to setup SWAP memory only if all the below conditions are met: + - allocate evaluates to True + - disk has enough space(> DISK_MEM_THRESHOLD) + - either system total memory < total_mem_threshold or system available memory < available_mem_threshold + + @param allocate: True to allocate SWAP memory if necessarry + @param swap_mem_size: the size of SWAP memory to allocate(in MiB) + @param total_mem_threshold: the system totla memory threshold(in MiB) + @param available_mem_threshold: the system available memory threshold(in MiB) + """ + self.allocate = allocate + self.swap_mem_size = SWAPAllocator.SWAP_MEM_SIZE if swap_mem_size is None else swap_mem_size + self.total_mem_threshold = SWAPAllocator.TOTAL_MEM_THRESHOLD if total_mem_threshold is None else total_mem_threshold + self.available_mem_threshold = SWAPAllocator.AVAILABLE_MEM_THRESHOLD if available_mem_threshold is None else available_mem_threshold + self.is_allocated = False + + @staticmethod + def get_disk_freespace(path): + """Return free disk space in bytes.""" + fs_stats = os.statvfs(path) + return fs_stats.f_bsize * fs_stats.f_bavail + + @staticmethod + def read_from_meminfo(): + """Read information from /proc/meminfo.""" + meminfo = {} + with open("/proc/meminfo") as fd: + for line in fd.readlines(): + if line: + fields = line.split() + if len(fields) >= 2 and fields[1].isdigit(): + meminfo[fields[0].rstrip(":")] = int(fields[1]) + return meminfo + + def setup_swapmem(self): + swapfile = SWAPAllocator.SWAP_FILE_PATH + with open(swapfile, 'wb') as fd: + os.posix_fallocate(fd.fileno(), 0, self.swap_mem_size * SWAPAllocator.MiB_TO_BYTES_FACTOR) + os.chmod(swapfile, 0o600) + run_command(f'mkswap {swapfile}; swapon {swapfile}') + + def remove_swapmem(self): + swapfile = SWAPAllocator.SWAP_FILE_PATH + run_command_or_raise(['swapoff', swapfile], raise_exception=False) + try: + os.unlink(swapfile) + finally: + pass + + def __enter__(self): + if self.allocate: + if self.get_disk_freespace('/host') < max(SWAPAllocator.DISK_FREESPACE_THRESHOLD, self.swap_mem_size) * SWAPAllocator.MiB_TO_BYTES_FACTOR: + echo_and_log("Failed to setup SWAP memory due to insufficient disk free space...", LOG_ERR) + return + meminfo = self.read_from_meminfo() + mem_total_in_bytes = meminfo["MemTotal"] * SWAPAllocator.KiB_TO_BYTES_FACTOR + mem_avail_in_bytes = meminfo["MemAvailable"] * SWAPAllocator.KiB_TO_BYTES_FACTOR + if (mem_total_in_bytes < self.total_mem_threshold * SWAPAllocator.MiB_TO_BYTES_FACTOR + or mem_avail_in_bytes < self.available_mem_threshold * SWAPAllocator.MiB_TO_BYTES_FACTOR): + echo_and_log("Setup SWAP memory") + swapfile = SWAPAllocator.SWAP_FILE_PATH + if os.path.exists(swapfile): + self.remove_swapmem() + try: + self.setup_swapmem() + except Exception: + self.remove_swapmem() + raise + self.is_allocated = True + + def __exit__(self, *exc_info): + if self.is_allocated: + self.remove_swapmem() + + +def validate_positive_int(ctx, param, value): + """Callback to validate param passed is a positive integer.""" + if isinstance(value, int) and value > 0: + return value + raise click.BadParameter("Must be a positive integer") + + # Main entrypoint @click.group(cls=AliasedGroup) def sonic_installer(): @@ -389,8 +486,22 @@ def sonic_installer(): help="Do not migrate current configuration to the newly installed image") @click.option('--skip-package-migration', is_flag=True, help="Do not migrate current packages to the newly installed image") +@click.option('--skip-setup-swap', is_flag=True, + help='Skip setup temporary SWAP memory used for installation') +@click.option('--swap-mem-size', default=1024, type=int, show_default='1024 MiB', + help='SWAP memory space size', callback=validate_positive_int, + cls=clicommon.MutuallyExclusiveOption, mutually_exclusive=['skip_setup_swap']) +@click.option('--total-mem-threshold', default=2048, type=int, show_default='2048 MiB', + help='If system total memory is lower than threshold, setup SWAP memory', + cls=clicommon.MutuallyExclusiveOption, mutually_exclusive=['skip_setup_swap'], + callback=validate_positive_int) +@click.option('--available-mem-threshold', default=1200, type=int, show_default='1200 MiB', + help='If system available memory is lower than threhold, setup SWAP memory', + cls=clicommon.MutuallyExclusiveOption, mutually_exclusive=['skip_setup_swap'], + callback=validate_positive_int) @click.argument('url') -def install(url, force, skip_migration=False, skip_package_migration=False): +def install(url, force, skip_migration=False, skip_package_migration=False, + skip_setup_swap=False, swap_mem_size=None, total_mem_threshold=None, available_mem_threshold=None): """ Install image from local binary or URL""" bootloader = get_bootloader() @@ -427,7 +538,8 @@ def install(url, force, skip_migration=False, skip_package_migration=False): raise click.Abort() echo_and_log("Installing image {} and setting it as default...".format(binary_image_version)) - bootloader.install_image(image_path) + with SWAPAllocator(not skip_setup_swap, swap_mem_size, total_mem_threshold, available_mem_threshold): + bootloader.install_image(image_path) # Take a backup of current configuration if skip_migration: echo_and_log("Skipping configuration migration as requested in the command option.") diff --git a/tests/swap_allocator_test.py b/tests/swap_allocator_test.py new file mode 100644 index 000000000000..033b215dd239 --- /dev/null +++ b/tests/swap_allocator_test.py @@ -0,0 +1,252 @@ +import click +import mock +import pytest +import pdb +import subprocess + +from sonic_installer.main import SWAPAllocator + + +class TestSWAPAllocator(object): + + @classmethod + def setup(cls): + print("SETUP") + + def test_read_from_meminfo(self): + proc_meminfo_lines = [ + "MemTotal: 32859496 kB", + "MemFree: 16275512 kB", + "HugePages_Total: 0", + "HugePages_Free: 0", + ] + + read_meminfo_expected_return = { + "MemTotal": 32859496, + "MemFree": 16275512, + "HugePages_Total": 0, + "HugePages_Free": 0 + } + + with mock.patch("builtins.open") as mock_open: + pseudo_fd = mock.MagicMock() + pseudo_fd.readlines = mock.MagicMock(return_value=proc_meminfo_lines) + mock_open.return_value.__enter__.return_value = pseudo_fd + read_meminfo_actual_return = SWAPAllocator.read_from_meminfo() + assert read_meminfo_actual_return == read_meminfo_expected_return + + def test_setup_swapmem(self): + with mock.patch("builtins.open") as mock_open, \ + mock.patch("os.posix_fallocate") as mock_fallocate, \ + mock.patch("os.chmod") as mock_chmod, \ + mock.patch("sonic_installer.main.run_command") as mock_run: + pseudo_fd = mock.MagicMock() + pseudo_fd_fileno = 10 + pseudo_fd.fileno.return_value = pseudo_fd_fileno + mock_open.return_value.__enter__.return_value = pseudo_fd + + swap_mem_size_in_mib = 2048 * 1024 + expected_swap_mem_size_in_bytes = swap_mem_size_in_mib * 1024 * 1024 + expected_swapfile_location = SWAPAllocator.SWAP_FILE_PATH + expected_swapfile_permission = 0o600 + swap_allocator = SWAPAllocator(allocate=True, swap_mem_size=swap_mem_size_in_mib) + swap_allocator.setup_swapmem() + + mock_fallocate.assert_called_once_with(pseudo_fd_fileno, 0, expected_swap_mem_size_in_bytes) + mock_chmod.assert_called_once_with(expected_swapfile_location, expected_swapfile_permission) + mock_run.assert_called_once_with(f'mkswap {expected_swapfile_location}; swapon {expected_swapfile_location}') + + def test_remove_swapmem(self): + with mock.patch("subprocess.Popen") as mock_popen, \ + mock.patch("os.unlink") as mock_unlink: + pseudo_subproc = mock.MagicMock() + mock_popen.return_value = pseudo_subproc + pseudo_subproc.communicate.return_value = ("swapoff: /home/swapfile: swapoff failed: No such file or directory", None) + pseudo_subproc.returncode = 255 + + swap_allocator = SWAPAllocator(allocate=True) + try: + swap_allocator.remove_swapmem() + except Exception as detail: + pytest.fail("SWAPAllocator.remove_swapmem should not raise exception %s" % repr(detail)) + + expected_swapfile_location = SWAPAllocator.SWAP_FILE_PATH + mock_popen.assert_called_once_with(['swapoff', expected_swapfile_location], stdout=subprocess.PIPE, text=True) + mock_unlink.assert_called_once_with(SWAPAllocator.SWAP_FILE_PATH) + + def test_swap_allocator_initialization_default_args(self): + expected_allocate = False + expected_swap_mem_size = SWAPAllocator.SWAP_MEM_SIZE + expected_total_mem_threshold = SWAPAllocator.TOTAL_MEM_THRESHOLD + expected_available_mem_threshold = SWAPAllocator.AVAILABLE_MEM_THRESHOLD + swap_allocator = SWAPAllocator(allocate=expected_allocate) + assert swap_allocator.allocate is expected_allocate + assert swap_allocator.swap_mem_size == expected_swap_mem_size + assert swap_allocator.total_mem_threshold == expected_total_mem_threshold + assert swap_allocator.available_mem_threshold == expected_available_mem_threshold + assert swap_allocator.is_allocated is False + + def test_swap_allocator_initialization_custom_args(self): + expected_allocate = True + expected_swap_mem_size = 2048 + expected_total_mem_threshold = 4096 + expected_available_mem_threshold = 1024 + swap_allocator = SWAPAllocator( + allocate=expected_allocate, + swap_mem_size=expected_swap_mem_size, + total_mem_threshold=expected_total_mem_threshold, + available_mem_threshold=expected_available_mem_threshold + ) + assert swap_allocator.allocate is expected_allocate + assert swap_allocator.swap_mem_size == expected_swap_mem_size + assert swap_allocator.total_mem_threshold == expected_total_mem_threshold + assert swap_allocator.available_mem_threshold == expected_available_mem_threshold + assert swap_allocator.is_allocated is False + + def test_swap_allocator_context_enter_allocate_true_insufficient_total_memory(self): + with mock.patch("sonic_installer.main.SWAPAllocator.get_disk_freespace") as mock_disk_free, \ + mock.patch("sonic_installer.main.SWAPAllocator.read_from_meminfo") as mock_meminfo, \ + mock.patch("sonic_installer.main.SWAPAllocator.setup_swapmem") as mock_setup, \ + mock.patch("sonic_installer.main.SWAPAllocator.remove_swapmem") as mock_remove, \ + mock.patch("os.path.exists") as mock_exists: + mock_disk_free.return_value = 10 * 1024 * 1024 * 1024 + mock_meminfo.return_value = { + "MemTotal": 2000000, + "MemAvailable": 1900000, + } + mock_exists.return_value = False + + swap_allocator = SWAPAllocator(allocate=True) + try: + swap_allocator.__enter__() + except Exception as detail: + pytest.fail("SWAPAllocator context manager should not raise exception %s" % repr(detail)) + mock_setup.assert_called_once() + mock_remove.assert_not_called() + assert swap_allocator.is_allocated is True + + def test_swap_allocator_context_enter_allocate_true_insufficient_available_memory(self): + with mock.patch("sonic_installer.main.SWAPAllocator.get_disk_freespace") as mock_disk_free, \ + mock.patch("sonic_installer.main.SWAPAllocator.read_from_meminfo") as mock_meminfo, \ + mock.patch("sonic_installer.main.SWAPAllocator.setup_swapmem") as mock_setup, \ + mock.patch("sonic_installer.main.SWAPAllocator.remove_swapmem") as mock_remove, \ + mock.patch("os.path.exists") as mock_exists: + mock_disk_free.return_value = 10 * 1024 * 1024 * 1024 + mock_meminfo.return_value = { + "MemTotal": 3000000, + "MemAvailable": 1000000, + } + mock_exists.return_value = False + + swap_allocator = SWAPAllocator(allocate=True) + try: + swap_allocator.__enter__() + except Exception as detail: + pytest.fail("SWAPAllocator context manager should not raise exception %s" % repr(detail)) + mock_setup.assert_called_once() + mock_remove.assert_not_called() + assert swap_allocator.is_allocated is True + + def test_swap_allocator_context_enter_allocate_true_insufficient_disk_space(self): + with mock.patch("sonic_installer.main.SWAPAllocator.get_disk_freespace") as mock_disk_free, \ + mock.patch("sonic_installer.main.SWAPAllocator.read_from_meminfo") as mock_meminfo, \ + mock.patch("sonic_installer.main.SWAPAllocator.setup_swapmem") as mock_setup, \ + mock.patch("sonic_installer.main.SWAPAllocator.remove_swapmem") as mock_remove, \ + mock.patch("os.path.exists") as mock_exists: + mock_disk_free.return_value = 1 * 1024 * 1024 * 1024 + mock_meminfo.return_value = { + "MemTotal": 32859496, + "MemAvailable": 16275512, + } + mock_exists.return_value = False + + swap_allocator = SWAPAllocator(allocate=True) + try: + swap_allocator.__enter__() + except Exception as detail: + pytest.fail("SWAPAllocator context manager should not raise exception %s" % repr(detail)) + mock_setup.assert_not_called() + mock_remove.assert_not_called() + assert swap_allocator.is_allocated is False + + def test_swap_allocator_context_enter_allocate_true_swapfile_present(self): + with mock.patch("sonic_installer.main.SWAPAllocator.get_disk_freespace") as mock_disk_free, \ + mock.patch("sonic_installer.main.SWAPAllocator.read_from_meminfo") as mock_meminfo, \ + mock.patch("sonic_installer.main.SWAPAllocator.setup_swapmem") as mock_setup, \ + mock.patch("sonic_installer.main.SWAPAllocator.remove_swapmem") as mock_remove, \ + mock.patch("os.path.exists") as mock_exists: + mock_disk_free.return_value = 10 * 1024 * 1024 * 1024 + mock_meminfo.return_value = { + "MemTotal": 32859496, + "MemAvailable": 1000000, + } + mock_exists.return_value = True + + swap_allocator = SWAPAllocator(allocate=True) + try: + swap_allocator.__enter__() + except Exception as detail: + pytest.fail("SWAPAllocator context manager should not raise exception %s" % repr(detail)) + mock_setup.assert_called_once() + mock_remove.assert_called_once() + assert swap_allocator.is_allocated is True + + def test_swap_allocator_context_enter_setup_error(self): + with mock.patch("sonic_installer.main.SWAPAllocator.get_disk_freespace") as mock_disk_free, \ + mock.patch("sonic_installer.main.SWAPAllocator.read_from_meminfo") as mock_meminfo, \ + mock.patch("sonic_installer.main.SWAPAllocator.setup_swapmem") as mock_setup, \ + mock.patch("sonic_installer.main.SWAPAllocator.remove_swapmem") as mock_remove, \ + mock.patch("os.path.exists") as mock_exists: + mock_disk_free.return_value = 10 * 1024 * 1024 * 1024 + mock_meminfo.return_value = { + "MemTotal": 32859496, + "MemAvailable": 1000000, + } + mock_exists.return_value = False + expected_err_str = "Pseudo Error" + mock_setup.side_effect = Exception(expected_err_str) + + swap_allocator = SWAPAllocator(allocate=True) + try: + swap_allocator.__enter__() + except Exception as detail: + assert expected_err_str in str(detail) + mock_setup.assert_called_once() + mock_remove.assert_called_once() + assert swap_allocator.is_allocated is False + + def test_swap_allocator_context_enter_allocate_false(self): + with mock.patch("sonic_installer.main.SWAPAllocator.get_disk_freespace") as mock_disk_free, \ + mock.patch("sonic_installer.main.SWAPAllocator.read_from_meminfo") as mock_meminfo, \ + mock.patch("sonic_installer.main.SWAPAllocator.setup_swapmem") as mock_setup, \ + mock.patch("sonic_installer.main.SWAPAllocator.remove_swapmem") as mock_remove, \ + mock.patch("os.path.exists") as mock_exists: + mock_disk_free.return_value = 10 * 1024 * 1024 * 1024 + mock_meminfo.return_value = { + "MemTotal": 32859496, + "MemAvailable": 1000000, + } + mock_exists.return_value = False + + swap_allocator = SWAPAllocator(allocate=False) + try: + swap_allocator.__enter__() + except Exception as detail: + pytest.fail("SWAPAllocator context manager should not raise exception %s" % repr(detail)) + mock_setup.assert_not_called() + mock_remove.assert_not_called() + assert swap_allocator.is_allocated is False + + def test_swap_allocator_context_exit_is_allocated_true(self): + with mock.patch("sonic_installer.main.SWAPAllocator.remove_swapmem") as mock_remove: + swap_allocator = SWAPAllocator(allocate=True) + swap_allocator.is_allocated = True + swap_allocator.__exit__(None, None, None) + mock_remove.assert_called_once() + + def test_swap_allocator_context_exit_is_allocated_false(self): + with mock.patch("sonic_installer.main.SWAPAllocator.remove_swapmem") as mock_remove: + swap_allocator = SWAPAllocator(allocate=True) + swap_allocator.is_allocated = False + swap_allocator.__exit__(None, None, None) + mock_remove.assert_not_called() diff --git a/utilities_common/cli.py b/utilities_common/cli.py index 1202bfdb9a36..72442878bebd 100644 --- a/utilities_common/cli.py +++ b/utilities_common/cli.py @@ -599,3 +599,33 @@ def is_interface_in_config_db(config_db, interface_name): return True + +class MutuallyExclusiveOption(click.Option): + """ + This option type is extended with `mutually_exclusive` parameter which make + CLI to ensure the other options specified in `mutually_exclusive` are not used. + """ + + def __init__(self, *args, **kwargs): + self.mutually_exclusive = set(kwargs.pop('mutually_exclusive', [])) + super(MutuallyExclusiveOption, self).__init__(*args, **kwargs) + + def get_help_record(self, ctx): + """Return help string with mutually_exclusive list added.""" + help_record = list(super(MutuallyExclusiveOption, self).get_help_record(ctx)) + if self.mutually_exclusive: + mutually_exclusive_str = 'NOTE: this argument is mutually exclusive with arguments: %s' % ', '.join(self.mutually_exclusive) + if help_record[-1]: + help_record[-1] += ' ' + mutually_exclusive_str + else: + help_record[-1] = mutually_exclusive_str + return tuple(help_record) + + def handle_parse_result(self, ctx, opts, args): + if self.name in opts and opts[self.name] is not None: + for opt_name in self.mutually_exclusive: + if opt_name in opts and opts[opt_name] is not None: + raise click.UsageError( + "Illegal usage: %s is mutually exclusive with arguments %s" % (self.name, ', '.join(self.mutually_exclusive)) + ) + return super(MutuallyExclusiveOption, self).handle_parse_result(ctx, opts, args)