From 2a330206421103176aa9ee0c598a13cb9f369770 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>
---
 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 72bffc8d5..fea082101 100644
--- a/subiquity/models/filesystem.py
+++ b/subiquity/models/filesystem.py
@@ -2053,6 +2053,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 ed7923ee8..418fe96f9 100644
--- a/subiquity/server/controllers/filesystem.py
+++ b/subiquity/server/controllers/filesystem.py
@@ -1765,5 +1765,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 373bae154..08f067a99 100644
--- a/subiquity/server/controllers/install.py
+++ b/subiquity/server/controllers/install.py
@@ -345,6 +345,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"
@@ -368,6 +387,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.use_snapd_install_api():
             await run_curtin_step(
                 name="partitioning",
@@ -387,6 +407,7 @@ async def run_curtin_step(name, stages, step_config, source=None):
                     device_map_path=logs_dir / "device-map-format.json",
                 ),
             )
+            register_shutdown_commands()
             if source is not None:
                 await run_curtin_step(
                     name="extract",
@@ -420,6 +441,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"],