refactor: move some common code to Suspend_

src/suspend_.ml

@@ -1,3 +1,5 @@
+module A = Atomic_
+
 type suspension = (unit, exn * Printexc.raw_backtrace) result -> unit
 type task = unit -> unit
 
@@ -32,9 +34,28 @@ let with_suspend ~(run : with_handler:bool -> task -> unit) (f : unit -> unit) :
 
   E.try_with f () { E.effc }
 
+(* DLA interop *)
+let prepare_for_await () : Dla_.t =
+  (* current state *)
+  let st : ((with_handler:bool -> task -> unit) * suspension) option A.t =
+    A.make None
+  in
+
+  let release () : unit =
+    match A.exchange st None with
+    | None -> ()
+    | Some (run, k) -> run ~with_handler:true (fun () -> k (Ok ()))
+  and await () : unit =
+    suspend { handle = (fun ~run k -> A.set st (Some (run, k))) }
+  in
+
+  let t = { Dla_.release; await } in
+  t
+
 [@@@ocaml.alert "+unstable"]
 [@@@else_]
 
 let with_suspend ~run:_ f = f ()
+let prepare_for_await () = { Dla_.release = ignore; await = ignore }
 
 [@@@endif]

src/suspend_.mli

@@ -50,6 +50,9 @@ val suspend : suspension_handler -> unit
 
 [@@@endif]
 
+val prepare_for_await : unit -> Dla_.t
+(** Our stub for DLA. Unstable. *)
+
 val with_suspend :
   run:(with_handler:bool -> task -> unit) -> (unit -> unit) -> unit
 (** [with_suspend ~run f] runs [f()] in an environment where [suspend]