refactor: boot_control: split reboot from post_update/post_rollback A…

…PI (#427)

This PR splits the reboot from boot_control post_update/post_rollback APIs, instead creating new finalizing_update/finalizing_rollback APIs for reboot operation.

src/otaclient/boot_control/_grub.py

@@ -40,7 +40,7 @@
 from pathlib import Path
 from pprint import pformat
 from subprocess import CalledProcessError
-from typing import ClassVar, Dict, Generator, List, Optional, Tuple
+from typing import ClassVar, Dict, List, NoReturn, Optional, Tuple
 
 from otaclient import errors as ota_errors
 from otaclient._types import OTAStatus
@@ -895,7 +895,7 @@ def pre_update(self, version: str, *, standby_as_ref: bool, erase_standby=False)
                 _err_msg, module=__name__
             ) from e
 
-    def post_update(self) -> Generator[None, None, None]:
+    def post_update(self) -> None:
         try:
             logger.info("grub_boot: post-update setup...")
             # ------ update fstab ------ #
@@ -916,11 +916,18 @@ def post_update(self) -> Generator[None, None, None]:
             # ------ pre-reboot ------ #
             self._mp_control.umount_all(ignore_error=True)
             self._boot_control.grub_reboot_to_standby()
+        except Exception as e:
+            _err_msg = f"failed on post_update: {e!r}"
+            logger.error(_err_msg)
+            raise ota_errors.BootControlPostUpdateFailed(
+                _err_msg, module=__name__
+            ) from e
 
-            yield  # hand over control to otaclient
+    def finalizing_update(self) -> NoReturn:
+        try:
             cmdhelper.reboot()
         except Exception as e:
-            _err_msg = f"failed on post_update: {e!r}"
+            _err_msg = f"reboot failed: {e!r}"
             logger.error(_err_msg)
             raise ota_errors.BootControlPostUpdateFailed(
                 _err_msg, module=__name__
@@ -944,10 +951,11 @@ def post_rollback(self):
             logger.info("grub_boot: post-rollback setup...")
             self._boot_control.grub_reboot_to_standby()
             self._mp_control.umount_all(ignore_error=True)
-            cmdhelper.reboot()
         except Exception as e:
             _err_msg = f"failed on pre_rollback: {e!r}"
             logger.error(_err_msg)
             raise ota_errors.BootControlPostRollbackFailed(
                 _err_msg, module=__name__
             ) from e
+
+    finalizing_rollback = finalizing_update

src/otaclient/boot_control/_jetson_cboot.py

@@ -22,7 +22,7 @@
 import logging
 import subprocess
 from pathlib import Path
-from typing import Generator, Optional
+from typing import NoReturn, Optional
 
 from otaclient import errors as ota_errors
 from otaclient._types import OTAStatus
@@ -616,7 +616,7 @@ def pre_update(self, version: str, *, standby_as_ref: bool, erase_standby: bool)
                 _err_msg, module=__name__
             ) from e
 
-    def post_update(self) -> Generator[None, None, None]:
+    def post_update(self) -> None:
         try:
             logger.info("jetson-cboot: post-update ...")
             # ------ update extlinux.conf ------ #
@@ -677,15 +677,23 @@ def post_update(self) -> Generator[None, None, None]:
             self._mp_control.umount_all(ignore_error=True)
             logger.info(f"[post-update]: \n{NVBootctrlJetsonCBOOT.dump_slots_info()}")
             logger.info("post update finished, wait for reboot ...")
-            yield  # hand over control back to otaclient
-            cmdhelper.reboot()
         except Exception as e:
             _err_msg = f"failed on post_update: {e!r}"
             logger.error(_err_msg)
             raise ota_errors.BootControlPostUpdateFailed(
                 _err_msg, module=__name__
             ) from e
 
