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

@@ -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()
 
 
 if __name__ == '__main__':

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)