Repository: ocaml-multicore/picos
Pull request #202 (2024-08-03 09:47:57)

Redesign resource management to use only one binding operator
This changes `Picos_finally` to use `let (let@) = (@@)` as the only binding operator.

Thanks to @c-cube and @edwintorok for discussion on this topic! 


Patch:
diff --git a/lib/picos_finally/picos_finally.ml b/lib/picos_finally/picos_finally.ml
index a8a8dd5..b6681e6 100644
--- a/lib/picos_finally/picos_finally.ml
+++ b/lib/picos_finally/picos_finally.ml
@@ -1,98 +1,119 @@
 open Picos
 
-type 'a template = { release : 'a -> unit; acquire : unit -> 'a }
-
 type ('a, _) tdt =
-  | Moved : ('a, [> `Moved ]) tdt
+  | Transferred : ('a, [> `Transferred ]) tdt
   | Borrowed : ('a, [> `Borrowed ]) tdt
-  | Released : ('a, [> `Released ]) tdt
+  | Dropped : ('a, [> `Dropped ]) tdt
   | Resource : {
       mutable resource : 'a;
       release : 'a -> unit;
-      moved_or_released : Trigger.t;
+      transferred_or_dropped : Trigger.t;
     }
       -> ('a, [> `Resource ]) tdt
 
 type 'a instance =
-  ('a, [ `Moved | `Borrowed | `Released | `Resource ]) tdt Atomic.t
+  ('a, [ `Transferred | `Borrowed | `Dropped | `Resource ]) tdt Atomic.t
+
+(* *)
 
-let[@inline never] error (case : (_, [< `Moved | `Borrowed | `Released ]) tdt) =
+let[@inline never] error
+    (case : (_, [< `Transferred | `Borrowed | `Dropped ]) tdt) =
   invalid_arg
     (match case with
-    | Moved -> "moved"
-    | Released -> "released"
+    | Transferred -> "transferred"
+    | Dropped -> "dropped"
     | Borrowed -> "borrowed")
 
 let[@inline never] check_released () =
   (* In case of cancelation we do not consider being released an error as the
      resource was released by (the |an)other party involved in the [move]. *)
   Fiber.check (Fiber.current ());
-  error Released
+  error Dropped
+
+(* *)
 
-let rec release instance =
+let rec drop instance =
   match Atomic.get instance with
-  | Moved | Released -> ()
+  | Transferred | Dropped -> ()
   | Borrowed as case -> error case
   | Resource r as before ->
-      if Atomic.compare_and_set instance before Released then begin
-        r.release r.resource;
-        Trigger.signal r.moved_or_released
+      if Atomic.compare_and_set instance before Dropped then begin
+        Trigger.signal r.transferred_or_dropped;
+        r.release r.resource
       end
-      else release instance
+      else drop instance
 
-let await_moved_or_released instance =
+(* *)
+
+let await_transferred_or_dropped instance =
   match Atomic.get instance with
-  | Moved | Released -> ()
+  | Transferred | Dropped -> ()
   | Borrowed as case ->
-      (* This should be impossible as [let&] should have restored the state. *)
+      (* This should be impossible as [let@ _ = borrow _ in _] should have
+         restored the state. *)
       error case
   | Resource r -> begin
-      match Trigger.await r.moved_or_released with
+      match Trigger.await r.transferred_or_dropped with
       | None ->
           (* We release in case we could not wait. *)
-          release instance
+          drop instance
       | Some exn_bt ->
           (* We have been canceled, so we release. *)
-          release instance;
+          drop instance;
           Exn_bt.raise exn_bt
     end
 
-(** This function is marked [@inline never] to ensure that there are no
-    allocations between the [acquire ()] and the [match ... with] nor before
-    [release]. *)
-let[@inline never] let_at acquire body ~on_return ~on_raise =
-  let x = acquire () in
-  match body x with
-  | y ->
-      on_return x;
-      y
+let[@inline never] instantiate release acquire scope =
+  let instance =
+    Sys.opaque_identity
+      begin
+        let transferred_or_dropped = Trigger.create () in
+        let state =
+          Resource { resource = Obj.magic (); release; transferred_or_dropped }
+        in
+        Atomic.make state
+      end
+  in
+  (* After this point there must be no allocations before [acquire ()]. *)
+  let (Resource r : (_, [ `Resource ]) tdt) = Obj.magic (Atomic.get instance) in
+  r.resource <- acquire ();
+  match scope instance with
+  | result ->
+      await_transferred_or_dropped instance;
+      result
   | exception exn ->
-      on_raise x;
+      drop instance;
       raise exn
 
-let ( let^ ) t body =
-  let moved_or_released = Trigger.create () in
-  let state =
-    Resource { resource = Obj.magic (); release = t.release; moved_or_released }
-  in
-  let instance = Atomic.make state in
-  (* [acquire] is called once by [let_at] before [instance] is returned, which
-     means the [Obj.magic] use below is safe. *)
-  let acquire () =
-    let (Resource r : (_, [ `Resource ]) tdt) =
-      Obj.magic (Atomic.get instance)
-    in
-    r.resource <- t.acquire ();
-    instance
-  in
-  let_at acquire body ~on_return:await_moved_or_released ~on_raise:release
+(* *)
+
+let[@inline never] rec transfer from scope =
+  match Atomic.get from with
+  | (Transferred | Borrowed) as case -> error case
+  | Dropped -> check_released ()
+  | Resource r as before ->
+      let into = Atomic.make Transferred in
+      if Atomic.compare_and_set from before Transferred then begin
+        Trigger.signal r.transferred_or_dropped;
+        Atomic.set into before;
+        match scope into with
+        | result ->
+            await_transferred_or_dropped into;
+            result
+        | exception exn ->
+            drop into;
+            raise exn
+      end
+      else transfer from scope
 
-let rec ( let& ) instance body =
+(* *)
+
+let[@inline never] rec borrow instance scope =
   match Atomic.get instance with
-  | (Moved | Released | Borrowed) as case -> error case
+  | (Transferred | Dropped | Borrowed) as case -> error case
   | Resource r as before ->
       if Atomic.compare_and_set instance before Borrowed then begin
-        match body r.resource with
+        match scope r.resource with
         | result ->
             Atomic.set instance before;
             result
@@ -100,28 +121,37 @@ let rec ( let& ) instance body =
             Atomic.set instance before;
             raise exn
       end
-      else ( let& ) instance body
+      else borrow instance scope
 
-let move instance =
-  match Atomic.get instance with
-  | (Moved | Borrowed) as case -> error case
-  | Released -> { release = ignore; acquire = check_released }
-  | Resource r ->
-      let rec acquire () =
-        match Atomic.get instance with
-        | (Moved | Borrowed) as case -> error case
-        | Released -> check_released ()
-        | Resource r as before ->
-            if Atomic.compare_and_set instance before Moved then begin
-              (* [Trigger.signal] should never raise exceptions. *)
-              Trigger.signal r.moved_or_released;
-              r.resource
-            end
-            else acquire ()
-      in
-      { release = r.release; acquire }
-
-let[@inline] ( let@ ) t body =
-  let_at t.acquire body ~on_return:t.release ~on_raise:t.release
-
-let[@inline] finally release acquire = { release; acquire }
+(* *)
+
+let[@inline never] rec move from scope =
+  match Atomic.get from with
+  | (Transferred | Borrowed) as case -> error case
+  | Dropped -> check_released ()
+  | Resource r as before ->
+      if Atomic.compare_and_set from before Transferred then begin
+        Trigger.signal r.transferred_or_dropped;
+        match scope r.resource with
+        | result ->
+            r.release r.resource;
+            result
+        | exception exn ->
+            r.release r.resource;
+            raise exn
+      end
+      else move from scope
+
+(* *)
+
+let[@inline never] finally release acquire scope =
+  let x = acquire () in
+  match scope x with
+  | y ->
+      release x;
+      y
+  | exception exn ->
+      release x;
+      raise exn
+
+external ( let@ ) : ('a -> 'b) -> 'a -> 'b = "%apply"
diff --git a/lib/picos_finally/picos_finally.mli b/lib/picos_finally/picos_finally.mli
index b68f48a..71c2432 100644
--- a/lib/picos_finally/picos_finally.mli
+++ b/lib/picos_finally/picos_finally.mli
@@ -16,73 +16,65 @@
 
 (** {1 API} *)
 
-(** {2 Resource templates} *)
+(** {2 Basics} *)
 
-type 'resource template
-(** Represents a resource template of release and acquire functions. *)
+val ( let@ ) : ('a -> 'b) -> 'a -> 'b
+(** [let@ resource = template in scope] is equivalent to
+    [template (fun resource -> scope)].
 
-val finally : ('r -> unit) -> (unit -> 'r) -> 'r template
-(** [finally release acquire] returns a resource template with the given
-    [release] and [acquire] functions. *)
+    ℹ️ You can use this binding operator with any [template] function that has a
+    type of the form [('r -> 'a) -> 'a]. *)
 
-(** {2 Managed instances}
+val finally : ('r -> unit) -> (unit -> 'r) -> ('r -> 'a) -> 'a
+(** [finally release acquire scope] calls [acquire ()] to obtain a [resource],
+    calls [scope resource], and then calls [release resource] after the scope
+    exits. *)
 
-    A resource instance can be first acquired in one scope and then moved to
-    another scope. *)
+(** {2 Instances} *)
 
-type 'resource instance
-(** Either contains a resource or is empty as the resource has been {{!move}
-    moved}, {{!release} released}, or has been {{!let&} borrowed}
+type 'r instance
+(** Either contains a resource or is empty as the resource has been {{!transfer}
+    transferred}, {{!drop} dropped}, or has been {{!borrow} borrowed}
     temporarily. *)
 
-val ( let^ ) : 'r template -> ('r instance -> 'a) -> 'a
-(** [let^ instance = finally release acquire in scope] calls [acquire ()] to
-    obtain a resource and stores it as a resource [instance].  Then, at the end
-    of [scope], awaits that the [instance] is {{!move} moved} or {{!release}
-    released}, or {{!release} releases} the resource in case an error has been
-    raised or the fiber has been canceled.
+val instantiate : ('r -> unit) -> (unit -> 'r) -> ('r instance -> 'a) -> 'a
+(** [instantiate release acquire scope] calls [acquire ()] to obtain a resource
+    and stores it as an {!instance}, calls [scope instance].  Then, if [scope]
+    returns normally, awaits until the {!instance} becomes empty.  In case
+    [scope] raises an exception or the fiber is canceled, the instance will be
+    {{!drop} dropped}. *)
 
-    💡 In other words, you are expected to either {!move} or {!release} the
-    resource.  [let^] only makes sure that the resource is released in case
-    something goes wrong. *)
-
-val release : 'r instance -> unit
-(** [release instance] releases the resource, if any, contained by the
-    [instance].
+val drop : 'r instance -> unit
+(** [drop instance] releases the resource, if any, contained by the {!instance}.
 
     @raise Invalid_argument if the resource has been {{!let&} borrowed} and
       hasn't yet been returned. *)
 
-val move : 'r instance -> 'r template
-(** [move instance] returns a resource {!template} where the acquire operation
-    atomically takes the resource from the [instance] and signals the previous
-    owner, and the release operation releases the resource.
-
-    ℹ️ [move] operates lazily such that either party in a transfer may call
-    [move].  The transfer takes place atomically at the point the {!template}
-    returned by [move] is {{!let@} bound} to a scope.
-
-    @raise Invalid_argument if the resource has been {{!let&} borrowed} and
-      hasn't yet been returned, has already been {{!move} moved}, or has been
-      {{!release} released} unless the current fiber has been canceled, in which
-      case the exception that the fiber was canceled with will be raised. *)
+val borrow : 'r instance -> ('r -> 'a) -> 'a
+(** [borrow instance scope] borrows the [resource] stored in the [instance],
+    calls [scope resource], and then returns the [resource] to the [instance]
+    after scope exits.
 
-(** {2 Naked resources}
+    @raise Invalid_argument if the resource has already been {{!borrow}
+      borrowed} and hasn't yet been returned, has already been {{!drop}
+      dropped}, or has already been {{!transfer} transferred}. *)
 
-    A naked resource cannot escape the scope in which it was acquired. *)
+val transfer : 'r instance -> ('r instance -> 'a) -> 'a
+(** [transfer source] transfers the [resource] stored in the [source] instance
+    into a new [target] instance, calls [scope target].  Then, if [scope]
+    returns normally, awaits until the [target] instance becomes empty.  In case
+    [scope] raises an exception or the fiber is canceled, the [target] instance
+    will be {{!drop} dropped}.
 
-val ( let@ ) : 'r template -> ('r -> 'a) -> 'a
-(** [let@ resource = finally release acquire in scope] calls [acquire ()] to
-    obtain a [resource], evaluates [scope], and calls [release resource]
-    whether [scope] returns normally or raises an exception. *)
+    @raise Invalid_argument if the resource has been {{!borrow} borrowed} and
+      hasn't yet been returned, has already been {{!transfer} transferred}, or
+      has been {{!drop} dropped} unless the current fiber has been canceled, in
+      which case the exception that the fiber was canceled with will be
+      raised. *)
 
-val ( let& ) : 'r instance -> ('r -> 'a) -> 'a
-(** [let& resource = instance in scope] borrows the [resource] held by the
-    [instance] for the duration of the [scope].
-
-    @raise Invalid_argument if the resource has already been {{!let&} borrowed}
-      and hasn't yet been returned, has already been {{!release} released}, or
-      has already been {{!move} moved}. *)
+val move : 'r instance -> ('r -> 'a) -> 'a
+(** [move instance scope] is equivalent to
+    {{!transfer} [transfer instance (fun instance -> borrow instance scope)]}. *)
 
 (** {1 Examples}
 
@@ -114,8 +106,8 @@ val ( let& ) : 'r instance -> ('r -> 'a) -> 'a
 
     {2 Looping server}
 
-    There is also a way to {!move} resources to allow forking fibers to handle
-    clients without leaks.
+    There is also a way to {{!move} move} {{!instantiate} instantiated}
+    resources to allow forking fibers to handle clients without leaks.
 
     Here is a sketch of a server that accepts in a loop and forks fibers to
     handle clients:
@@ -126,8 +118,8 @@ val ( let& ) : 'r instance -> ('r -> 'a) -> 'a
 
         (* loop to accept clients *)
         while true do
-          let^ client_fd =
-            finally Unix.close @@ fun () ->
+          let@ client_fd =
+            instantiate Unix.close @@ fun () ->
             Unix.accept ~cloexec:true server_fd
             |> fst
           in
@@ -143,9 +135,10 @@ val ( let& ) : 'r instance -> ('r -> 'a) -> 'a
 
     {2 Move resource from child to parent}
 
-    You can {!move} a resource between any two fibers and {{!let&} borrow} it
-    before moving it.  For example, you can create a resource in a child fiber,
-    use it there, and then move it to the parent fiber:
+    You can {{!move} move} an {{!instantiate} instantiated} resource between any
+    two fibers and {{!borrow} borrow} it before moving it.  For example, you can
+    create a resource in a child fiber, use it there, and then move it to the
+    parent fiber:
 
     {[
       let move_from_child_to_parent () =
@@ -160,13 +153,13 @@ val ( let& ) : 'r instance -> ('r -> 'a) -> 'a
           and pretend_acquire () = () in
 
           (* allocate a resource *)
-          let^ instance =
-            finally pretend_release pretend_acquire
+          let@ instance =
+            instantiate pretend_release pretend_acquire
           in
 
           begin
             (* borrow the resource *)
-            let& resource = instance in
+            let@ resource = borrow instance in
 
             (* use the resource... omitted *)
             ()
diff --git a/lib/picos_structured/picos_structured.mli b/lib/picos_structured/picos_structured.mli
index bdf54bc..7d2fa6c 100644
--- a/lib/picos_structured/picos_structured.mli
+++ b/lib/picos_structured/picos_structured.mli
@@ -476,8 +476,8 @@ end
 
         Bundle.join_after begin fun bundle ->
           while true do
-            let^ client_fd =
-              finally Unix.close @@ fun () ->
+            let@ client_fd =
+              instantiate Unix.close @@ fun () ->
               Unix.accept
                 ~cloexec:true server_fd |> fst
             in


Test patch:
diff --git a/test/test_finally.ml b/test/test_finally.ml
index e62f48e..8669e9e 100644
--- a/test/test_finally.ml
+++ b/test/test_finally.ml
@@ -2,10 +2,10 @@ open Picos_finally
 
 let test_move_is_lazy () =
   Test_scheduler.run @@ fun () ->
-  let^ moveable = finally Fun.id Fun.id in
-  release moveable;
+  let@ moveable = instantiate Fun.id Fun.id in
+  drop moveable;
   begin
-    match move moveable with _ -> () | exception _ -> assert false
+    match (move moveable : _ -> _) with _ -> () | exception _ -> assert false
   end;
   begin
     match
@@ -18,9 +18,9 @@ let test_move_is_lazy () =
 
 let test_borrow_returns_resource () =
   Test_scheduler.run @@ fun () ->
-  let^ moveable = finally Fun.id Fun.id in
+  let@ moveable = instantiate Fun.id Fun.id in
   begin
-    let& _ = moveable in
+    let@ _ = borrow moveable in
     ()
   end;
   let@ _ = move moveable in
@@ -29,7 +29,7 @@ let test_borrow_returns_resource () =
 let () =
   [
     ("move", [ Alcotest.test_case "is lazy" `Quick test_move_is_lazy ]);
-    ( "let&",
+    ( "borrow",
       [
         Alcotest.test_case "returns resource" `Quick
           test_borrow_returns_resource;
diff --git a/test/test_server_and_client.ml b/test/test_server_and_client.ml
index 7dc7751..2570636 100644
--- a/test/test_server_and_client.ml
+++ b/test/test_server_and_client.ml
@@ -39,8 +39,8 @@ let main () =
         Printf.printf "  Server listening\n%!";
         Bundle.join_after @@ fun bundle ->
         while true do
-          let^ client =
-            finally Unix.close @@ fun () ->
+          let@ client =
+            instantiate Unix.close @@ fun () ->
             Printf.printf "  Server accepting\n%!";
             Unix.accept ~cloexec:true socket |> fst
           in