Move 'with_socket' in the 'Socket' module

4 files changed, 26 insertions, 16 deletions

diff --git a/lib/mltp_ipc/socket.ml b/lib/mltp_ipc/socket.ml
index e85b7b2..0d5333d 100644
--- a/lib/mltp_ipc/socket.ml
+++ b/lib/mltp_ipc/socket.ml
@@ -5,6 +5,8 @@
 type socket = Lwt_io.input_channel * Lwt_io.output_channel * Lwt_unix.file_descr
 type error = Bad_magic_string of string | Connection_closed
 
+let ( let*! ) x k = Lwt.bind x k
+
 let connect socket_path : socket Lwt.t =
   let open Lwt.Syntax in
   let socket = Lwt_unix.socket PF_UNIX SOCK_STREAM 0 in
@@ -13,9 +15,21 @@ let connect socket_path : socket Lwt.t =
   let socket_out = Lwt_io.of_fd ~mode:Output socket in
   (socket_in, socket_out, socket)
 
-let ( let*! ) x k = Lwt.bind x k
 let close (_, _, s) = Lwt_unix.close s
 
+let with_socket socket_path f =
+  let open Lwt.Syntax in
+  let* socket = connect socket_path in
+  Lwt.try_bind
+    (fun () ->
+      let* res = f socket in
+      let* () = close socket in
+      Lwt.return res)
+    Lwt.return
+    (fun exn ->
+      let* () = close socket in
+      raise exn)
+
 let catch_end_of_file f =
   Lwt.try_bind f Lwt_result.return @@ function
   | End_of_file -> Lwt_result.fail Connection_closed
diff --git a/lib/mltp_ipc/socket.mli b/lib/mltp_ipc/socket.mli
index 46dda86..9e8cf9f 100644
--- a/lib/mltp_ipc/socket.mli
+++ b/lib/mltp_ipc/socket.mli
@@ -19,6 +19,12 @@ val connect : string -> socket Lwt.t
 val close : socket -> unit Lwt.t
 (** Close a bi-directional connection with a peer. *)
 
+val with_socket : string -> (socket -> 'a Lwt.t) -> 'a Lwt.t
+(** [with_socket path k] establishes a bi-connection with a peer using
+    the UNIX socket located at [path], hands over the socket to the
+    continuation [k], and takes care of closing the connection prior
+    to returning the result, even in case of an exception. *)
+
 type error =
   | Bad_magic_string of string
       (** When trying to read a MTLP message, the magic string was not
diff --git a/lib/sway_ipc/sway_ipc.ml b/lib/sway_ipc/sway_ipc.ml
index be33e9f..13ec1af 100644
--- a/lib/sway_ipc/sway_ipc.ml
+++ b/lib/sway_ipc/sway_ipc.ml
@@ -24,18 +24,7 @@ let trust_sway f =
   let* x = f () in
   match x with Ok x -> Lwt.return x | Error e -> raise (Sway_ipc_error e)
 
-let with_socket f =
-  let open Lwt.Syntax in
-  let* socket = connect () in
-  Lwt.try_bind
-    (fun () ->
-      let* res = f socket in
-      let* () = Socket.close socket in
-      Lwt.return res)
-    Lwt.return
-    (fun exn ->
-      let* () = Socket.close socket in
-      raise exn)
+let with_socket f = Socket.with_socket (sway_sock_path ()) f
 
 let socket_from_option = function
   | Some socket -> Lwt.return socket
diff --git a/lib/sway_ipc/sway_ipc.mli b/lib/sway_ipc/sway_ipc.mli
index d5d7e72..753a0df 100644
--- a/lib/sway_ipc/sway_ipc.mli
+++ b/lib/sway_ipc/sway_ipc.mli
@@ -17,9 +17,10 @@ val close : socket -> unit Lwt.t
 (** [close socket] puts an end to a connection with Sway. *)
 
 val with_socket : (socket -> 'a Lwt.t) -> 'a Lwt.t
-(** [with_socket k] establishes a connection with Sway, hands over the
-    socket to the continuation [k], and takes care of closing the
-    connection with Sway prior to returning the result. *)
+(** [with_socket k] establishes a bi-connection with Sway, hands over
+    the socket to the continuation [k], and takes care of closing the
+    connection prior to returning the result, even in case of an
+    exception. *)
 
 val send_command : ?socket:socket -> 'a Sway_ipc_types.Message.t -> 'a Lwt.t
 (** [send_command ?socket msg] sends the command [msg] to Sway (by