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

(************************************************************************)
(*  v      *   The Coq Proof Assistant  /  The Coq Development Team     *)
(* <O___,, *   INRIA - CNRS - LIX - LRI - PPS - Copyright 1999-2017     *)
(*   \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)