From a45d8c46a6e926ae1f48acf0bd1917c01a3eaf2c Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Tue, 7 Mar 2017 09:35:52 +0100
Subject: [PATCH] add `CCUnix.with_file_lock` for locking whole files

---
 src/unix/CCUnix.ml  | 43 +++++++++++++++++++++++++++++++++++++++++++
 src/unix/CCUnix.mli |  9 +++++++++
 2 files changed, 52 insertions(+)

diff --git a/src/unix/CCUnix.ml b/src/unix/CCUnix.ml
index 216522a3..4d2b4d55 100644
--- a/src/unix/CCUnix.ml
+++ b/src/unix/CCUnix.ml
@@ -224,6 +224,49 @@ let establish_server sockaddr ~f =
       continue := false
   done
 
+
+(** {6 Locking} *)
+
+let with_file_lock ~kind filename f =
+  let lock_file = Unix.openfile filename [Unix.O_CREAT; Unix.O_WRONLY] 0o644 in
+  let lock_action = match kind with
+    | `Read -> Unix.F_RLOCK
+    | `Write -> Unix.F_LOCK
+  in
+  Unix.lockf lock_file lock_action 0;
+  try
+    let x = f () in
+    Unix.lockf lock_file Unix.F_ULOCK 0;
+    Unix.close lock_file;
+    x
+  with e ->
+    Unix.lockf lock_file Unix.F_ULOCK 0;
+    Unix.close lock_file;
+    raise e
+
+(*$R
+  let m = 200 in
+  let n = 50 in
+  let write_atom filename s =
+    with_file_lock ~kind:`Write filename
+      (fun () ->
+        CCIO.with_out ~flags:[Open_append; Open_creat]
+          filename (fun oc -> output_string oc s; flush oc))
+  in
+  let f filename =
+    for j=1 to m do
+      write_atom filename "foo\n"
+    done
+  in
+  CCIO.File.with_temp ~prefix:"containers_" ~suffix:".txt"
+    (fun filename ->
+        let a = Array.init n (fun _ -> Thread.create f filename) in
+        Array.iter Thread.join a;
+        let lines = CCIO.with_in filename CCIO.read_lines_l in
+        assert_equal ~printer:string_of_int (n * m) (List.length lines);
+        assert_bool "all valid" (List.for_all ((=) "foo") lines))
+*)
+
 module Infix = struct
   let (?|) fmt = call_full fmt
 
diff --git a/src/unix/CCUnix.mli b/src/unix/CCUnix.mli
index 237df841..6e6be9a2 100644
--- a/src/unix/CCUnix.mli
+++ b/src/unix/CCUnix.mli
@@ -161,6 +161,15 @@ val establish_server : Unix.sockaddr -> f:(in_channel -> out_channel -> _) -> un
     The callback should raise {!ExitServer} to stop the loop.
     @since 0.16 *)
 
+val with_file_lock : kind:[`Read|`Write] -> string -> (unit -> 'a) -> 'a
+(** [with_file_lock ~kind filename f] puts a lock on the offset 0
+    of the file named [filename], calls [f] and returns its result after
+    the file is unlocked. If [f ()] raises an exception the exception is
+    re-raised after the file is unlocked.
+
+    @param kind specifies whether the lock is read-only or read-write.
+    @since NEXT_RELEASE *)
+
 (** {2 Infix Functions} *)
 
 module Infix : sig