+    def finalizing_update(self) -> NoReturn:
+        try:
+            cmdhelper.reboot()
+        except Exception as e:
+            _err_msg = f"reboot failed: {e!r}"
+            logger.error(_err_msg)
+            raise ota_errors.BootControlPostUpdateFailed(
+                _err_msg, module=__name__
+            ) from e
+
     def pre_rollback(self):
         try:
             logger.info("jetson-cboot: pre-rollback setup ...")
@@ -704,14 +712,15 @@ def post_rollback(self):
             logger.info("jetson-cboot: post-rollback setup...")
             self._mp_control.umount_all(ignore_error=True)
             self._cboot_control.switch_boot_to_standby()
-            cmdhelper.reboot()
         except Exception as e:
             _err_msg = f"failed on post_rollback: {e!r}"
             logger.error(_err_msg)
             raise ota_errors.BootControlPostRollbackFailed(
                 _err_msg, module=__name__
             ) from e
 
+    finalizing_rollback = finalizing_update
+
     def on_operation_failure(self):
         """Failure registering and cleanup at failure."""
         logger.warning("on failure try to unmounting standby slot...")

src/otaclient/boot_control/_jetson_uefi.py

@@ -27,7 +27,7 @@
 import shutil
 import subprocess
 from pathlib import Path
-from typing import Any, ClassVar, Generator, Literal
+from typing import Any, ClassVar, Generator, Literal, NoReturn
 
 from pydantic import BaseModel
 from typing_extensions import Self
@@ -982,7 +982,7 @@ def pre_update(self, version: str, *, standby_as_ref: bool, erase_standby: bool)
                 _err_msg, module=__name__
             ) from e
 
-    def post_update(self) -> Generator[None, None, None]:
+    def post_update(self) -> None:
         try:
             logger.info("jetson-uefi: post-update ...")
             # ------ update extlinux.conf ------ #
@@ -1063,15 +1063,23 @@ def post_update(self) -> Generator[None, None, None]:
             # ------ prepare to reboot ------ #
             self._mp_control.umount_all(ignore_error=True)
             logger.info("post update finished, wait for reboot ...")
-            yield  # hand over control back to otaclient
-            cmdhelper.reboot()
         except Exception as e:
             _err_msg = f"jetson-uefi: failed on post_update: {e!r}"
             logger.error(_err_msg)
             raise ota_errors.BootControlPostUpdateFailed(
                 _err_msg, module=__name__
             ) from e
 
+    def finalizing_update(self) -> NoReturn:
+        try:
+            cmdhelper.reboot()
+        except Exception as e:
+            _err_msg = f"reboot failed: {e!r}"
+            logger.error(_err_msg)
+            raise ota_errors.BootControlPostUpdateFailed(
+                _err_msg, module=__name__
+            ) from e
+
     def pre_rollback(self):
         try:
             logger.info("jetson-uefi: pre-rollback setup ...")
@@ -1090,14 +1098,15 @@ def post_rollback(self):
             logger.info("jetson-uefi: post-rollback setup...")
             self._mp_control.umount_all(ignore_error=True)
             self._uefi_control.switch_boot_to_standby()
-            cmdhelper.reboot()
         except Exception as e:
             _err_msg = f"jetson-uefi: failed on post_rollback: {e!r}"
             logger.error(_err_msg)
             raise ota_errors.BootControlPostRollbackFailed(
                 _err_msg, module=__name__
             ) from e
 
+    finalizing_rollback = finalizing_update
+
     def on_operation_failure(self):
         """Failure registering and cleanup at failure."""
         logger.warning("on failure try to unmounting standby slot...")

src/otaclient/boot_control/_rpi_boot.py

@@ -22,7 +22,7 @@
 import subprocess
 from pathlib import Path
 from string import Template
-from typing import Any, Generator, Literal
+from typing import Any, Generator, Literal, NoReturn
 
 from typing_extensions import Self
 
