From cf828eeb8ddaf022f01fb4ca5f6685daaa3c8184 Mon Sep 17 00:00:00 2001
From: Olivier Gayot <olivier.gayot@canonical.com>
Date: Tue, 19 Sep 2023 15:31:43 +0200
Subject: [PATCH] s390x: ensure chreipl is called before unmounting /target

For ZFS, we recently introduced a call to $(umount --recursive /target)
slighly before shutting down or rebooting. Unfortunately, on s390x, we
also had a very late call to chreipl to make the firmware boot from the
installed system.

The call to chreipl reads data from /target/boot, and it fails if the
filesystem is no longer mounted.

Fixed by calling chreipl earlier in the installation, during the
postinst phase rather than after the user clicks "reboot".

Signed-off-by: Olivier Gayot <olivier.gayot@canonical.com>
---
 subiquity/server/controllers/install.py       | 13 +++++++
 subiquity/server/controllers/shutdown.py      |  6 ---
 .../server/controllers/tests/test_install.py  | 38 +++++++++++++++++++
 3 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/subiquity/server/controllers/install.py b/subiquity/server/controllers/install.py
index 4240eda2..26dde35e 100644
--- a/subiquity/server/controllers/install.py
+++ b/subiquity/server/controllers/install.py
@@ -19,6 +19,7 @@ import glob
 import json
 import logging
 import os
+import platform
 import re
 import shlex
 import shutil
@@ -666,6 +667,17 @@ class InstallController(SubiquityController):
             )
             raise
 
+    async def platform_postinstall(self):
+        """Run architecture specific commands/quirks"""
+        if platform.machine() == "s390x":
+            try:
+                # Ensure we boot from the installed system.
+                await arun_command(["chreipl", "/target/boot"])
+            except subprocess.CalledProcessError as cpe:
+                if cpe.stderr is not None:
+                    log.warning("chreipl stderr:\n%s", cpe.stderr)
+                raise
+
     @with_context(
         description="final system configuration", level="INFO", childlevel="DEBUG"
     )
@@ -700,6 +712,7 @@ class InstallController(SubiquityController):
                     hostname = f.read().strip()
 
             await self.app.controllers.Ad.join_domain(hostname, context)
+        await self.platform_postinstall()
         self.model.filesystem.copy_artifacts_to_target()
 
     @with_context(description="configuring cloud-init")
diff --git a/subiquity/server/controllers/shutdown.py b/subiquity/server/controllers/shutdown.py
index 7cf05304..971bbe41 100644
--- a/subiquity/server/controllers/shutdown.py
+++ b/subiquity/server/controllers/shutdown.py
@@ -16,7 +16,6 @@
 import asyncio
 import logging
 import os
-import platform
 import subprocess
 
 from subiquity.common.apidef import API
@@ -137,11 +136,6 @@ class ShutdownController(SubiquityController):
         if self.opts.dry_run:
             self.app.exit()
         else:
-            # On shutdown, it seems that the asynchronous command runners are not
-            # reliable.  Best to keep using the traditional sort.
-            if self.app.state == ApplicationState.DONE:
-                if platform.machine() == "s390x":
-                    run_command(["chreipl", "/target/boot"])
             if self.mode == ShutdownMode.REBOOT:
                 run_command(["/sbin/reboot"])
             elif self.mode == ShutdownMode.POWEROFF:
diff --git a/subiquity/server/controllers/tests/test_install.py b/subiquity/server/controllers/tests/test_install.py
index 20197584..00c91e01 100644
--- a/subiquity/server/controllers/tests/test_install.py
+++ b/subiquity/server/controllers/tests/test_install.py
@@ -350,3 +350,41 @@ class TestInstallController(unittest.IsolatedAsyncioTestCase):
         self.assertIn("--fs-uuid fsuuid", cfg)
         self.assertIn("rp-partuuid=partuuid", cfg)
         self.assertIn("uuid=casper-uuid", cfg)
+
+    @patch("platform.machine", return_value="s390x")
+    @patch("subiquity.server.controllers.install.arun_command")
+    async def test_postinstall_platform_s390x(self, arun, machine):
+        await self.controller.platform_postinstall()
+        arun.assert_called_once_with(["chreipl", "/target/boot"])
+
+    @patch("platform.machine", return_value="s390x")
+    async def test_postinstall_platform_s390x_fail(self, machine):
+        cpe = subprocess.CalledProcessError(
+            cmd=["chreipl", "/target/boot"],
+            returncode=1,
+            stderr="chreipl: No valid target specified",
+        )
+
+        arun_patch = patch(
+            "subiquity.server.controllers.install.arun_command", side_effect=cpe
+        )
+
+        assert_logs = self.assertLogs(
+            "subiquity.server.controllers.install", level="WARNING"
+        )
+        assert_raises = self.assertRaises(subprocess.CalledProcessError)
+
+        with assert_logs as errors, assert_raises, arun_patch as arun:
+            await self.controller.platform_postinstall()
+
+        arun.assert_called_once_with(["chreipl", "/target/boot"])
+        self.assertIn(
+            ("chreipl stderr:\n%s", ["chreipl: No valid target specified"]),
+            [(record.msg, list(record.args)) for record in errors.records],
+        )
+
+    @patch("platform.machine", return_value="amd64")
+    @patch("subiquity.server.controllers.install.arun_command")
+    async def test_postinstall_platform_amd64(self, arun, machine):
+        await self.controller.platform_postinstall()
+        arun.assert_not_called()