[sonic_installer] Add swap setup support

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.")

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()