slot_mnt_helper: register atexit hook for mount points

src/otaclient/boot_control/_slot_mnt_helper.py

@@ -16,103 +16,18 @@
 
 from __future__ import annotations
 
+import atexit
 import logging
 import shutil
+from functools import partial
 from pathlib import Path
-from subprocess import CalledProcessError
-from time import sleep
 
 from otaclient.configs.cfg import cfg
 from otaclient_common import cmdhelper, replace_root
 from otaclient_common.typing import StrOrPath
 
 logger = logging.getLogger(__name__)
 
-MAX_RETRY_COUNT = 6
-RETRY_INTERVAL = 2
-
-
-def ensure_mount(
-    target: StrOrPath, mnt_point: StrOrPath, *, mount_func, raise_exception: bool
-) -> None:  # pragma: no cover
-    """Ensure the <target> mounted on <mnt_point> by our best.
-
-    Raises:
-        If <raise_exception> is True, raises the last failed attemp's CalledProcessError.
-    """
-    for _retry in range(MAX_RETRY_COUNT + 1):
-        try:
-            mount_func(target=target, mount_point=mnt_point)
-            cmdhelper.is_target_mounted(mnt_point, raise_exception=True)
-            return
-        except CalledProcessError as e:
-            logger.error(
-                f"retry#{_retry} failed to mount {target} on {mnt_point}: {e!r}"
-            )
-            logger.error(f"{e.stderr=}\n{e.stdout=}")
-
-            if _retry >= MAX_RETRY_COUNT:
-                logger.error(
-                    f"exceed max retry count mounting {target} on {mnt_point}, abort"
-                )
-                if raise_exception:
-                    raise
-                return
-
-            sleep(RETRY_INTERVAL)
-            continue
-
-
-def ensure_umount(
-    mnt_point: StrOrPath, *, ignore_error: bool
-) -> None:  # pragma: no cover
-    """Try to umount the <mnt_point> at our best.
-
-    Raises:
-        If <ignore_error> is False, raises the last failed attemp's CalledProcessError.
-    """
-    for _retry in range(MAX_RETRY_COUNT + 1):
-        try:
-            if not cmdhelper.is_target_mounted(mnt_point, raise_exception=False):
-                break
-            cmdhelper.umount(mnt_point, raise_exception=True)
-        except CalledProcessError as e:
-            logger.warning(f"retry#{_retry} failed to umount {mnt_point}: {e!r}")
-            logger.warning(f"{e.stderr}\n{e.stdout}")
-
-            if _retry >= MAX_RETRY_COUNT:
-                logger.error(f"reached max retry on umounting {mnt_point}, abort")
-                if not ignore_error:
-                    raise
-                return
-
-            sleep(RETRY_INTERVAL)
-            continue
-
-
-def ensure_mointpoint(
-    mnt_point: Path, *, ignore_error: bool
-) -> None:  # pragma: no cover
-    """Ensure the <mnt_point> exists, has no mount on it and ready for mount.
-
-    If the <mnt_point> is valid, but we failed to umount any previous mounts on it,
-        we still keep use the mountpoint as later mount will override the previous one.
-    """
-    if mnt_point.is_symlink() or not mnt_point.is_dir():
-        mnt_point.unlink(missing_ok=True)
-
-    if not mnt_point.exists():
-        mnt_point.mkdir(exist_ok=True, parents=True)
-        return
-
-    try:
-        ensure_umount(mnt_point, ignore_error=ignore_error)
-    except Exception:
-        logger.warning(
-            f"{mnt_point} still has other mounts on it, "
-            f"but still use {mnt_point} and override the previous mount"
-        )
-
 
 class SlotMountHelper:  # pragma: no cover
     """Helper class that provides methods for mounting slots."""
@@ -141,17 +56,34 @@ def __init__(
             )
         )
 
+        # ensure the each mount points being umounted at termination
+        atexit.register(
+            partial(
+                cmdhelper.ensure_umount,
+                self.active_slot_mount_point,
+                ignore_error=True,
+            )
+        )
+        atexit.register(
+            partial(
+                cmdhelper.ensure_umount,
+                self.standby_slot_mount_point,
+                ignore_error=True,
+                max_retry=3,
+            )
+        )
+
     def mount_standby(self) -> None:
         """Mount standby slot dev rw to <standby_slot_mount_point>.
 
         Raises:
             CalledProcessedError on the last failed attemp.
         """
         logger.debug("mount standby slot rootfs dev...")
-        ensure_mointpoint(self.standby_slot_mount_point, ignore_error=True)
-        ensure_umount(self.standby_slot_dev, ignore_error=False)
+        cmdhelper.ensure_mointpoint(self.standby_slot_mount_point, ignore_error=True)
+        cmdhelper.ensure_umount(self.standby_slot_dev, ignore_error=False)
 
-        ensure_mount(
+        cmdhelper.ensure_mount(
             target=self.standby_slot_dev,
             mnt_point=self.standby_slot_mount_point,
             mount_func=cmdhelper.mount_rw,
@@ -165,8 +97,8 @@ def mount_active(self) -> None:
             CalledProcessedError on the last failed attemp.
         """
         logger.debug("mount active slot rootfs dev...")
-        ensure_mointpoint(self.active_slot_mount_point, ignore_error=True)
-        ensure_mount(
+        cmdhelper.ensure_mointpoint(self.active_slot_mount_point, ignore_error=True)
+        cmdhelper.ensure_mount(
             target=self.active_rootfs,
             mnt_point=self.active_slot_mount_point,
             mount_func=cmdhelper.bind_mount_ro,
@@ -195,7 +127,7 @@ def prepare_standby_dev(
         erase_standby: bool = False,
         fslabel: str | None = None,
     ) -> None:
-        ensure_umount(self.standby_slot_dev, ignore_error=True)
+        cmdhelper.ensure_umount(self.standby_slot_dev, ignore_error=True)
         if erase_standby:
             return cmdhelper.mkfs_ext4(self.standby_slot_dev, fslabel=fslabel)
 
@@ -206,5 +138,7 @@ def prepare_standby_dev(
 
     def umount_all(self, *, ignore_error: bool = True):
         logger.debug("unmount standby slot and active slot mount point...")
-        ensure_umount(self.active_slot_mount_point, ignore_error=ignore_error)
-        ensure_umount(self.standby_slot_mount_point, ignore_error=ignore_error)
+        cmdhelper.ensure_umount(self.active_slot_mount_point, ignore_error=ignore_error)
+        cmdhelper.ensure_umount(
+            self.standby_slot_mount_point, ignore_error=ignore_error
+        )