Make critical sections safe in the presence of exceptions

We introduce the `with_lock` combinator that locks a mutex in an atomic fashion.
This ensures that exceptions thrown by signals will not leave the system in a
deadlocked state.

9 files changed, 96 insertions, 92 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
diff --git a/ide/coqide/idetop.ml b/ide/coqide/idetop.ml
index a6a7f7d742..2fe57ec66c 100644
--- a/ide/coqide/idetop.ml
+++ b/ide/coqide/idetop.ml
@@ -490,11 +490,11 @@ let eval_call c =
 let print_xml =
   let m = Mutex.create () in
   fun oc xml ->
-    Mutex.lock m;
-    if !Flags.xml_debug then
-      Printf.printf "SENT --> %s\n%!" (Xml_printer.to_string_fmt xml);
-    try Control.protect_sigalrm (Xml_printer.print oc) xml; Mutex.unlock m
-    with e -> let e = Exninfo.capture e in Mutex.unlock m; Exninfo.iraise e
+    CThread.with_lock m ~scope:(fun () ->
+        if !Flags.xml_debug then
+          Printf.printf "SENT --> %s\n%!" (Xml_printer.to_string_fmt xml);
+        try Control.protect_sigalrm (Xml_printer.print oc) xml
+        with e -> let e = Exninfo.capture e in Exninfo.iraise e)
 
 let slave_feeder fmt xml_oc msg =
   let xml = Xmlprotocol.(of_feedback fmt msg) in
diff --git a/lib/future.ml b/lib/future.ml
index 23d089fb6b..247b139798 100644
--- a/lib/future.ml
+++ b/lib/future.ml
@@ -112,8 +112,8 @@ let create_delegate ?(blocking=true) ~name fix_exn =
     if not blocking then (fun () -> raise (NotReady name)), ignore else
     let lock = Mutex.create () in
     let cond = Condition.create () in
-    (fun () -> Mutex.lock lock; Condition.wait cond lock; Mutex.unlock lock),
-    (fun () -> Mutex.lock lock; Condition.broadcast cond; Mutex.unlock lock) in
+    (fun () -> CThread.with_lock lock ~scope:(fun () -> Condition.wait cond lock)),
+    (fun () -> CThread.with_lock lock ~scope:(fun () -> Condition.broadcast cond)) in
   let ck = create ~name ~fix_exn (Delegated wait) in
   ck, assignment signal ck
 
diff --git a/lib/remoteCounter.ml b/lib/remoteCounter.ml
index 5e1150146e..9ea751eef9 100644
--- a/lib/remoteCounter.ml
+++ b/lib/remoteCounter.ml
@@ -28,10 +28,10 @@ let new_counter ~name a ~incr ~build =
          managers (that are threads) and the main thread, hence the mutex *)
     if Flags.async_proofs_is_worker () then
       CErrors.anomaly(Pp.str"Slave processes must install remote counters.");
-    Mutex.lock m; let x = f () in Mutex.unlock m;
+    let x = CThread.with_lock m ~scope:f in
     build x in
   let mk_thsafe_remote_getter f () =
-    Mutex.lock m; let x = f () in Mutex.unlock m; x in
+    CThread.with_lock m ~scope:f in
   let getter = ref(mk_thsafe_local_getter (fun () -> !data := incr !!data; !!data)) in
   let installer f =
     if not (Flags.async_proofs_is_worker ()) then
diff --git a/stm/stm.ml b/stm/stm.ml
index 9480bbdc2e..f5768726c3 100644
--- a/stm/stm.ml
+++ b/stm/stm.ml
@@ -98,8 +98,7 @@ let forward_feedback, forward_feedback_hook =
   let m = Mutex.create () in
   Hook.make ~default:(function
     | { doc_id = did; span_id = id; route; contents } ->
-        try Mutex.lock m; feedback ~did ~id ~route contents; Mutex.unlock m
-        with e -> Mutex.unlock m; raise e) ()
+        CThread.with_lock m ~scope:(fun () -> feedback ~did ~id ~route contents)) ()
 
 let unreachable_state, unreachable_state_hook = Hook.make
  ~default:(fun ~doc:_ _ _ -> ()) ()
