From 83f03089aaa661268e9b5432f67eeaaaa57ccfb3 Mon Sep 17 00:00:00 2001
From: Michael Hudson-Doyle <michael.hudson@canonical.com>
Date: Tue, 6 Oct 2020 17:21:51 +1300
Subject: [PATCH] handle automatically run server processes in dry-run mode
 better

otherwise, when we start adding server restarts to the mix, it is easy to
end up with multiple server processes running
---
 subiquity/cmd/tui.py | 14 +++++---------
 subiquity/core.py    | 18 +++++++++++-------
 2 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/subiquity/cmd/tui.py b/subiquity/cmd/tui.py
index c20d6fa9..ae6780f4 100755
--- a/subiquity/cmd/tui.py
+++ b/subiquity/cmd/tui.py
@@ -102,6 +102,7 @@ def make_client_args_parser():
         '--snap-section', action='store', default='server',
         help=("Show snaps from this section of the store in the snap "
               "list screen."))
+    parser.add_argument('--server-pid')
     return parser
 
 
@@ -115,7 +116,6 @@ def main():
     from subiquity.core import Subiquity
     parser = make_client_args_parser()
     args = sys.argv[1:]
-    server_proc = None
     if '--dry-run' in args:
         opts, unknown = parser.parse_known_args(args)
         if opts.socket is None:
@@ -130,7 +130,10 @@ def main():
                 server_args
             server_proc = subprocess.Popen(
                 server_cmd, stdout=server_output, stderr=subprocess.STDOUT)
+            opts.server_pid = str(server_proc.pid)
             print("running server pid {}".format(server_proc.pid))
+        elif opts.server_pid is not None:
+            print("reconnecting to server pid {}".format(opts.server_pid))
         else:
             opts = parser.parse_args(args)
     else:
@@ -217,20 +220,13 @@ def main():
             opts.answers = None
 
     subiquity_interface = Subiquity(opts, block_log_dir)
-    subiquity_interface.server_proc = server_proc
 
     subiquity_interface.note_file_for_apport(
         "InstallerLog", logfiles['debug'])
     subiquity_interface.note_file_for_apport(
         "InstallerLogInfo", logfiles['info'])
 
-    try:
-        subiquity_interface.run()
-    finally:
-        if server_proc is not None:
-            print('killing server {}'.format(server_proc.pid))
-            server_proc.send_signal(2)
-            server_proc.wait()
+    subiquity_interface.run()
 
 
 if __name__ == '__main__':
diff --git a/subiquity/core.py b/subiquity/core.py
index dc4434e9..35b09531 100644
--- a/subiquity/core.py
+++ b/subiquity/core.py
@@ -219,19 +219,17 @@ class Subiquity(TuiApplication):
             self.ui.block_input = True
             self.aio_loop.create_task(self._restart_server())
             return
-        if remove_last_screen:
-            self._remove_last_screen()
         if self.urwid_loop is not None:
             self.urwid_loop.stop()
         cmdline = ['snap', 'run', 'subiquity']
         if self.opts.dry_run:
-            if self.server_proc is not None and not restart_server:
-                print('killing server {}'.format(self.server_proc.pid))
-                self.server_proc.send_signal(2)
-                self.server_proc.wait()
             cmdline = [
                 sys.executable, '-m', 'subiquity.cmd.tui',
-                ] + sys.argv[1:]
+                ] + sys.argv[1:] + ['--socket', self.opts.socket]
+            if self.opts.server_pid is not None:
+                cmdline.extend(['--server-pid', self.opts.server_pid])
+            log.debug("restarting %r", cmdline)
+
         os.execvp(cmdline[0], cmdline)
 
     def get_primary_tty(self):
@@ -376,6 +374,12 @@ class Subiquity(TuiApplication):
             else:
                 traceback.print_exc()
                 signal.pause()
+        finally:
+            if self.opts.server_pid:
+                print('killing server {}'.format(self.opts.server_pid))
+                pid = int(self.opts.server_pid)
+                os.kill(pid, 2)
+                os.waitpid(pid, 0)
 
     def add_event_listener(self, listener):
         self.event_listeners.append(listener)