From c0f1bbf5a1913634ee2ddf8645274f0da944f282 Mon Sep 17 00:00:00 2001
From: Michael Hudson-Doyle <michael.hudson@canonical.com>
Date: Tue, 21 Sep 2021 14:41:49 +1200
Subject: [PATCH] cope with the server going away earlier too

---
 subiquity/client/client.py | 49 +++++++++++++++++---------------------
 subiquity/cmd/tui.py       |  2 ++
 subiquity/server/server.py |  2 +-
 3 files changed, 25 insertions(+), 28 deletions(-)

diff --git a/subiquity/client/client.py b/subiquity/client/client.py
index 45e3ce1b..3627add3 100644
--- a/subiquity/client/client.py
+++ b/subiquity/client/client.py
@@ -229,6 +229,22 @@ class SubiquityClient(TuiApplication):
             answer = await run_in_thread(input)
         await self.confirm_install()
 
+    async def _status_get(self, cur=None):
+        while True:
+            try:
+                return await self.client.meta.status.GET(cur=cur)
+            except aiohttp.ClientError:
+                try:
+                    fp = open(self.state_path("server-state"))
+                except FileNotFoundError:
+                    pass
+                else:
+                    with fp:
+                        state = getattr(ApplicationState, fp.read(), None)
+                    if state == ApplicationState.EXITED:
+                        self.exit()
+                await asyncio.sleep(1)
+
     async def noninteractive_watch_app_state(self, initial_status):
         app_status = initial_status
         confirm_task = None
@@ -246,21 +262,7 @@ class SubiquityClient(TuiApplication):
                 print("An error occurred. Press enter to start a shell")
                 await run_in_thread(input)
                 os.execvp("/bin/bash", ["/bin/bash"])
-            try:
-                app_status = await self.client.meta.status.GET(
-                    cur=app_state)
-            except aiohttp.ClientError:
-                try:
-                    fp = open(self.state_path("server-state"))
-                except FileNotFoundError:
-                    pass
-                else:
-                    with fp:
-                        state = getattr(ApplicationState, fp.read(), None)
-                    if state == ApplicationState.EXITED:
-                        self.exit()
-                await asyncio.sleep(1)
-                continue
+            app_status = await self._status_get(app_state)
 
     def subiquity_event_noninteractive(self, event):
         if event['SUBIQUITY_EVENT_TYPE'] == 'start':
@@ -288,28 +290,21 @@ class SubiquityClient(TuiApplication):
                 spinner.cancel()
                 p('\x08 \n')
 
-        async def _connect():
-            while True:
-                try:
-                    return await self.client.meta.status.GET()
-                except aiohttp.ClientError:
-                    await asyncio.sleep(1)
-
-        status = await spinning_wait("connecting", _connect())
+        status = await spinning_wait("connecting", self._status_get())
         journald_listen(
             self.aio_loop,
             [status.echo_syslog_id],
             lambda e: print(e['MESSAGE']))
         if status.state == ApplicationState.STARTING_UP:
             status = await spinning_wait(
-                "starting up", self.client.meta.status.GET(cur=status.state))
+                "starting up", self._status_get(cur=status.state))
         if status.state == ApplicationState.CLOUD_INIT_WAIT:
             status = await spinning_wait(
-                "waiting for cloud-init",
-                self.client.meta.status.GET(cur=status.state))
+                "waiting for cloud-init", self._status_get(cur=status.state))
         if status.state == ApplicationState.EARLY_COMMANDS:
             print("running early commands")
-            status = await self.client.meta.status.GET(cur=status.state)
+            status = await spinning_wait(
+                "waiting for cloud-init", self._status_get(cur=status.state))
             await asyncio.sleep(0.5)
         return status
 
diff --git a/subiquity/cmd/tui.py b/subiquity/cmd/tui.py
index 8bf6901d..be01cea0 100755
--- a/subiquity/cmd/tui.py
+++ b/subiquity/cmd/tui.py
@@ -87,6 +87,8 @@ def main():
         opts, unknown = parser.parse_known_args(args)
         if opts.socket is None:
             os.makedirs('.subiquity', exist_ok=True)
+            if os.path.exists('.subiquity/run/subiquity/server-state'):
+                os.unlink('.subiquity/run/subiquity/server-state')
             sock_path = '.subiquity/socket'
             opts.socket = sock_path
             server_args = ['--dry-run', '--socket=' + sock_path] + unknown
diff --git a/subiquity/server/server.py b/subiquity/server/server.py
index 5cec6275..e4f94611 100644
--- a/subiquity/server/server.py
+++ b/subiquity/server/server.py
@@ -256,8 +256,8 @@ class SubiquityServer(Application):
         self.block_log_dir = block_log_dir
         self.cloud = None
         self.cloud_init_ok = None
-        self._state = ApplicationState.STARTING_UP
         self.state_event = asyncio.Event()
+        self.update_state(ApplicationState.STARTING_UP)
         self.interactive = None
         self.confirming_tty = ''
         self.fatal_error = None