From c1f21e7ca437e1e8a712ec9bed3bfdf3f4c893e5 Mon Sep 17 00:00:00 2001
From: Anton Hvornum <anton@hvornum.se>
Date: Mon, 19 Sep 2022 23:42:08 +0200
Subject: Add compression to /etc/fstab for btrfs subvolumes (#1473)

* Adding a btrfs compression plugin to genfstab

* Allowing the genfstab plugin to break on success
---
 archinstall/lib/disk/btrfs/__init__.py      |  4 ----
 archinstall/lib/disk/btrfs/btrfs_helpers.py | 34 +++++++++++++++++++++++++++++
 archinstall/lib/installer.py                |  3 ++-
 3 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/archinstall/lib/disk/btrfs/__init__.py b/archinstall/lib/disk/btrfs/__init__.py
index 3c183112..a26e0160 100644
--- a/archinstall/lib/disk/btrfs/__init__.py
+++ b/archinstall/lib/disk/btrfs/__init__.py
@@ -54,7 +54,3 @@ def create_subvolume(installation: Installer, subvolume_location :Union[pathlib.
 	log(f"Creating a subvolume on {target}", level=logging.INFO)
 	if (cmd := SysCommand(f"btrfs subvolume create {target}")).exit_code != 0:
 		raise DiskError(f"Could not create a subvolume at {target}: {cmd}")
-
-
-def manage_btrfs_subvolumes(installation :Installer, partition :Dict[str, str]) -> list:
-	raise Deprecated("Use setup_subvolumes() instead.")
diff --git a/archinstall/lib/disk/btrfs/btrfs_helpers.py b/archinstall/lib/disk/btrfs/btrfs_helpers.py
index ab528388..f6d2734a 100644
--- a/archinstall/lib/disk/btrfs/btrfs_helpers.py
+++ b/archinstall/lib/disk/btrfs/btrfs_helpers.py
@@ -1,4 +1,5 @@
 import logging
+import re
 from pathlib import Path
 from typing import Optional, Dict, Any, TYPE_CHECKING
 
@@ -6,6 +7,7 @@ from ...models.subvolume import Subvolume
 from ...exceptions import SysCallError, DiskError
 from ...general import SysCommand
 from ...output import log
+from ...plugins import plugins
 from ..helpers import get_mount_info
 from .btrfssubvolumeinfo import BtrfsSubvolumeInfo
 
@@ -14,6 +16,35 @@ if TYPE_CHECKING:
 	from ...installer import Installer
 
 
+class fstab_btrfs_compression_plugin():
+	def __init__(self, partition_dict):
+		self.partition_dict = partition_dict
+
+	def on_genfstab(self, installation):
+		with open(f"{installation.target}/etc/fstab", 'r') as fh:
+			fstab = fh.read()
+
+		# Replace the {installation}/etc/fstab with entries
+		# using the compress=zstd where the mountpoint has compression set.
+		with open(f"{installation.target}/etc/fstab", 'w') as fh:
+			for line in fstab.split('\n'):
+				# So first we grab the mount options by using subvol=.*? as a locator.
+				# And we also grab the mountpoint for the entry, for instance /var/log
+				if (subvoldef := re.findall(',.*?subvol=.*?[\t ]', line)) and (mountpoint := re.findall('[\t ]/.*?[\t ]', line)):
+					for subvolume in self.partition_dict.get('btrfs', {}).get('subvolumes', []):
+						# We then locate the correct subvolume and check if it's compressed
+						if subvolume.compress and subvolume.mountpoint == mountpoint[0].strip():
+							# We then sneak in the compress=zstd option if it doesn't already exist:
+							# We skip entries where compression is already defined
+							if ',compress=zstd,' not in line:
+								line = line.replace(subvoldef[0], f",compress=zstd{subvoldef[0]}")
+								break
+
+				fh.write(f"{line}\n")
+
+		return True
+
+
 def mount_subvolume(installation: 'Installer', device: 'BTRFSPartition', subvolume: Subvolume):
 	# we normalize the subvolume name (getting rid of slash at the start if exists.
 	# In our implementation has no semantic load.
@@ -61,6 +92,9 @@ def setup_subvolumes(installation: 'Installer', partition_dict: Dict[str, Any]):
 				if (cmd := SysCommand(f"chattr +c {installation.target}/{name}")).exit_code != 0:
 					raise DiskError(f"Could not set compress attribute at {installation.target}/{name}: {cmd}")
 
+			if 'fstab_btrfs_compression_plugin' not in plugins:
+				plugins['fstab_btrfs_compression_plugin'] = fstab_btrfs_compression_plugin(partition_dict)
+
 
 def subvolume_info_from_path(path: Path) -> Optional[BtrfsSubvolumeInfo]:
 	try:
diff --git a/archinstall/lib/installer.py b/archinstall/lib/installer.py
index c7682655..1270959e 100644
--- a/archinstall/lib/installer.py
+++ b/archinstall/lib/installer.py
@@ -432,7 +432,8 @@ class Installer:
 
 		for plugin in plugins.values():
 			if hasattr(plugin, 'on_genfstab'):
-				plugin.on_genfstab(self)
+				if plugin.on_genfstab(self) is True:
+					break
 
 		return True
 
-- 
cgit v1.2.3-70-g09d2