/
lifedb_out_tasks.ml
345 lines (318 loc) · 12.7 KB
/
lifedb_out_tasks.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
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
(* Copyright (C) 2009 Anil Madhavapeddy <anil@recoil.org>
This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License along
with this program; if not, write to the Free Software Foundation, Inc.,
51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
*)
(* Task scheduling *)
open Printf
open Utils
module LS=Lifedb_schema
module SS=Sync_schema
exception Task_error of string
exception Internal_task_error of string
let m = Mutex.create ()
type task_state = {
cmd: string;
plugin: string;
mtype: string;
cwd: string;
start_time: float;
secret: (string * string) option;
args : string list option;
mutable outfd: Unix.file_descr option;
mutable errfd: Unix.file_descr option;
mutable running: Fork_helper.task;
mutable uids: string list;
mutable files: string list;
}
let task_list = Hashtbl.create 1
let task_table_limit = 20
let task_poll_period = ref 20.
let task_throttle () = Thread.delay 0.1
let json_of_task name t : Lifedb.Rpc.Task.out_r =
let secret = match t.secret with |None -> None
|Some (s,u) -> Some (object method service=s method username=u end) in
object
method name=name
method plugin=t.plugin
method pltype=t.mtype
method secret=secret
method args=t.args
method duration=Unix.gettimeofday () -. t.start_time
method pid=Fork_helper.pid_of_task t.running
end
let string_of_task t =
let running = Fork_helper.string_of_task t.running in
sprintf ": `%s` %s" t.cmd running
let log_task_table () =
Hashtbl.iter (fun name state ->
Log.logmod "Tasks" "%s: %s" name (string_of_task state)
) task_list
let find_task name =
try
Some (Hashtbl.find task_list name)
with
Not_found -> None
let find_task_by_mtype mtype =
let f = ref None in
Hashtbl.iter (fun n t ->
if t.mtype = mtype then
f := Some (n,t)
) task_list;
!f
(* create task descriptor and leave it unstarted *)
let create_task task_name (p:Lifedb.Rpc.Task.out_t) =
assert(not (Mutex.try_lock m));
if Hashtbl.length task_list >= task_table_limit then
raise (Task_error "too many tasks already registered");
if String.contains task_name '.' || (String.contains task_name '/') then
raise (Task_error "task name cant contain . or /");
let pl = match Lifedb_plugin.find_plugin p#plugin with
|None -> raise (Task_error (sprintf "plugin %s not found" p#plugin))
|Some x -> x in
let secret = match p#secret with
|None -> None
|Some s -> Some (s#service, s#username) in
let now_time = Unix.gettimeofday () in
let task = { cmd=pl#cmd; outfd=None; errfd=None; cwd=pl#dir; plugin=pl#name; secret=secret; start_time=now_time; running=(Fork_helper.blank_task ()); args=p#args; uids=[]; files=[]; mtype=p#pltype } in
Hashtbl.add task_list task_name task;
Log.logmod "Tasks" "Created outbound task '%s' %s" task_name (string_of_task task)
let find_or_create_task name (t:Lifedb.Rpc.Task.out_t) =
match find_task name with
|Some _ -> ()
|None -> create_task name t
(* close any logging fds, reset fields *)
let reset_task name =
assert(not (Mutex.try_lock m));
let closeopt task = function
|None -> ()
|Some fd ->
let lg = sprintf "[%s] Log closing: %s\n" (current_datetime()) (Fork_helper.string_of_task task.running) in
ignore(Unix.handle_unix_error (Unix.write fd lg 0) (String.length lg));
Unix.handle_unix_error Unix.close fd;
in
match find_task name with
|Some task ->
let time_taken = (Unix.gettimeofday ()) -. task.start_time in
let exit_code = Fork_helper.exit_code_of_task task.running in
Log.push (`Plugin (name, time_taken, exit_code));
closeopt task task.outfd;
closeopt task task.errfd;
task.outfd <- None;
task.errfd <- None;
task.running <- Fork_helper.blank_task ();
task.uids <- [];
task.files <- [];
|None -> ()
(* remove the task entirely *)
let destroy_task name =
assert(not (Mutex.try_lock m));
match find_task name with
|Some task -> begin
let final_status = Fork_helper.destroy task.running in
reset_task name;
Hashtbl.remove task_list name;
Log.logmod "Tasks" "Outbound task %s destroyed: %s" name
(Fork_helper.string_of_status final_status);
end
|None -> raise (Task_error "task not found")
(* split up a list of entris into a hashtable of their respective mtypes *)
let partition_entries_into_mtypes lifedb es =
let h = Hashtbl.create 1 in
List.iter (fun e ->
if not (Hashtbl.mem h e#mtype#name) then
Hashtbl.add h e#mtype#name [];
Hashtbl.replace h e#mtype#name (e :: (Hashtbl.find h e#mtype#name))
) es;
h
let start_task name =
let t = Hashtbl.find task_list name in
assert(not (Mutex.try_lock m));
let env = match t.secret with
|None -> [||]
|Some (s,u) -> begin
match Lifedb_passwd.lookup_passwd s u with
|Some p -> [| ("LIFEDB_PASSWORD=" ^ p); ("LIFEDB_USERNAME="^u) |]
|None -> Log.logmod "Tasks" "WARNING: unable to find passwd for task '%s'" name; [||]
end in
(* add environment arguments *)
let args = match t.args with None -> [||] | Some a -> Array.of_list a in
let env = Array.append env args in
let logdir = Lifedb_config.Dir.log() in
let logfile = sprintf "%s/%s.log" logdir name in
let errlogfile = sprintf "%s/%s.err" logdir name in
let openfdfn f = Unix.handle_unix_error (Unix.openfile f [ Unix.O_APPEND; Unix.O_CREAT; Unix.O_WRONLY]) 0o600 in
let outfd = openfdfn logfile in
let errfd = openfdfn errlogfile in
let logfn fd s = ignore(Unix.write fd s 0 (String.length s)) in
let tmstr = current_datetime () in
logfn outfd (sprintf "[%s] Stdout log started\n" tmstr);
logfn errfd (sprintf "[%s] Stderr log started\n" tmstr);
let env = Array.append env [| "LIFEDB_SYNC_DIR=out"; (sprintf "LIFEDB_UID_MAP=%s" (Lifedb_config.Dir.uidmap ()));
(sprintf "HOME=%s" (Sys.getenv "HOME"));
(sprintf "USER=%s" (Sys.getenv "USER")) |] in
let cmd =
if Lifedb_config.test_mode () then
sprintf "sleep %d" (Random.int 5 + 3)
else
(* XXX check shell escaping here!! *)
sprintf "%s %s" t.cmd (String.concat " " (List.map String.escaped t.files))
in
let ts = Fork_helper.create cmd env t.cwd (logfn outfd) (logfn errfd) in
Log.logmod "Tasks" "Executing outbound command '%s' (%s)" name cmd;
task_throttle ();
t.running <- ts;
t.outfd <- Some outfd;
t.errfd <- Some errfd
(* Look for items in the INBOX with a pltype matching an active plugin, and schedule it
if so *)
let task_sweep lifedb syncdb () =
(* for each user, look for entries in the inbox to them *)
List.iter (fun (user:SS.User.t) ->
let es = LS.Entry.get_by_inbox_delivered ~inbox:(Some user#uid) ~delivered:0L lifedb in
match es with
|[] -> ()
|es -> begin
(* constrain es to only 50 entries at a time to avoid overloading output plugin *)
let es = list_max_size 50 es in
(* we have inbox entries, look for a plugin to handle each mtype *)
let h = partition_entries_into_mtypes lifedb es in
Hashtbl.iter (fun mtype_name es ->
(* look for an output task to handle this mtype name *)
match find_task_by_mtype mtype_name with
|None ->
Log.logmod "Task" "Unable to find output task for <- %s : %s" mtype_name user#uid
|Some (name,t) -> begin
match Fork_helper.status_of_task t.running with
|Fork_helper.Not_started ->
(* set the entry UIDs and kick the command off *)
t.uids <- List.map (fun e -> e#uid) es;
t.files <- List.map (fun e -> e#file_name) es;
start_task name
|_ ->
Log.logmod "Task" "Pending INBOX items, but already running %s" name
end
) h
end
) (SS.User.get syncdb);
Hashtbl.iter (fun name task ->
let td = string_of_task task in
match Fork_helper.status_of_task task.running with
|Fork_helper.Running pid ->
Log.logmod "Sweep" "%s ... %s" name td
|Fork_helper.Not_started -> ()
|Fork_helper.Done exit_code ->
Log.logmod "Sweep" "%s ... finished %s" name td;
if exit_code = 0 then begin
(* successfully delivered msgs, so mark them in the DB as delivered *)
List.iter (fun uid ->
match LS.Entry.get_by_uid ~uid lifedb with
|[e] ->
Log.logmod "Task" "Successfully delivered: %s" e#file_name;
e#set_delivered 1L;
ignore(e#save)
|_ -> ()
) task.uids;
end;
reset_task name;
|Fork_helper.Killed signal ->
Log.logmod "Sweep" "%s ... crashed %s" name td;
reset_task name;
) task_list
let dispatch cgi = function
|`Create (name,p) ->
let params = Lifedb.Rpc.Task.out_t_of_json (Json_io.json_of_string p) in
with_lock m (fun () ->
match find_task name with
|Some state ->
Lifedb_rpc.return_error cgi `Bad_request "Task already exists" "Use a different id"
|None -> begin
try
create_task name params
with
|Task_error err ->
Lifedb_rpc.return_error cgi `Bad_request "Task error" err
end
)
|`Get name ->
with_lock m (fun () ->
match find_task name with
|Some state ->
cgi#output#output_string (Json_io.string_of_json (Lifedb.Rpc.Task.json_of_out_r (json_of_task name state)))
|None ->
Lifedb_rpc.return_error cgi `Not_found "Task error" "Task not found"
)
|`List ->
with_lock m (fun () ->
let r = Hashtbl.fold (fun name state a -> json_of_task name state :: a) task_list [] in
let res = object method results=List.length r method rows=r end in
cgi#output#output_string (Json_io.string_of_json (Lifedb.Rpc.Task.json_of_out_rs res))
)
|`Destroy name ->
with_lock m (fun () ->
try
destroy_task name
with |Task_error err ->
Lifedb_rpc.return_error cgi `Bad_request "Task error" err
)
(* task thread which waits on a condition to do a sweep. is signalled regularly
or via a process exiting and delivering a SIGCHLD *)
let c = Condition.create ()
let cm = Mutex.create ()
let task_thread () =
let lifedb = LS.Init.t (Lifedb_config.Dir.lifedb_db ()) in
let syncdb = SS.Init.t (Lifedb_config.Dir.sync_db ()) in
while true do
with_lock cm (fun () ->
Condition.wait c cm;
with_lock m (task_sweep lifedb syncdb);
)
done
(* thread to kick the sweeping thread regularly to update task status. *)
let task_regular_kick () =
while true do
with_lock cm (fun () ->
Condition.signal c
);
Thread.delay !task_poll_period
done
(* scan the config directory and spawn tasks *)
let config_file_extension = ".outconf"
let scan_config_file config_file =
Log.logmod "Tasks" "Scanning config file %s" config_file;
let task = Lifedb.Rpc.Task.out_t_of_json (Json_io.load_json config_file) in
let task_name = Filename.chop_suffix (Filename.basename config_file) config_file_extension in
match Lifedb_plugin.find_plugin task#plugin with
|None -> Log.logmod "Tasks" "Plugin '%s' not found for task '%s', skipping it" task#plugin task_name;
|Some _ ->
Log.logmod "Tasks" "Added '%s' (plugin %s)" task_name task#plugin;
let task : Lifedb.Rpc.Task.out_t = object
method plugin=task#plugin
method secret=task#secret
method args=task#args
method pltype=task#pltype
end in
with_lock m (fun () -> find_or_create_task task_name task)
let do_scan () =
let config_dir = Lifedb_config.Dir.config () in
let dh = Unix.opendir config_dir in
try_final (fun () ->
repeat_until_eof (fun () ->
let next_entry = Unix.readdir dh in
if Filename.check_suffix next_entry config_file_extension then
scan_config_file (Filename.concat config_dir next_entry)
)
) (fun () -> Unix.closedir dh)
let init () =
let _ = Thread.create task_thread () in
let _ = Thread.create task_regular_kick () in
Sys.set_signal Sys.sigchld (Sys.Signal_handle (fun _ ->
with_lock cm (fun () -> Condition.signal c)))