From 331e71be9659f810050cc06358dd2946fbde64ef Mon Sep 17 00:00:00 2001
From: Olivier Gayot <olivier.gayot@canonical.com>
Date: Fri, 23 Aug 2024 17:28:10 +0200
Subject: [PATCH] storage: lean on systemd-shutdown to export zpools on
 shutdown

People have reported issues involving non exported zpools at the end of
Ubuntu installations. Diagnosing such issues is tricky because the
`zpool export` commands have to run after the target is unmounted. And
once the target is unmounted, logs can no longer be copied over to the
target so bug reports don't include traces of potential `zpool export`
failures (or failures to unmount target).

Instead of relying on Subiquity's _pre_shutdown hook, we now make use of
systemd-shutdown to export the zpools. The chance of successfully
exporting the zpools is expected to be higher since it will happen later
in the shutdown sequence ; after more mounts have been detached.

Signed-off-by: Olivier Gayot <olivier.gayot@canonical.com>
(cherry picked from commit 2a330206421103176aa9ee0c598a13cb9f369770)
---
 subiquity/models/filesystem.py             |  7 +++++++
 subiquity/server/controllers/filesystem.py |  2 --
 subiquity/server/controllers/install.py    | 22 ++++++++++++++++++++++
 3 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/subiquity/models/filesystem.py b/subiquity/models/filesystem.py
index 2e5930280..c0e3fd1f6 100644
--- a/subiquity/models/filesystem.py
+++ b/subiquity/models/filesystem.py
@@ -1986,6 +1986,13 @@ def render(self, mode: ActionRenderMode = ActionRenderMode.DEFAULT):
             config["grub"] = self.grub
         return config
 
+    def systemd_shutdown_commands(self) -> list[list[str]]:
+        """Return a list of commands meant to be executed by systemd-shutdown.
+        We entrust the execution of `zpool export` commands to systemd-shutdown
+        instead of subiquity's _pre_shutdown hook, in hope that the commands
+        will more likely succeed."""
+        return [["zpool", "export", zpool.name] for zpool in self._all(type="zpool")]
+
     def load_probe_data(self, probe_data):
         for devname, devdata in probe_data["blockdev"].items():
             if int(devdata["attrs"]["size"]) != 0:
diff --git a/subiquity/server/controllers/filesystem.py b/subiquity/server/controllers/filesystem.py
index f79d08b2d..0ae0f5a80 100644
--- a/subiquity/server/controllers/filesystem.py
+++ b/subiquity/server/controllers/filesystem.py
@@ -1644,5 +1644,3 @@ async def _pre_shutdown(self):
                 )
             else:
                 await self.app.command_runner.run(["umount", "--recursive", "/target"])
-        if len(self.model._all(type="zpool")) > 0:
-            await self.app.command_runner.run(["zpool", "export", "-a"])
diff --git a/subiquity/server/controllers/install.py b/subiquity/server/controllers/install.py
index 282a7eedb..2424b3acc 100644
--- a/subiquity/server/controllers/install.py
+++ b/subiquity/server/controllers/install.py
@@ -326,6 +326,25 @@ async def curtin_install(self, *, context, source):
 
         fs_controller = self.app.controllers.Filesystem
 
+        def register_shutdown_commands() -> None:
+            cmds: list[list[str]] = fs_controller.model.systemd_shutdown_commands()
+
+            if not cmds:
+                return
+
+            shutdown_dir = root / "usr/lib/systemd/system-shutdown"
+            shutdown_script = shutdown_dir / "subiquity-storage.shutdown"
+            shutdown_dir.mkdir(parents=True, exist_ok=True)
+            with shutdown_script.open(mode="w", encoding="utf-8") as stream:
+                # Generating a bash script is .. ewww
+                # Should we place the commands in a JSON file and hardcode a
+                # script that calls jq or something?
+                stream.write("#!/bin/bash\n")
+                for cmd in cmds:
+                    stream.write(shlex.join(cmd))
+                    stream.write("\n")
+            shutdown_script.chmod(0o755)
+
         async def run_curtin_step(name, stages, step_config, source=None):
             config = copy.deepcopy(base_config)
             filename = f"subiquity-{name.replace(' ', '-')}.conf"
@@ -349,6 +368,7 @@ async def run_curtin_step(name, stages, step_config, source=None):
                     device_map_path=logs_dir / "device-map.json",
                 ),
             )
+            register_shutdown_commands()
         elif fs_controller.is_core_boot_classic():
             await run_curtin_step(
                 name="partitioning",
@@ -368,6 +388,7 @@ async def run_curtin_step(name, stages, step_config, source=None):
                     device_map_path=logs_dir / "device-map-format.json",
                 ),
             )
+            register_shutdown_commands()
             await run_curtin_step(
                 name="extract",
                 stages=["extract"],
@@ -400,6 +421,7 @@ async def run_curtin_step(name, stages, step_config, source=None):
                 ),
                 source=source,
             )
+            register_shutdown_commands()
             await run_curtin_step(
                 name="extract",
                 stages=["extract"],