More fine grained mutex_condvar for channels

Author: ctk21

lib/chan.ml

@@ -1,38 +1,46 @@
+type mutex_condvar = {
+  mutex: Domain.Mutex.t;
+  condition: Domain.Condition.t
+}
+
+type waiting_notified =
+  | Waiting
+  | Notified
+
 type 'a contents =
-  | Empty of {receivers: ('a option ref * Domain.Condition.t) Fun_queue.t}
-  | NotEmpty of {senders: ('a * Domain.Condition.t) Fun_queue.t; messages: 'a Fun_queue.t}
+  | Empty of {receivers: ('a option ref * mutex_condvar) Fun_queue.t}
+  | NotEmpty of {senders: ('a * waiting_notified ref * mutex_condvar) Fun_queue.t; messages: 'a Fun_queue.t}
 
 type 'a t = {
-  mutex: Domain.Mutex.t;
   buffer_size: int option;
   contents: 'a contents Atomic.t
 }
 
-let condition_key : Domain.Condition.t Domain.DLS.key = Domain.DLS.new_key ()
+let mutex_condvar_key : mutex_condvar Domain.DLS.key = Domain.DLS.new_key ()
 
-let get_condvar m =
-  match Domain.DLS.get condition_key with
+let get_mutex_condvar () =
+  match Domain.DLS.get mutex_condvar_key with
   | None ->
+      let m = Domain.Mutex.create () in
       let c = Domain.Condition.create m in
-      Domain.DLS.set condition_key c;
-      c
-  | Some c -> c
+      let mc = {mutex=m; condition=c} in
+      Domain.DLS.set mutex_condvar_key mc;
+      mc
+  | Some mc -> mc
 
 let make_bounded n =
   if n < 0 then raise (Invalid_argument "Chan.make_bounded") ;
   {buffer_size= Some n;
-   contents = Atomic.make (Empty {receivers= Fun_queue.empty; });
-   mutex = Domain.Mutex.create ();}
+   contents = Atomic.make (Empty {receivers= Fun_queue.empty; })}
 
 let make_unbounded () =
   {buffer_size= None;
-   contents = Atomic.make (Empty {receivers= Fun_queue.empty});
-   mutex = Domain.Mutex.create ();}
+   contents = Atomic.make (Empty {receivers= Fun_queue.empty})}
 
 (* [send'] is shared by both the blocking and polling versions. Returns a
  * boolean indicating whether the send was successful. Hence, it always returns
  * [true] if [polling] is [false]. *)
-let send' {mutex; buffer_size; contents} v ~polling =
+let send' {buffer_size; contents} v ~polling =
   let open Fun_queue in
   let rec loop () =
     let old_contents = Atomic.get contents in
@@ -48,15 +56,21 @@ let send' {mutex; buffer_size; contents} v ~polling =
             begin if not polling then begin
               (* The channel is empty (no senders), no waiting receivers,
                 * buffer size is 0 and we're not polling *)
-              let condition = get_condvar mutex in
+              let mc = get_mutex_condvar () in
+              let cond_slot = ref Waiting in
               let new_contents =
                 NotEmpty
-                  {messages= empty; senders= push empty (v, condition)}
+                  {messages= empty; senders= push empty (v, cond_slot, mc)}
               in
-              Domain.Mutex.lock mutex;
               if Atomic.compare_and_set contents old_contents new_contents
-              then (Domain.Condition.wait condition; true)
-              else (Domain.Mutex.unlock mutex; loop ())
+              then begin
+                Domain.Mutex.lock mc.mutex;
+                while !cond_slot = Waiting do
+                  Domain.Condition.wait mc.condition
+                done;
+                Domain.Mutex.unlock mc.mutex;
+                true
+              end else loop ()
             end else
               (* The channel is empty (no senders), no waiting receivers,
                 * buffer size is 0 and we're polling *)
@@ -71,18 +85,18 @@ let send' {mutex; buffer_size; contents} v ~polling =
             if Atomic.compare_and_set contents old_contents new_contents
             then true
             else loop ()
-      | Some ((r, condition), receivers') ->
+      | Some ((r, mc), receivers') ->
           (* The channel is empty (no senders) and there are waiting receivers
            * *)
           let new_contents = Empty {receivers= receivers'} in
-          Domain.Mutex.lock mutex;
           if Atomic.compare_and_set contents old_contents new_contents
           then begin
+            Domain.Mutex.lock mc.mutex;
             r := Some v;
-            Domain.Condition.signal condition;
-            Domain.Mutex.unlock mutex;
+            Domain.Condition.signal mc.condition;
+            Domain.Mutex.unlock mc.mutex;
             true
-           end else (Domain.Mutex.unlock mutex; loop ())
+           end else loop ()
     end
     | NotEmpty {senders; messages} ->
         (* The channel is not empty *)
@@ -91,19 +105,19 @@ let send' {mutex; buffer_size; contents} v ~polling =
           begin if not polling then
             (* The channel is not empty, the buffer is full and we're not
               * polling *)
-            let condition = get_condvar mutex in
+            let cond_slot = ref Waiting in
+            let mc = get_mutex_condvar () in
             let new_contents =
-              NotEmpty {senders= push senders (v, condition); messages}
+              NotEmpty {senders= push senders (v, cond_slot, mc); messages}
             in
-            Domain.Mutex.lock mutex;
             if Atomic.compare_and_set contents old_contents new_contents then begin
-              Domain.Condition.wait condition;
-              Domain.Mutex.unlock mutex;
+              Domain.Mutex.lock mc.mutex;
+              while !cond_slot = Waiting do
+                Domain.Condition.wait mc.condition;
+              done;
+              Domain.Mutex.unlock mc.mutex;
               true
-            end else begin
-              Domain.Mutex.unlock mutex;
-              loop ()
-            end
+            end else loop ()
           else
             (* The channel is not empty, the buffer is full and we're
               * polling *)
@@ -129,7 +143,7 @@ let send_poll c v = send' c v ~polling:true
 (* [recv'] is shared by both the blocking and polling versions. Returns a an
  * optional value indicating whether the receive was successful. Hence, it
  * always returns [Some v] if [polling] is [false]. *)
-let recv' {mutex; buffer_size; contents} ~polling =
+let recv' {buffer_size; contents} ~polling =
   let open Fun_queue in
   let rec loop () =
     let old_contents = Atomic.get contents in
@@ -139,16 +153,19 @@ let recv' {mutex; buffer_size; contents} ~polling =
         if not polling then begin
           (* The channel is empty (no senders), and we're not polling *)
           let msg_slot = ref None in
-          let condition = get_condvar mutex in
+          let mc = get_mutex_condvar () in
           let new_contents =
-            Empty {receivers= push receivers (msg_slot, condition)}
+            Empty {receivers= push receivers (msg_slot, mc)}
           in
-          Domain.Mutex.lock mutex;
           if Atomic.compare_and_set contents old_contents new_contents then
-            (Domain.Condition.wait condition;
-             Domain.Mutex.unlock mutex;
-             !msg_slot)
-          else (Domain.Mutex.unlock mutex; loop ())
+          begin
+            Domain.Mutex.lock mc.mutex;
+            while !msg_slot = None do
+              Domain.Condition.wait mc.condition;
+            done;
+            Domain.Mutex.unlock mc.mutex;
+            !msg_slot
+          end else loop ()
         end else
           (* The channel is empty (no senders), and we're polling *)
           None
@@ -170,7 +187,7 @@ let recv' {mutex; buffer_size; contents} ~polling =
             if Atomic.compare_and_set contents old_contents new_contents
             then Some m
             else loop ()
-        | None, Some ((m, condition), senders') ->
+        | None, Some ((m, c, mc), senders') ->
             (* The channel is not empty, there are no messages, and there
               * is a waiting sender. This is only possible is the buffer
               * size is 0. *)
@@ -181,26 +198,28 @@ let recv' {mutex; buffer_size; contents} ~polling =
               else
                 NotEmpty {messages; senders= senders'}
             in
-            Domain.Mutex.lock mutex;
             if Atomic.compare_and_set contents old_contents new_contents
             then begin
-              Domain.Condition.signal condition;
-              Domain.Mutex.unlock mutex;
+              Domain.Mutex.lock mc.mutex;
+              c := Notified;
+              Domain.Condition.signal mc.condition;
+              Domain.Mutex.unlock mc.mutex;
               Some m
-            end else (Domain.Mutex.unlock mutex; loop ())
-        | Some (m, messages'), Some ((ms, condition), senders') ->
+            end else loop ()
+        | Some (m, messages'), Some ((ms, sc, mc), senders') ->
             (* The channel is not empty, there is a message, and there is a
               * waiting sender. *)
             let new_contents =
               NotEmpty {messages= push messages' ms; senders= senders'}
             in
-            Domain.Mutex.lock mutex;
             if Atomic.compare_and_set contents old_contents new_contents
             then begin
-              Domain.Condition.signal condition;
-              Domain.Mutex.unlock mutex;
+              Domain.Mutex.lock mc.mutex;
+              sc := Notified;
+              Domain.Condition.signal mc.condition;
+              Domain.Mutex.unlock mc.mutex;
               Some m
-            end else (Domain.Mutex.unlock mutex; loop ())
+            end else loop ()
   in
   loop ()