CA-361220: xenopsd: introduce TASK.destroy_on_finish

There are certain tasks that are run asynchronously without anyone
waiting for the result (e.g. import_metadata_async).
Allow setting a flag on these tasks, so that they are cleaned up when
finished (either successfully or not).

Needed to avoid space leaks due to an ever growing tasks/updates list.

Signed-off-by: Edwin Török <[email protected]>

Author: robhoes

lib/task_server.ml

@@ -93,6 +93,7 @@ functor
       ; test_cancel_at: int option  (** index of the cancel point to trigger *)
       ; mutable backtrace: Backtrace.t  (** on error, a backtrace *)
       ; mutable cancellable: bool
+      ; mutable destroy_on_finish: bool
     }
 
     and tasks = {
@@ -149,26 +150,44 @@ functor
             )
         ; backtrace= Backtrace.empty
         ; cancellable= true
+        ; destroy_on_finish= false
         }
       in
       Mutex.execute tasks.m (fun () ->
           tasks.task_map := SMap.add t.id t !(tasks.task_map)) ;
       t
 
+    (* Remove the task from the id -> task mapping. NB any active thread will
+       still continue. *)
+    let destroy task =
+      let tasks = task.tasks in
+      Mutex.execute tasks.m (fun () ->
+          tasks.task_map := SMap.remove task.id !(tasks.task_map)
+      )
+
+    let task_finished item =
+      if item.destroy_on_finish then (
+        debug "Auto-destroying task %s" item.id ;
+        destroy item
+      )
+
     (* [run t] executes the task body, updating the fields of [t] *)
     let run item =
-      try
-        let start = Unix.gettimeofday () in
-        let result = item.f item in
-        let duration = Unix.gettimeofday () -. start in
-        item.state <- Interface.Task.Completed {Interface.Task.duration; result} ;
-        debug "Task %s completed; duration = %.0f" item.id duration
-      with e ->
-        Backtrace.is_important e ;
-        error "Task %s failed; %s" item.id (Printexc.to_string e) ;
-        item.backtrace <- Backtrace.remove e ;
-        let e = e |> Interface.marshal_exn in
-        item.state <- Interface.Task.Failed e
+      ( try
+          let start = Unix.gettimeofday () in
+          let result = item.f item in
+          let duration = Unix.gettimeofday () -. start in
+          item.state <-
+            Interface.Task.Completed {Interface.Task.duration; result} ;
+          debug "Task %s completed; duration = %.0f" item.id duration
+        with e ->
+          Backtrace.is_important e ;
+          error "Task %s failed; %s" item.id (Printexc.to_string e) ;
+          item.backtrace <- Backtrace.remove e ;
+          let e = e |> Interface.marshal_exn in
+          item.state <- Interface.Task.Failed e
+      ) ;
+      task_finished item
 
     let find_locked tasks id =
       try SMap.find id !(tasks.task_map)
@@ -219,13 +238,6 @@ functor
       Mutex.execute tasks.m (fun () ->
           SMap.bindings !(tasks.task_map) |> List.map snd)
 
-    (* Remove the task from the id -> task mapping. NB any active thread will
-       still continue. *)
-    let destroy task =
-      let tasks = task.tasks in
-      Mutex.execute tasks.m (fun () ->
-          tasks.task_map := SMap.remove task.id !(tasks.task_map))
-
     let cancel task =
       let callbacks =
         Mutex.execute task.tm (fun () ->
@@ -286,5 +298,19 @@ functor
           (* If task is cancelling, just cancel it before setting it to not
              cancellable *)
           check_cancelling_locked task ;
-          task.cancellable <- false)
+          task.cancellable <- false
+      )
+
+    let destroy_on_finish t =
+      t.destroy_on_finish <- true ;
+      let already_finished =
+        Mutex.execute t.tm @@ fun () ->
+        t.destroy_on_finish <- true ;
+        match t.state with
+        | Interface.Task.Pending _ ->
+            false
+        | Interface.Task.Completed _ | Interface.Task.Failed _ ->
+            true
+      in
+      if already_finished then task_finished t
   end

lib/task_server.mli

@@ -113,4 +113,9 @@ module Task : functor (Interface : INTERFACE) -> sig
 
   (* Set a task not cancellable *)
   val prohibit_cancellation : task_handle -> unit
+
+  (* When the task finishes automatically destroy it to avoid resource leaks.
+     Useful for asynchronous tasks that we don't wait for.
+  *)
+  val destroy_on_finish : task_handle -> unit
 end

xen/xenops_interface.ml

@@ -551,6 +551,11 @@ module XenopsAPI (R : RPC) = struct
       declare "Task.list"
         ["List all the current tasks"]
         (debug_info_p @-> returning task_list_p err)
+
+    let destroy_on_finish =
+      declare "Task.destroy_on_finish"
+        ["Ensures the task will be destroyed when it finishes"]
+        (debug_info_p @-> task_id_p @-> returning unit_p err)
   end
 
   module HOST = struct