Add Picos_select.await_on for convenience

polytypic · polytypic · commit e409f808b2fd · 2024-03-17T11:08:09.000+02:00
diff --git a/lib/picos/select/picos_select.ml b/lib/picos/select/picos_select.ml
@@ -273,8 +273,16 @@ let[@alert "-handler"] rec insert_fd s fds (Return_on r as op) =
       wakeup s `Alive
     else insert_fd s fds op
 
-let return_on computation file_descr event value =
+let return_on computation file_descr op value =
   let s = get () in
   insert_fd s
-    (match event with `R -> s.new_rd | `W -> s.new_wr | `E -> s.new_ex)
+    (match op with `R -> s.new_rd | `W -> s.new_wr | `E -> s.new_ex)
     (Return_on { computation; file_descr; value })
+
+let await_on file_descr op =
+  let computation = Computation.create () in
+  return_on computation file_descr op file_descr;
+  try Computation.await computation
+  with exn ->
+    Computation.cancel computation exit_exn_bt;
+    raise exn
diff --git a/lib/picos/select/picos_select.mli b/lib/picos/select/picos_select.mli
@@ -7,6 +7,8 @@
 
 open Picos
 
+(** {1 Timeouts} *)
+
 val cancel_after : _ Computation.t -> seconds:float -> Exn_bt.t -> unit
 (** [cancel_after computation ~seconds exn_bt] arranges for [computation] to be
     {{!Picos.Computation.cancel} canceled} with given [exn_bt] after given time
@@ -16,6 +18,8 @@ val cancel_after : _ Computation.t -> seconds:float -> Exn_bt.t -> unit
     ℹ️ You can use [cancel_after] to implement the handler for the
     {{!Picos.Computation.Cancel_after} [Cancel_after]} effect. *)
 
+(** {1 IO} *)
+
 val return_on :
   'a Computation.t -> Unix.file_descr -> [ `R | `W | `E ] -> 'a -> unit
 (** [return_on computation file_descr op value] arranges for [computation] to be
@@ -25,3 +29,7 @@ val return_on :
 
     ℹ️ Using [Unix.set_nonblock] and [return_on] you can implement direct-style
     transparently asynchronous IO on top of the [Unix] module. *)
+
+val await_on : Unix.file_descr -> [ `R | `W | `E ] -> Unix.file_descr
+(** [await_on file_descr op] awaits until [file_descr] becomes available for
+    [op]. *)
diff --git a/test/lib/elements/async_unix.ml b/test/lib/elements/async_unix.ml
@@ -1,25 +1,16 @@
-open Picos
 include Unix
 
-let await file_descr op =
-  let c = Computation.create () in
-  Picos_select.return_on c file_descr op file_descr;
-  try Computation.await c
-  with exn ->
-    Computation.cancel c (Exn_bt.get_callstack 0 exn);
-    raise exn
-
 let read file_descr bytes pos len =
   try Unix.read file_descr bytes pos len
   with Unix.Unix_error ((EAGAIN | EWOULDBLOCK), _, _) ->
-    Unix.read (await file_descr `R) bytes pos len
+    Unix.read (Picos_select.await_on file_descr `R) bytes pos len
 
 let write file_descr bytes pos len =
   try Unix.write file_descr bytes pos len
   with Unix.Unix_error ((EAGAIN | EWOULDBLOCK), _, _) ->
-    Unix.write (await file_descr `W) bytes pos len
+    Unix.write (Picos_select.await_on file_descr `W) bytes pos len
 
 let accept ?cloexec file_descr =
   try Unix.accept ?cloexec file_descr
   with Unix.Unix_error ((EAGAIN | EWOULDBLOCK), _, _) ->
-    Unix.accept ?cloexec (await file_descr `R)
+    Unix.accept ?cloexec (Picos_select.await_on file_descr `R)
diff --git a/test/lib/elements/async_unix.mli b/test/lib/elements/async_unix.mli
@@ -1,7 +1,3 @@
 (** A transparently asynchronous [Unix] module for {!Picos}. *)
 
 include module type of Unix
-
-val await : file_descr -> [ `R | `W | `E ] -> file_descr
-(** [await file_descr op] awaits for [file_descr] to be ready for [op] and
-    returns [file_descr]. *)
