1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

(************************************************************************)
(*         *   The Coq Proof Assistant / The Coq Development Team       *)
(*  v      *   INRIA, CNRS and contributors - Copyright 1999-2019       *)
(* <O___,, *       (see CREDITS file for the list of authors)           *)
(*   \VV/  **************************************************************)
(*    //   *    This file is distributed under the terms of the         *)
(*         *     GNU Lesser General Public License Version 2.1          *)
(*         *     (see LICENSE file for the text of the license)         *)
(************************************************************************)

let debug = false

type priority = Low | High

(* Default priority *)

let default_async_proofs_worker_priority = Low

let string_of_priority = function Low -> "low" | High -> "high"
let priority_of_string = function
  | "low" -> Low
  | "high" -> High
  | _ -> raise (Invalid_argument "priority_of_string")

type request =
  | Hello of 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 Low
  | [ "HELLO"; "HIGH" ] -> Hello 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 Low -> "HELLO LOW\n"
  | Hello 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 () -> n)
  (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 n)
  (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)