@@ -523,15 +523,14 @@ def post_rollback(self):
             logger.info("rpi_boot: post-rollback setup...")
             self._rpiboot_control.prepare_tryboot_txt()
             self._mp_control.umount_all(ignore_error=True)
-            self._rpiboot_control.reboot_tryboot()
         except Exception as e:
             _err_msg = f"failed on post_rollback: {e!r}"
             logger.error(_err_msg)
             raise ota_errors.BootControlPostRollbackFailed(
                 _err_msg, module=__name__
             ) from e
 
-    def post_update(self) -> Generator[None, None, None]:
+    def post_update(self) -> None:
         try:
             logger.info("rpi_boot: post-update setup...")
             self._mp_control.preserve_ota_folder_to_standby()
@@ -542,15 +541,25 @@ def post_update(self) -> Generator[None, None, None]:
             )
             self._rpiboot_control.prepare_tryboot_txt()
             self._mp_control.umount_all(ignore_error=True)
-            yield  # hand over control back to otaclient
-            self._rpiboot_control.reboot_tryboot()
         except Exception as e:
             _err_msg = f"failed on post_update: {e!r}"
             logger.error(_err_msg)
             raise ota_errors.BootControlPostUpdateFailed(
                 _err_msg, module=__name__
             ) from e
 
+    def finalizing_update(self) -> NoReturn:
+        try:
+            self._rpiboot_control.reboot_tryboot()
+        except Exception as e:
+            _err_msg = f"reboot failed: {e!r}"
+            logger.error(_err_msg)
+            raise ota_errors.BootControlPostUpdateFailed(
+                _err_msg, module=__name__
+            ) from e
+
+    finalizing_rollback = finalizing_update
+
     def on_operation_failure(self):
         """Failure registering and cleanup at failure."""
         logger.warning("on failure try to unmounting standby slot...")

src/otaclient/boot_control/protocol.py

@@ -17,7 +17,7 @@
 
 from abc import abstractmethod
 from pathlib import Path
-from typing import Generator, Protocol
+from typing import Protocol
 
 from typing_extensions import deprecated
 
@@ -55,23 +55,40 @@ def get_standby_boot_dir(self) -> Path:
                         them to actual boot dir.
         """
 
+    @abstractmethod
+    def load_version(self) -> str:
+        """Read the version info from the current slot."""
+
+    @abstractmethod
+    def on_operation_failure(self) -> None:
+        """Cleanup by boot_control implementation when OTA failed."""
+
+    #
+    # ------ update ------ #
+    #
+
     @abstractmethod
     def pre_update(
         self, version: str, *, standby_as_ref: bool, erase_standby: bool
     ): ...
 
     @abstractmethod
-    def pre_rollback(self): ...
+    def post_update(self) -> None: ...
 
     @abstractmethod
-    def post_update(self) -> Generator[None, None, None]: ...
+    def finalizing_update(self) -> None:
+        """Normally this method only reboots the device."""
+
+    #
+    # ------ rollback ------ #
+    #
 
     @abstractmethod
-    def post_rollback(self): ...
+    def pre_rollback(self) -> None: ...
 
     @abstractmethod
-    def load_version(self) -> str:
-        """Read the version info from the current slot."""
+    def post_rollback(self): ...
 
     @abstractmethod
-    def on_operation_failure(self): ...
+    def finalizing_rollback(self) -> None:
+        """Normally this method only reboots the device."""

src/otaclient/ota_core.py

@@ -548,9 +548,7 @@ def _execute_update(self):
         )
         # NOTE(20240219): move persist file handling here
         self._process_persistents(otameta)
-
-        # boot controller postupdate
-        next(_postupdate_gen := self._boot_controller.post_update())
+        self._boot_controller.post_update()
 
         # ------ finalizing update ------ #
         logger.info("local update finished, wait on all subecs...")
@@ -571,7 +569,7 @@ def _execute_update(self):
 
         logger.info(f"device will reboot in {WAIT_BEFORE_REBOOT} seconds!")
         time.sleep(WAIT_BEFORE_REBOOT)
-        next(_postupdate_gen, None)  # reboot
+        self._boot_controller.finalizing_update()
 
     # API
 
@@ -600,6 +598,7 @@ def execute(self):
         try:
             self._boot_controller.pre_rollback()
             self._boot_controller.post_rollback()
+            self._boot_controller.finalizing_rollback()
         except ota_errors.OTAError as e:
             logger.error(f"rollback failed: {e!r}")
             self._boot_controller.on_operation_failure()

src/otaclient_common/cmdhelper.py

@@ -443,16 +443,15 @@ def reboot(args: list[str] | None = None) -> NoReturn:  # pragma: no cover
     Args:
         args (Optional[list[str]], optional): args passed to reboot command.
             Defaults to None, not passing any args.
+
+    Raises:
+        CalledProcessError for the reboot call, or SystemExit on sys.exit(0).
     """
     cmd = ["reboot"]
     if args:
         logger.info(f"will reboot with argument: {args=}")
         cmd.extend(args)
 
