Source file parallel_managed.ml
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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
open Core
open Async
module type Worker = sig
type t [@@deriving sexp_of]
type unmanaged_t
type 'a functions
val functions : unmanaged_t functions
type worker_state_init_arg
type connection_state_init_arg
module Id : Identifiable
val id : t -> Id.t
val spawn
: ?where:Executable_location.t
-> ?name:string
-> ?env:(string * string) list
-> ?connection_timeout:Time.Span.t
-> ?cd:string
-> ?umask:int
-> redirect_stdout:Fd_redirection.t
-> redirect_stderr:Fd_redirection.t
-> worker_state_init_arg
-> connection_state_init_arg
-> on_failure:(Error.t -> unit)
-> t Or_error.t Deferred.t
val spawn_exn
: ?where:Executable_location.t
-> ?name:string
-> ?env:(string * string) list
-> ?connection_timeout:Time.Span.t
-> ?cd:string
-> ?umask:int
-> redirect_stdout:Fd_redirection.t
-> redirect_stderr:Fd_redirection.t
-> worker_state_init_arg
-> connection_state_init_arg
-> on_failure:(Error.t -> unit)
-> t Deferred.t
val run
: t
-> f:(unmanaged_t, 'query, 'response) Parallel.Function.t
-> arg:'query
-> 'response Or_error.t Deferred.t
val run_exn
: t
-> f:(unmanaged_t, 'query, 'response) Parallel.Function.t
-> arg:'query
-> 'response Deferred.t
val kill : t -> unit Or_error.t Deferred.t
val kill_exn : t -> unit Deferred.t
end
module Make (S : Parallel.Worker_spec) = struct
module Unmanaged = Parallel.Make (S)
module Id = Utils.Worker_id
type nonrec t =
{ unmanaged : Unmanaged.t
; connection_state_init_arg : S.Connection_state.init_arg
; id : Id.t
}
type unmanaged_t = Unmanaged.t
type conn =
[ `Pending of Unmanaged.Connection.t Or_error.t Ivar.t
| `Connected of Unmanaged.Connection.t
]
let sexp_of_t t = [%sexp_of: Unmanaged.t] t.unmanaged
let id t = t.id
let functions = Unmanaged.functions
let workers : conn Id.Table.t = Id.Table.create ()
let get_connection { unmanaged = t; connection_state_init_arg; id } =
match Hashtbl.find workers id with
| Some (`Pending ivar) -> Ivar.read ivar
| Some (`Connected conn) -> Deferred.Or_error.return conn
| None ->
let ivar = Ivar.create () in
Hashtbl.add_exn workers ~key:id ~data:(`Pending ivar);
(match%map Unmanaged.Connection.client t connection_state_init_arg with
| Error e ->
Ivar.fill ivar (Error e);
Hashtbl.remove workers id;
Error e
| Ok conn ->
Ivar.fill ivar (Ok conn);
Hashtbl.set workers ~key:id ~data:(`Connected conn);
(Unmanaged.Connection.close_finished conn >>> fun () -> Hashtbl.remove workers id);
Ok conn)
;;
let with_shutdown_on_error worker ~f =
match%bind f () with
| Ok _ as ret -> return ret
| Error _ as ret ->
let%bind (_ : unit Or_error.t) = Unmanaged.shutdown worker in
return ret
;;
let spawn
?where
?name
?env
?connection_timeout
?cd
?umask
~redirect_stdout
~redirect_stderr
worker_state_init_arg
connection_state_init_arg
~on_failure
=
Unmanaged.spawn
?where
?env
?name
?connection_timeout
?cd
?umask
~shutdown_on:Heartbeater_timeout
~redirect_stdout
~redirect_stderr
worker_state_init_arg
~on_failure
>>=? fun worker ->
with_shutdown_on_error worker ~f:(fun () ->
Unmanaged.Connection.client worker connection_state_init_arg)
>>|? fun connection ->
let id = Id.create () in
Hashtbl.add_exn workers ~key:id ~data:(`Connected connection);
(Unmanaged.Connection.close_finished connection
>>> fun () ->
match Hashtbl.find workers id with
| None ->
()
| Some _ ->
Hashtbl.remove workers id;
let error = Error.createf !"Lost connection with worker" in
on_failure error);
{ unmanaged = worker; connection_state_init_arg; id }
;;
let spawn_exn
?where
?name
?env
?connection_timeout
?cd
?umask
~redirect_stdout
~redirect_stderr
worker_state_init_arg
connection_init_arg
~on_failure
=
spawn
?where
?name
?env
?connection_timeout
?cd
?umask
~redirect_stdout
~redirect_stderr
worker_state_init_arg
connection_init_arg
~on_failure
>>| Or_error.ok_exn
;;
let kill t =
Hashtbl.remove workers t.id;
Unmanaged.shutdown t.unmanaged
;;
let kill_exn t = kill t >>| Or_error.ok_exn
let run t ~f ~arg =
get_connection t >>=? fun conn -> Unmanaged.Connection.run conn ~f ~arg
;;
let run_exn t ~f ~arg = run t ~f ~arg >>| Or_error.ok_exn
end