From 472e42717e0aa4a478bf2746c2ff42a6adc7c02e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Reynir=20Bj=C3=B6rnsson?= <reynir@reynir.dk>
Date: Mon, 30 Nov 2020 15:26:29 +0100
Subject: [PATCH] Implement killall waiting

---
 daemon/albatrossd.ml | 12 +++++-------
 src/vmm_vmmd.ml      | 16 +++++++++++-----
 src/vmm_vmmd.mli     |  2 +-
 3 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/daemon/albatrossd.ml b/daemon/albatrossd.ml
index d9953ed..38c9995 100644
--- a/daemon/albatrossd.ml
+++ b/daemon/albatrossd.ml
@@ -175,13 +175,11 @@ let jump _ systemd influx tmpdir dbdir retries enable_stats =
        let self_destruct_mutex = Lwt_mutex.create () in
        let self_destruct () =
          Lwt_mutex.with_lock self_destruct_mutex (fun () ->
-             (if Vmm_vmmd.killall !state then
-                (* not too happy about the sleep here, but cleaning up resources
-                   is really important (fifos, vm images, tap devices) - which
-                   is done asynchronously (in the task waitpid() on the pid) *)
-                Lwt_unix.sleep 1.
-              else
-                Lwt.return_unit) >>= fun () ->
+             (let state', tasks = Vmm_vmmd.killall !state Lwt.task in
+              state := state';
+              Lwt_list.iter_s (fun exit_code ->
+                  exit_code >>= fun (_ : process_exit) -> Lwt.return_unit)
+                tasks) >>= fun () ->
              Vmm_lwt.safe_close ss)
        in
        Sys.(set_signal sigterm
diff --git a/src/vmm_vmmd.ml b/src/vmm_vmmd.ml
index bca0d4a..ac77635 100644
--- a/src/vmm_vmmd.ml
+++ b/src/vmm_vmmd.ml
@@ -17,11 +17,6 @@ type 'a t = {
 
 let in_shutdown = ref false
 
-let killall t =
-  match List.map snd (Vmm_trie.all t.resources.Vmm_resources.unikernels) with
-  | [] -> false
-  | vms -> in_shutdown := true ; List.iter Vmm_unix.destroy vms ; true
-
 let remove_resources t name =
   let resources = match Vmm_resources.remove_vm t.resources name with
     | Error (`Msg e) ->
@@ -63,6 +58,17 @@ let register_restart t id create =
   | Some _ -> Logs.err (fun m -> m "restart attempted to overwrite waiter"); None
   | _ -> Some (register t id create)
 
+let killall t create =
+  let vms = Vmm_trie.all t.resources.Vmm_resources.unikernels in
+  in_shutdown := true ;
+  let t, xs = List.fold_left
+      (fun (t, acc) (id, _) ->
+         let (t, a) = register t id create in
+         (t, a :: acc))
+      (t, []) vms in
+  List.iter Vmm_unix.destroy (List.map snd vms) ;
+  t, xs
+
 let init () =
   let t = {
     console_counter = 1L ;
diff --git a/src/vmm_vmmd.mli b/src/vmm_vmmd.mli
index 787a8b1..48043d4 100644
--- a/src/vmm_vmmd.mli
+++ b/src/vmm_vmmd.mli
@@ -32,7 +32,7 @@ val handle_command : 'a t -> Vmm_commands.wire ->
    | `Wait_and_create of Name.t * (Name.t * Unikernel.config) ],
    Vmm_commands.res) result
 
-val killall : 'a t -> bool
+val killall : 'a t -> (unit -> 'b * 'a) -> 'a t * ('b list)
 
 val restore_unikernels : unit -> (Unikernel.config Vmm_trie.t, [> `Msg of string ]) result