-    try:
-        logger.warning("system will reboot now!")
-        subprocess_call(cmd, raise_exception=True)
-        sys.exit(0)
-    except CalledProcessError:
-        logger.exception("failed to reboot")
-        raise
+    logger.warning("system will reboot now!")
+    subprocess_call(cmd, raise_exception=True)
+    sys.exit(0)

tests/test_otaclient/test_boot_control/test_grub.py

@@ -20,6 +20,7 @@
 import shutil
 import typing
 from pathlib import Path
+from typing import Any
 
 import pytest
 import pytest_mock
@@ -368,10 +369,12 @@ def test_grub_normal_update(
         shutil.copy(slot_a_ota_partition_dir / _initrd, slot_b / "boot")
 
         logger.info("pre-update completed, entering post-update...")
+
         # test post-update
-        _post_updater = grub_controller.post_update()
-        next(_post_updater)
-        next(_post_updater, None)
+        grub_controller: Any  # for typing
+        grub_controller.post_update()
+        grub_controller.finalizing_update()
+
         assert (
             slot_b / Path(cfg.FSTAB_FILE).relative_to("/")
         ).read_text().strip() == self.FSTAB_UPDATED.strip()

tests/test_otaclient/test_boot_control/test_rpi_boot.py

@@ -19,6 +19,7 @@
 import typing
 from pathlib import Path
 from string import Template
+from typing import Any
 
 import pytest
 import pytest_mock
@@ -262,10 +263,10 @@ def test_rpi_boot_normal_update(self, mocker: pytest_mock.MockerFixture):
         mocker.patch(f"{MODULE}.cfg", _mock_otaclient_cfg)
 
         # ------ boot_controller_inst1.stage1: init ------ #
-        rpi_boot_controller1 = RPIBootController()
+        rpi_boot_controller = RPIBootController()
 
         # ------ boot_controller_inst1.stage2: pre_update ------ #
-        rpi_boot_controller1.pre_update(
+        rpi_boot_controller.pre_update(
             version=VERSION,
             standby_as_ref=False,
             erase_standby=False,
@@ -299,9 +300,9 @@ def test_rpi_boot_normal_update(self, mocker: pytest_mock.MockerFixture):
         shutil.copy(os.path.realpath(_initrd_img), self.slot_b_boot_dir)
 
         # ------ boot_controller_inst1.stage3: post_update, reboot switch boot ------ #
-        _post_updater = rpi_boot_controller1.post_update()
-        next(_post_updater)
-        next(_post_updater, None)  # actual reboot here
+        rpi_boot_controller: Any  # for typing only
+        rpi_boot_controller.post_update()
+        rpi_boot_controller.finalizing_update()
 
         # --- assertion --- #
         self.reboot_tryboot_mock.assert_called_once()