1 files changed, 140 insertions, 0 deletions
diff --git a/stm/coqworkmgrApi.ml b/stm/coqworkmgrApi.ml
new file mode 100644
index 00000000..c34d447e
--- /dev/null
+++ b/stm/coqworkmgrApi.ml
@@ -0,0 +1,140 @@
+(************************************************************************)
+(*  v      *   The Coq Proof Assistant  /  The Coq Development Team     *)
+(* <O___,, *   INRIA - CNRS - LIX - LRI - PPS - Copyright 1999-2015     *)
+(*   \VV/  **************************************************************)
+(*    //   *      This file is distributed under the terms of the       *)
+(*         *       GNU Lesser General Public License Version 2.1        *)
+(************************************************************************)
+
+let debug = false
+
+type request =
+  | Hello of Flags.priority
+  | Get of int
+  | TryGet of int
+  | GiveBack of int
+  | Ping
+
+type response =
+  | Tokens of int
+  | Noluck
+  | Pong of int * int * int
+
+exception ParseError
+
+(* make it work with telnet: strip trailing \r  *)
+let strip_r s =
+  let len = String.length s in
+  if s.[len - 1] <> '\r' then s else String.sub s 0 (len - 1)
+
+let positive_int_of_string n =
+  try
+    let n = int_of_string n in
+    if n <= 0 then raise ParseError else n
+  with Invalid_argument _ | Failure _ -> raise ParseError
+
+let parse_request s =
+  if debug then Printf.eprintf "parsing '%s'\n" s;
+  match Str.split (Str.regexp " ") (strip_r s) with
+  | [ "HELLO"; "LOW" ] -> Hello Flags.Low
+  | [ "HELLO"; "HIGH" ] -> Hello Flags.High
+  | [ "GET"; n ] -> Get (positive_int_of_string n)
+  | [ "TRYGET"; n ] -> TryGet (positive_int_of_string n)
+  | [ "GIVEBACK"; n ] -> GiveBack (positive_int_of_string n)
+  | [ "PING" ] -> Ping
+  | _ -> raise ParseError
+
+let parse_response s =
+  if debug then Printf.eprintf "parsing '%s'\n" s;
+  match Str.split (Str.regexp " ") (strip_r s) with
+  | [ "TOKENS"; n ] -> Tokens (positive_int_of_string n)
+  | [ "NOLUCK" ] -> Noluck
+  | [ "PONG"; n; m; p ] ->
+      let n = try int_of_string n with _ -> raise ParseError in
+      let m = try int_of_string m with _ -> raise ParseError in
+      let p = try int_of_string p with _ -> raise ParseError in
+      Pong (n,m,p)
+  | _ -> raise ParseError
+       
+let print_request = function
+  | Hello Flags.Low -> "HELLO LOW\n"
+  | Hello Flags.High -> "HELLO HIGH\n"
+  | Get n -> Printf.sprintf "GET %d\n" n
+  | TryGet n -> Printf.sprintf "TRYGET %d\n" n
+  | GiveBack n -> Printf.sprintf "GIVEBACK %d\n" n
+  | Ping -> "PING\n"
+
+let print_response = function
+  | Tokens n -> Printf.sprintf "TOKENS %d\n" n
+  | Noluck -> "NOLUCK\n"
+  | Pong (n,m,p) -> Printf.sprintf "PONG %d %d %d\n" n m p
+
+let connect s =
+  try
+    match Str.split (Str.regexp ":") s with
+    | [ h; p ] ->
+        let open Unix in
+        let s = socket PF_INET SOCK_STREAM 0 in
+        connect s (ADDR_INET (inet_addr_of_string h,int_of_string p));
+        Some s
+    | _ -> None
+  with Unix.Unix_error _ -> None
+
+let manager = ref None
+
+let option_map f = function None -> None | Some x -> Some (f x)
+
+let init p =
+  try
+    let sock = Sys.getenv "COQWORKMGR_SOCK" in
+    manager := option_map (fun s -> 
+      let cout = Unix.out_channel_of_descr s in
+      set_binary_mode_out cout true;
+      let cin = Unix.in_channel_of_descr s in
+      set_binary_mode_in cin true;
+      output_string cout (print_request (Hello p)); flush cout;
+      cin, cout) (connect sock)
+  with Not_found | End_of_file -> ()
+
+let with_manager f g =
+  try
+    match !manager with
+    | None -> f ()
+    | Some (cin, cout) -> g cin cout
+  with
+  | ParseError | End_of_file -> manager := None; f ()
+
+let get n =
+  with_manager
+  (fun () ->
+    min n (min !Flags.async_proofs_n_workers !Flags.async_proofs_n_tacworkers))
+  (fun cin cout ->
+    output_string cout (print_request (Get n));
+    flush cout;
+    let l = input_line cin in
+    match parse_response l with
+    | Tokens m -> m
+    | _ -> raise (Failure "coqworkmgr protocol error"))
+
+let tryget n =
+  with_manager
+  (fun () ->
+    Some
+     (min n
+       (min !Flags.async_proofs_n_workers !Flags.async_proofs_n_tacworkers)))
+  (fun cin cout ->
+    output_string cout (print_request (TryGet n));
+    flush cout;
+    let l = input_line cin in
+    match parse_response l with
+    | Tokens m -> Some m
+    | Noluck -> None
+    | _ -> raise (Failure "coqworkmgr protocol error"))
+
+let giveback n =
+  with_manager
+  (fun () -> ())
+  (fun cin cout ->
+    output_string cout (print_request (GiveBack n));
+    flush cout)
+