@@ -758,17 +757,16 @@ end = struct (* {{{ *)
     let worker = ref None
 
     let set_last_job j =
-      Mutex.lock m;
-      job := Some j;
-      Condition.signal c;
-      Mutex.unlock m
+      CThread.with_lock m ~scope:(fun () ->
+          job := Some j;
+          Condition.signal c)
 
     let get_last_job () =
-      Mutex.lock m;
-      while Option.is_empty !job do Condition.wait c m; done;
-      match !job with
-      | None -> assert false
-      | Some x -> job := None; Mutex.unlock m; x
+      CThread.with_lock m ~scope:(fun () ->
+          while Option.is_empty !job do Condition.wait c m; done;
+          match !job with
+          | None -> assert false
+          | Some x -> job := None; x)
 
     let run_command () =
       try while true do get_last_job () () done
diff --git a/stm/tQueue.ml b/stm/tQueue.ml
index 2aaca85582..f5bd726dde 100644
--- a/stm/tQueue.ml
+++ b/stm/tQueue.ml
@@ -68,61 +68,54 @@ let create () = {
 
 let pop ?(picky=(fun _ -> true)) ?(destroy=ref false) tq =
   let { queue = q; lock = m; cond = c; cond_waiting = cn } = tq in
-  Mutex.lock m;
-  if tq.release then (Mutex.unlock m; raise BeingDestroyed);
-  while not (PriorityQueue.exists picky q || !destroy) do
-    tq.nwaiting <- tq.nwaiting + 1;
-    Condition.broadcast cn;
-    Condition.wait c m;
-    tq.nwaiting <- tq.nwaiting - 1;
-    if tq.release || !destroy then (Mutex.unlock m; raise BeingDestroyed)
-  done;
-  if !destroy then (Mutex.unlock m; raise BeingDestroyed);
-  let x = PriorityQueue.pop ~picky q in
-  Condition.signal c;
-  Condition.signal cn;
-  Mutex.unlock m;
-  x
+  CThread.with_lock m ~scope:(fun () ->
+      if tq.release then raise BeingDestroyed;
+      while not (PriorityQueue.exists picky q || !destroy) do
+        tq.nwaiting <- tq.nwaiting + 1;
+        Condition.broadcast cn;
+        Condition.wait c m;
+        tq.nwaiting <- tq.nwaiting - 1;
+        if tq.release || !destroy then raise BeingDestroyed
+      done;
+      if !destroy then raise BeingDestroyed;
+      let x = PriorityQueue.pop ~picky q in
+      Condition.signal c;
+      Condition.signal cn;
+      x)
 
 let broadcast tq =
   let { lock = m; cond = c } = tq in
-  Mutex.lock m;
-  Condition.broadcast c;
-  Mutex.unlock m
+  CThread.with_lock m ~scope:(fun () ->
+      Condition.broadcast c)
 
 let push tq x =
   let { queue = q; lock = m; cond = c; release } = tq in
   if release then CErrors.anomaly(Pp.str
     "TQueue.push while being destroyed! Only 1 producer/destroyer allowed.");
-  Mutex.lock m;
-  PriorityQueue.push q x;
-  Condition.broadcast c;
-  Mutex.unlock m
+  CThread.with_lock m ~scope:(fun () ->
+      PriorityQueue.push q x;
+      Condition.broadcast c)
 
 let length tq =
   let { queue = q; lock = m } = tq in
-  Mutex.lock m;
-  let n = PriorityQueue.length q in
-  Mutex.unlock m;
-  n
+  CThread.with_lock m ~scope:(fun () ->
+      PriorityQueue.length q)
 
 let clear tq =
   let { queue = q; lock = m; cond = c } = tq in
-  Mutex.lock m;
-  PriorityQueue.clear q;
-  Mutex.unlock m
+  CThread.with_lock m ~scope:(fun () ->
+      PriorityQueue.clear q)
 
 let clear_saving tq f =
   let { queue = q; lock = m; cond = c } = tq in
-  Mutex.lock m;
   let saved = ref [] in
-  while not (PriorityQueue.is_empty q) do
-    let elem = PriorityQueue.pop q in
-    match f elem with
-    | Some x -> saved := x :: !saved
-    | None -> ()
-  done;
-  Mutex.unlock m;
+  CThread.with_lock m ~scope:(fun () ->
+      while not (PriorityQueue.is_empty q) do
+        let elem = PriorityQueue.pop q in
+        match f elem with
+        | Some x -> saved := x :: !saved
+        | None -> ()
+      done);
   List.rev !saved
 
 let is_empty tq = PriorityQueue.is_empty tq.queue
@@ -130,32 +123,28 @@ let is_empty tq = PriorityQueue.is_empty tq.queue
 let destroy tq =
   tq.release <- true;
   while tq.nwaiting > 0 do
-    Mutex.lock tq.lock;
-    Condition.broadcast tq.cond;
-    Mutex.unlock tq.lock;
+    CThread.with_lock tq.lock ~scope:(fun () ->
+        Condition.broadcast tq.cond)
   done;
   tq.release <- false
 
 let wait_until_n_are_waiting_and_queue_empty j tq =
-  Mutex.lock tq.lock;
-  while not (PriorityQueue.is_empty tq.queue) || tq.nwaiting < j do
-    Condition.wait tq.cond_waiting tq.lock
-  done;
-  Mutex.unlock tq.lock
+  CThread.with_lock tq.lock ~scope:(fun () ->
+      while not (PriorityQueue.is_empty tq.queue) || tq.nwaiting < j do
+        Condition.wait tq.cond_waiting tq.lock
+      done)
 
 let wait_until_n_are_waiting_then_snapshot j tq =
   let l = ref [] in
-  Mutex.lock tq.lock;
-  while not (PriorityQueue.is_empty tq.queue) do
-    l := PriorityQueue.pop tq.queue :: !l
-  done;
-  while tq.nwaiting < j do Condition.wait tq.cond_waiting tq.lock done;
-  List.iter (PriorityQueue.push tq.queue) (List.rev !l);
-  if !l <> [] then Condition.broadcast tq.cond;
-  Mutex.unlock tq.lock;
+  CThread.with_lock tq.lock ~scope:(fun () ->
+      while not (PriorityQueue.is_empty tq.queue) do
+        l := PriorityQueue.pop tq.queue :: !l
+      done;
+      while tq.nwaiting < j do Condition.wait tq.cond_waiting tq.lock done;
+      List.iter (PriorityQueue.push tq.queue) (List.rev !l);
+      if !l <> [] then Condition.broadcast tq.cond);
   List.rev !l
 
 let set_order tq rel =
-  Mutex.lock tq.lock;
-  PriorityQueue.set_rel rel tq.queue;
-  Mutex.unlock tq.lock
+  CThread.with_lock tq.lock ~scope:(fun () ->
+      PriorityQueue.set_rel rel tq.queue)
diff --git a/stm/workerPool.ml b/stm/workerPool.ml
index 45c92c3748..fef9300377 100644
--- a/stm/workerPool.ml
+++ b/stm/workerPool.ml
@@ -72,12 +72,7 @@ let worker_handshake slave_ic slave_oc =
     exit 1
 
 let locking { lock; pool = p } f =
-  try
-    Mutex.lock lock;
-    let x = f p in
-    Mutex.unlock lock;
-    x
-  with e -> Mutex.unlock lock; raise e
+  CThread.with_lock lock ~scope:(fun () -> f p)
 
 let rec create_worker extra pool priority id =
   let cancel = ref false in