fix preemptive test: but needs new primitive in preemptive

raphael-proust · raphael-proust · commit 205581f8832c · 2025-09-15T11:15:08.000+02:00
diff --git a/src/unix/lwt_preemptive.ml b/src/unix/lwt_preemptive.ml
@@ -43,14 +43,15 @@ sig
   type 'a t
 
   val make : unit -> 'a t
-  val get : 'a t -> 'a
+  val get : 'a t -> ('a, unit) result
   val set : 'a t -> 'a -> unit
+  val kill : 'a t -> unit
 end =
 struct
   type 'a t = {
     m  : Mutex.t;
     cv : Condition.t;
-    mutable cell : 'a option;
+    mutable cell : ('a, unit) result option;
   }
 
   let make () = { m = Mutex.create (); cv = Condition.create (); cell = None }
@@ -71,7 +72,13 @@ struct
 
   let set t v =
     Mutex.lock t.m;
-    t.cell <- Some v;
+    t.cell <- Some (Ok v);
+    Mutex.unlock t.m;
+    Condition.signal t.cv
+
+  let kill t =
+    Mutex.lock t.m;
+    t.cell <- Some (Error ());
     Mutex.unlock t.m;
     Condition.signal t.cv
 end
@@ -97,14 +104,16 @@ let waiters : thread Lwt.u Lwt_sequence.t Domain.DLS.key = Domain.DLS.new_key Lw
 
 (* Code executed by a worker: *)
 let rec worker_loop worker =
-  let id, task = CELL.get worker.task_cell in
-  task ();
-  (* If there is too much threads, exit. This can happen if the user
-     decreased the maximum: *)
-  if Domain.DLS.get threads_count > Domain.DLS.get max_threads then worker.reuse <- false;
-  (* Tell the main thread that work is done: *)
-  Lwt_unix.send_notification id;
-  if worker.reuse then worker_loop worker
+  match CELL.get worker.task_cell with
+  | Error () -> ()
+  | Ok (id, task) ->
+      task ();
+      (* If there is too much threads, exit. This can happen if the user
+         decreased the maximum: *)
+      if Domain.DLS.get threads_count > Domain.DLS.get max_threads then worker.reuse <- false;
+      (* Tell the main thread that work is done: *)
+      Lwt_unix.send_notification id;
+      if worker.reuse then worker_loop worker
 
 (* create a new worker: *)
 let make_worker () =
@@ -258,10 +267,14 @@ let run_in_domain d f =
   run_in_domain_dont_wait d job;
   (* Wait for the result. *)
   match CELL.get cell with
-  | Result.Ok ret -> ret
-  | Result.Error exn -> raise exn
+  | Ok (Ok ret) -> ret
+  | Ok (Error exn) -> raise exn
+  | Error () -> assert false
 
 (* This version shadows the one above, adding an exception handler *)
 let run_in_domain_dont_wait d f handler =
   let f () = Lwt.catch f (fun exc -> handler exc; Lwt.return_unit) in
   run_in_domain_dont_wait d f
+
+let kill_all () =
+  Queue.iter (fun thread -> CELL.kill thread.task_cell) (Domain.DLS.get workers)
diff --git a/src/unix/lwt_preemptive.mli b/src/unix/lwt_preemptive.mli
@@ -87,3 +87,7 @@ val get_max_number_of_threads_queued : unit -> int
 val nbthreads : unit -> int
 val nbthreadsbusy : unit -> int
 val nbthreadsqueued : unit -> int
+
+(* kill_all is to be called before joining the domain, not satisfying UI for
+   now, searching for a better way *)
+val kill_all : unit -> unit
diff --git a/test/multidomain/preempting.ml b/test/multidomain/preempting.ml
@@ -12,10 +12,7 @@ let simulate_work data =
 
 let () = Lwt_unix.init_domain ()
 
-(* atomic just for debugging: to record when domains are finished *)
-let x = Atomic.make 0
-
-let domain_go_brrrrrrr n input = Domain.spawn (fun () ->
+let domain_go_brrrrrrr input = Domain.spawn (fun () ->
   flush_all ();
   Lwt_unix.init_domain ();
   let v = Lwt_main.run (
@@ -24,35 +21,23 @@ let domain_go_brrrrrrr n input = Domain.spawn (fun () ->
       Lwt_list.map_p (Lwt_preemptive.detach simulate_work) input
     )
   in
-  (* printing just for debug: to see when different domains are finished *)
-  Printf.printf "domain #%d %d scheduler returned\n" n (Domain.self () :> int);
-  flush_all ();
-  Atomic.incr x;
+  Lwt_preemptive.kill_all ();
   v
 )
 
 let () =
-  let rec go n acc = function
+  let rec go acc = function
     | [_] | [] ->
-        Printf.printf "all domain started\n"; flush_all ();
         acc
     | (_ :: more) as wrk ->
         let expected = List.map String.length wrk in
-        let acc = (n, expected, domain_go_brrrrrrr n wrk) :: acc in
-        go (n + 1) acc more
-  in
-  let results = go 1 [] input in
-  Unix.sleepf 5.; (* sleeping for debug: to observe that atomic is at max value *)
-  Printf.printf "done debug-sleeping, about to join all domains (atomic=%d)\n" (Atomic.get x); flush_all ();
-  let results = List.map (fun (n, e, d) ->
-    Printf.printf "joining domain #%d (atomic=%d)\n" n (Atomic.get x); flush_all ();
-    let d = Domain.join d in
-    Printf.printf "joined domain #%d (atomic=%d)\n" n (Atomic.get x); flush_all ();
-    (e, d)) results
+        let acc = (expected, domain_go_brrrrrrr wrk) :: acc in
+        go acc more
   in
+  let results = go [] input in
   let success =
     List.for_all
-      (fun (expected, d) -> List.for_all2 Int.equal expected d)
+      (fun (expected, d) -> List.for_all2 Int.equal expected (Domain.join d))
       results
   in
   let code =
