Feat: Zram algorithm config (#4042)

* Add configuration for swap algorithm. Backward compatible implementation

* Fix interaction to default to Yes and show (default)

* Fix mypy error

* Any -> str, str

* feedback Enums

* test file

* line length warning

* Renames

* Fix default values in TUI menu for display

* Address feedback

* More feedback, really appreciate it.

* Adapt to use same | None = None pattern

* Pytests

* Add missing import for Zram

archinstall/lib/args.py

@@ -13,7 +13,7 @@ from urllib.request import Request, urlopen
 from pydantic.dataclasses import dataclass as p_dataclass
 
 from archinstall.lib.crypt import decrypt
-from archinstall.lib.models.application import ApplicationConfiguration
+from archinstall.lib.models.application import ApplicationConfiguration, ZramConfiguration
 from archinstall.lib.models.authentication import AuthenticationConfiguration
 from archinstall.lib.models.bootloader import Bootloader, BootloaderConfiguration
 from archinstall.lib.models.device import DiskEncryption, DiskLayoutConfiguration
@@ -67,12 +67,12 @@ class ArchConfig:
 	bootloader_config: BootloaderConfiguration | None = None
 	app_config: ApplicationConfiguration | None = None
 	auth_config: AuthenticationConfiguration | None = None
+	swap: ZramConfiguration | None = None
 	hostname: str = 'archlinux'
 	kernels: list[str] = field(default_factory=lambda: ['linux'])
 	ntp: bool = True
 	packages: list[str] = field(default_factory=list)
 	parallel_downloads: int = 0
-	swap: bool = True
 	timezone: str = 'UTC'
 	services: list[str] = field(default_factory=list)
 	custom_commands: list[str] = field(default_factory=list)
@@ -211,7 +211,9 @@ class ArchConfig:
 		if parallel_downloads := args_config.get('parallel_downloads', 0):
 			arch_config.parallel_downloads = parallel_downloads
 
-		arch_config.swap = args_config.get('swap', True)
+		swap_arg = args_config.get('swap')
+		if swap_arg is not None:
+			arch_config.swap = ZramConfiguration.parse_arg(swap_arg)
 
 		if timezone := args_config.get('timezone', 'UTC'):
 			arch_config.timezone = timezone

archinstall/lib/global_menu.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from typing import override
 
 from archinstall.lib.disk.disk_menu import DiskLayoutConfigurationMenu
-from archinstall.lib.models.application import ApplicationConfiguration
+from archinstall.lib.models.application import ApplicationConfiguration, ZramConfiguration
 from archinstall.lib.models.authentication import AuthenticationConfiguration
 from archinstall.lib.models.device import DiskLayoutConfiguration, DiskLayoutType, EncryptionType, FilesystemType, PartitionModification
 from archinstall.lib.packages import list_available_packages
@@ -80,7 +80,7 @@ class GlobalMenu(AbstractMenu[None]):
 			),
 			MenuItem(
 				text=tr('Swap'),
-				value=True,
+				value=ZramConfiguration(enabled=True),
 				action=ask_for_swap,
 				preview_action=self._prev_swap,
 				key='swap',
@@ -372,7 +372,9 @@ class GlobalMenu(AbstractMenu[None]):
 	def _prev_swap(self, item: MenuItem) -> str | None:
 		if item.value is not None:
 			output = f'{tr("Swap on zram")}: '
-			output += tr('Enabled') if item.value else tr('Disabled')
+			output += tr('Enabled') if item.value.enabled else tr('Disabled')
+			if item.value.enabled:
+				output += f'\n{tr("Compression algorithm")}: {item.value.algorithm.value}'
 			return output
 		return None
 

archinstall/lib/installer.py

@@ -16,6 +16,7 @@ from typing import Any
 from archinstall.lib.disk.device_handler import device_handler
 from archinstall.lib.disk.fido import Fido2
 from archinstall.lib.disk.utils import get_lsblk_by_mountpoint, get_lsblk_info
+from archinstall.lib.models.application import ZramAlgorithm
 from archinstall.lib.models.device import (
 	DiskEncryption,
 	DiskLayoutConfiguration,
@@ -991,7 +992,7 @@ class Installer:
 			self._configure_grub_btrfsd(snapshot_type)
 			self.enable_service('grub-btrfsd.service')
 
-	def setup_swap(self, kind: str = 'zram') -> None:
+	def setup_swap(self, kind: str = 'zram', algo: ZramAlgorithm = ZramAlgorithm.ZSTD) -> None:
 		if kind == 'zram':
 			info('Setting up swap on zram')
 			self.pacman.strap('zram-generator')
@@ -1000,12 +1001,12 @@ class Installer:
 			# Convert KB to MB and divide by 2, with minimum of 4096 MB
 			size_mb = max(ram_kb // 2048, 4096)
 			info(f'Zram size: {size_mb} from RAM: {ram_kb}')
-			# We could use the default example below, but maybe not the best idea: https://github.com/archlinux/archinstall/pull/678#issuecomment-962124813
-			# zram_example_location = '/usr/share/doc/zram-generator/zram-generator.conf.example'
-			# shutil.copy2(f"{self.target}{zram_example_location}", f"{self.target}/usr/lib/systemd/zram-generator.conf")
+			info(f'Zram compression algorithm: {algo.value}')
+
 			with open(f'{self.target}/etc/systemd/zram-generator.conf', 'w') as zram_conf:
 				zram_conf.write('[zram0]\n')
 				zram_conf.write(f'zram-size = {size_mb}\n')
+				zram_conf.write(f'compression-algorithm = {algo.value}\n')
 
 			self.enable_service('systemd-zram-setup@zram0.service')
 

archinstall/lib/interactions/system_conf.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from archinstall.lib.models.application import ZramAlgorithm, ZramConfiguration
 from archinstall.lib.translationhandler import tr
 from archinstall.tui.curses_menu import SelectMenu
 from archinstall.tui.menu_item import MenuItem, MenuItemGroup
@@ -89,16 +90,12 @@ def select_driver(options: list[GfxDriver] = [], preset: GfxDriver | None = None
 			return result.get_value()
 
 
-def ask_for_swap(preset: bool = True) -> bool:
-	if preset:
-		default_item = MenuItem.yes()
-	else:
-		default_item = MenuItem.no()
-
+def ask_for_swap(preset: ZramConfiguration = ZramConfiguration(enabled=True)) -> ZramConfiguration:
 	prompt = tr('Would you like to use swap on zram?') + '\n'
 
 	group = MenuItemGroup.yes_no()
-	group.set_focus_by_value(default_item)
+	group.set_default_by_value(True)
+	group.set_focus_by_value(preset.enabled)
 
 	result = SelectMenu[bool](
 		group,
@@ -113,6 +110,29 @@ def ask_for_swap(preset: bool = True) -> bool:
 		case ResultType.Skip:
 			return preset
 		case ResultType.Selection:
-			return result.item() == MenuItem.yes()
+			enabled = result.item() == MenuItem.yes()
+			if not enabled:
+				return ZramConfiguration(enabled=False)
+
+			# Ask for compression algorithm
+			algo_group = MenuItemGroup.from_enum(ZramAlgorithm, sort_items=False)
+			algo_group.set_default_by_value(ZramAlgorithm.ZSTD)
+			algo_group.set_focus_by_value(preset.algorithm)
+
+			algo_result = SelectMenu[ZramAlgorithm](
+				algo_group,
+				header=tr('Select zram compression algorithm:') + '\n',
+				alignment=Alignment.CENTER,
+				allow_skip=True,
+			).run()
+
+			algo: ZramAlgorithm
+			match algo_result.type_:
+				case ResultType.Skip:
+					algo = preset.algorithm
+				case ResultType.Selection:
+					algo = algo_result.get_value()
+
+			return ZramConfiguration(enabled=True, algorithm=algo)
 		case ResultType.Reset:
 			raise ValueError('Unhandled result type')

archinstall/lib/models/application.py

@@ -21,6 +21,14 @@ class PrintServiceConfigSerialization(TypedDict):
 	enabled: bool
 
 
+class ZramAlgorithm(StrEnum):
+	ZSTD = 'zstd'
+	LZO_RLE = 'lzo-rle'
+	LZO = 'lzo'
+	LZ4 = 'lz4'
+	LZ4HC = 'lz4hc'
+
+
 class ApplicationSerialization(TypedDict):
 	bluetooth_config: NotRequired[BluetoothConfigSerialization]
 	audio_config: NotRequired[AudioConfigSerialization]
@@ -67,6 +75,21 @@ class PrintServiceConfiguration:
 		return PrintServiceConfiguration(arg['enabled'])
 
 
+@dataclass(frozen=True)
+class ZramConfiguration:
+	enabled: bool
+	algorithm: ZramAlgorithm = ZramAlgorithm.ZSTD
+
+	@staticmethod
+	def parse_arg(arg: bool | dict[str, Any]) -> 'ZramConfiguration':
+		if isinstance(arg, bool):
+			return ZramConfiguration(enabled=arg)
+
+		enabled = arg.get('enabled', True)
+		algo = arg.get('algorithm', arg.get('algo', ZramAlgorithm.ZSTD.value))
+		return ZramConfiguration(enabled=enabled, algorithm=ZramAlgorithm(algo))
+
+
 @dataclass
 class ApplicationConfiguration:
 	bluetooth_config: BluetoothConfiguration | None = None

archinstall/scripts/guided.py

@@ -98,8 +98,8 @@ def perform_installation(mountpoint: Path) -> None:
 		if mirror_config := config.mirror_config:
 			installation.set_mirrors(mirror_config, on_target=True)
 
-		if config.swap:
-			installation.setup_swap('zram')
+		if config.swap and config.swap.enabled:
+			installation.setup_swap('zram', algo=config.swap.algorithm)
 
 		if config.bootloader_config and config.bootloader_config.bootloader != Bootloader.NO_BOOTLOADER:
 			if config.bootloader_config.bootloader == Bootloader.Grub and SysInfo.has_uefi():

examples/config-sample.json

@@ -146,7 +146,10 @@
   },
   "script": "guided",
   "silent": false,
-  "swap": true,
+  "swap": {
+    "enabled": true,
+    "algorithm": "zstd"
+  },
   "timezone": "UTC",
   "version": "2.8.6"
 }

tests/test_args.py

@@ -7,7 +7,14 @@ from pytest import MonkeyPatch
 from archinstall.default_profiles.profile import GreeterType
 from archinstall.lib.args import ArchConfig, ArchConfigHandler, Arguments
 from archinstall.lib.hardware import GfxDriver
-from archinstall.lib.models.application import ApplicationConfiguration, Audio, AudioConfiguration, BluetoothConfiguration, PrintServiceConfiguration
+from archinstall.lib.models.application import (
+	ApplicationConfiguration,
+	Audio,
+	AudioConfiguration,
+	BluetoothConfiguration,
+	PrintServiceConfiguration,
+	ZramConfiguration,
+)
 from archinstall.lib.models.authentication import AuthenticationConfiguration, U2FLoginConfiguration, U2FLoginMethod
 from archinstall.lib.models.bootloader import Bootloader, BootloaderConfiguration
 from archinstall.lib.models.device import DiskLayoutConfiguration, DiskLayoutType
@@ -225,7 +232,7 @@ def test_config_file_parsing(
 		ntp=True,
 		packages=['firefox'],
 		parallel_downloads=66,
-		swap=False,
+		swap=ZramConfiguration(enabled=False),
 		timezone='UTC',
 		services=['service_1', 'service_2'],
 		custom_commands=["echo 'Hello, World!'"],