Merge PR #14046: make critical sections safe in the presence of exceptions

Reviewed-by: ejgallego
Ack-by: gares
Ack-by: SkySkimmer
Ack-by: gadmm

3 files changed, 32 insertions, 10 deletions

diff --git a/clib/cThread.ml b/clib/cThread.ml
index 89ca2f7d83..3796fdf788 100644
--- a/clib/cThread.ml
+++ b/clib/cThread.ml
@@ -107,3 +107,22 @@ let mask_sigalrm f x =
 
 let create f x =
   Thread.create (mask_sigalrm f) x
+
+(*
+  Atomic mutex lock taken from https://gitlab.com/gadmm/memprof-limits/-/blob/master/src/thread_map.ml#L23-34
+  Critical sections :
+   - Mutex.lock does not poll on leaving the blocking section
+     since 4.12.
+   - Never inline, to avoid theoretically-possible reorderings with
+     flambda.
+     (workaround to the lack of masking)
+*)
+
+(* We inline the call to Mutex.unlock to avoid polling in bytecode mode *)
+external unlock: Mutex.t -> unit = "caml_mutex_unlock"
+
+let[@inline never] with_lock m ~scope =
+  let () = Mutex.lock m (* BEGIN ATOMIC *) in
+  match (* END ATOMIC *) scope () with
+  | (* BEGIN ATOMIC *) x -> unlock m ; (* END ATOMIC *) x
+  | (* BEGIN ATOMIC *) exception e -> unlock m ; (* END ATOMIC *) raise e
diff --git a/clib/cThread.mli b/clib/cThread.mli
index 87889f3356..d974135d43 100644
--- a/clib/cThread.mli
+++ b/clib/cThread.mli
@@ -29,3 +29,8 @@ val thread_friendly_really_read_line : thread_ic -> string
 (* Wrapper around Thread.create that blocks signals such as Sys.sigalrm (used
  * for Timeout *)
 val create : ('a -> 'b) -> 'a -> Thread.t
+
+(*
+  Atomic mutex lock taken from https://gitlab.com/gadmm/memprof-limits/-/blob/master/src/thread_map.ml#L23-34
+*)
+val with_lock : Mutex.t -> scope:(unit -> 'a) -> 'a
diff --git a/clib/exninfo.ml b/clib/exninfo.ml
index 07b7f47529..4c1f47df30 100644
--- a/clib/exninfo.ml
+++ b/clib/exninfo.ml
@@ -71,10 +71,9 @@ let record_backtrace b =
 let get_backtrace e = get e backtrace_info
 
 let iraise (e,i) =
-  let () = Mutex.lock lock in
-  let id = Thread.id (Thread.self ()) in
-  let () = current := (id, (e,i)) :: remove_assoc id !current in
-  let () = Mutex.unlock lock in
+  CThread.with_lock lock ~scope:(fun () ->
+      let id = Thread.id (Thread.self ()) in
+      current := (id, (e,i)) :: remove_assoc id !current);
   match get i backtrace_info with
   | None ->
     raise e
@@ -82,12 +81,11 @@ let iraise (e,i) =
     Printexc.raise_with_backtrace e bt
 
 let find_and_remove () =
-  let () = Mutex.lock lock in
-  let id = Thread.id (Thread.self ()) in
-  let (v, l) = find_and_remove_assoc id !current in
-  let () = current := l in
-  let () = Mutex.unlock lock in
-  v
+  CThread.with_lock lock ~scope:(fun () ->
+      let id = Thread.id (Thread.self ()) in
+      let (v, l) = find_and_remove_assoc id !current in
+      let () = current := l in
+      v)
 
 let info e =
   let (src, data) = find_and_remove () in