From f9ba3566575566bb263ed7bb0c42a9f36c36c97f Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Wed, 31 May 2023 00:58:51 -0400
Subject: [PATCH] Fut.wait_block: a bit of spinning before blocking

---
 src/gen/gen.ml  |  7 +++++--
 src/moonpool.ml | 32 +++++++++++++++++++++++---------
 2 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/src/gen/gen.ml b/src/gen/gen.ml
index 5e5d54bb..e7a85e89 100644
--- a/src/gen/gen.ml
+++ b/src/gen/gen.ml
@@ -55,6 +55,8 @@ let get_id (self:t) : int = Thread.id self
 
 let spawn f : t =
   Thread.create f ()
+
+let relax () = Thread.yield ()
 |}
 
 let domain_post_5 =
@@ -65,8 +67,9 @@ type t = unit Domain.t
 
 let get_id (self:t) : int = (Domain.get_id self :> int)
 
-let spawn f : t =
-  Domain.spawn f
+let spawn : _ -> t = Domain.spawn
+
+let relax = Domain.cpu_relax
 |}
 
 let p_version s = Scanf.sscanf s "%d.%d" (fun x y -> x, y)
diff --git a/src/moonpool.ml b/src/moonpool.ml
index 6657f693..018f8c5d 100644
--- a/src/moonpool.ml
+++ b/src/moonpool.ml
@@ -346,15 +346,29 @@ module Fut = struct
   (* ### blocking ### *)
 
   let wait_block (self : 'a t) : 'a or_error =
-    match peek self with
-    | Some x ->
-      (* fast path *)
-      x
-    | None ->
-      (* use queue only once *)
-      let q = S_queue.create () in
-      on_result self (fun r -> S_queue.push q r);
-      S_queue.pop q
+    match A.get self.st with
+    | Done x -> x (* fast path *)
+    | Waiting _ ->
+      let real_block () =
+        (* use queue only once *)
+        let q = S_queue.create () in
+        on_result self (fun r -> S_queue.push q r);
+        S_queue.pop q
+      in
+
+      (* a bit of spinlock *)
+      let rec loop i =
+        if i = 0 then
+          real_block ()
+        else (
+          match A.get self.st with
+          | Done x -> x
+          | Waiting _ ->
+            Domain_.relax ();
+            (loop [@tailcall]) (i - 1)
+        )
+      in
+      loop 50
 
   let wait_block_exn self =
     match wait_block self with