fd986fb6f191c7dbd32b5bf59740a3c03e27a1e6
[whenjobs.git] / daemon / daemon.ml
1 (* whenjobs
2  * Copyright (C) 2012 Red Hat Inc.
3  *
4  * This program is free software; you can redistribute it and/or modify
5  * it under the terms of the GNU General Public License as published by
6  * the Free Software Foundation; either version 2 of the License, or
7  * (at your option) any later version.
8  *
9  * This program is distributed in the hope that it will be useful,
10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12  * GNU General Public License for more details.
13  *
14  * You should have received a copy of the GNU General Public License along
15  * with this program; if not, write to the Free Software Foundation, Inc.,
16  * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
17  *)
18
19 open Whenutils
20
21 open Unix
22 open Printf
23
24 (* See [exit.c]. *)
25 external _exit : int -> 'a = "whenjobs__exit"
26
27 (* All jobs that are loaded.  Maps name -> [job] structure. *)
28 let jobs = ref StringMap.empty
29
30 (* Map variable names to jobs which depend on that variable.  This
31  * gives us a quick way to tell which jobs might need to be reevaluated
32  * when a variable is set.
33  *)
34 let dependencies = ref StringMap.empty
35
36 (* Current values of variables.  Using the referentially transparent
37  * type Map is very useful here because it lets us cheaply keep
38  * previous values of variables.
39  *)
40 let variables : variables ref = ref StringMap.empty
41
42 (* $HOME/.whenjobs *)
43 let jobsdir = ref ""
44
45 (* Jobs that are running; map of PID -> (job, other data).  Note that
46  * the job may no longer exist *OR* it may have been renamed,
47  * eg. if the jobs file was reloaded.
48  *)
49 let running = ref IntMap.empty
50
51 (* Was debugging requested on the command line? *)
52 let debug = ref false
53
54 (* The server. *)
55 let server = ref None
56
57 let esys = Unixqueue.standard_event_system ()
58
59 let rec init j d =
60   jobsdir := j;
61   debug := d;
62
63   Whenlock.create_lock !jobsdir;
64
65   (* Remove old socket if it exists. *)
66   let addr = sprintf "%s/socket" !jobsdir in
67   (try unlink addr with Unix_error _ -> ());
68
69   server := Some (
70     Whenproto_srv.When.V1.create_server
71       ~proc_reload_file
72       ~proc_set_variable
73       ~proc_get_variable
74       ~proc_get_variable_names
75       ~proc_exit_daemon
76       (Rpc_server.Unix addr)
77       Rpc.Tcp (* not TCP, this is the same as SOCK_STREAM *)
78       Rpc.Socket
79       esys
80   );
81
82   (* Handle SIGCHLD to clean up jobs. *)
83   Sys.set_signal Sys.sigchld (Sys.Signal_handle handle_sigchld)
84
85 and proc_reload_file () =
86   if !debug then Syslog.notice "remote call: reload_file";
87
88   try reload_file (); `ok
89   with Failure err -> `error err
90
91 and proc_set_variable (name, value) =
92   if !debug then Syslog.notice "remote call: set_variable %s" name;
93
94   let value = variable_of_rpc value in
95   variables := StringMap.add name value !variables;
96
97   (* Which jobs need to be re-evaluated? *)
98   let jobnames = try StringMap.find name !dependencies with Not_found -> [] in
99   reevaluate_whenjobs jobnames
100
101 and proc_get_variable name =
102   if !debug then Syslog.notice "remote call: get_variable %s" name;
103
104   try rpc_of_variable (StringMap.find name !variables)
105   with (* all non-existent variables are empty strings *)
106     Not_found -> `string_t ""
107
108 and proc_get_variable_names () =
109   if !debug then Syslog.notice "remote call: get_variable_names";
110
111   (* Only return variables that are non-empty. *)
112   let vars = StringMap.fold (
113     fun name value xs -> if value <> T_string "" then name :: xs else xs
114   ) !variables [] in
115   let vars = Array.of_list vars in
116   Array.sort compare vars;
117   vars
118
119 and proc_exit_daemon () =
120   if !debug then Syslog.notice "remote call: exit_daemon";
121
122   match !server with
123   | None ->
124     `error "exit_daemon: no server handle"
125   | Some s ->
126     Rpc_server.stop_server ~graceful:true s;
127     server := None;
128     `ok
129
130 (* Reload the jobs file. *)
131 and reload_file () =
132   let file = sprintf "%s/jobs.cmo" !jobsdir in
133   Whenfile.init ();
134
135   let js =
136     try
137       Dynlink.loadfile file;
138       let jobs = Whenfile.get_jobs () in
139       Syslog.notice "loaded %d job(s) from %s" (List.length jobs) file;
140       jobs
141     with
142     | Dynlink.Error err ->
143       let err = Dynlink.error_message err in
144       Syslog.error "error loading jobs: %s" err;
145       failwith err
146     | exn ->
147       failwith (Printexc.to_string exn) in
148
149   (* Set 'jobs' and related global variables. *)
150   let () =
151     let map = List.fold_left (
152       fun map j ->
153         let name = j.job_name in
154         StringMap.add name j map
155     ) StringMap.empty js in
156     jobs := map in
157
158   let () =
159     let map = List.fold_left (
160       fun map j ->
161         let deps = dependencies_of_job j in
162         let name = j.job_name in
163         List.fold_left (
164           fun map d ->
165             let names = try StringMap.find d map with Not_found -> [] in
166             StringMap.add d (name :: names) map
167         ) map deps
168     ) StringMap.empty js in
169     dependencies := map in
170
171   (* Re-evaluate all when jobs. *)
172   reevaluate_whenjobs (StringMap.keys !jobs);
173
174   (* Schedule the next every job to run. *)
175   schedule_next_everyjob ()
176
177 (* Re-evaluate each named when-statement job, in a loop until we reach
178  * a fixpoint.  Run those that need to be run.  every-statement jobs
179  * are ignored here.
180  *)
181 and reevaluate_whenjobs jobnames =
182   let rec loop set jobnames =
183     let set' =
184       List.fold_left (
185         fun set jobname ->
186           let job =
187             try StringMap.find jobname !jobs
188             with Not_found -> assert false in
189           assert (jobname = job.job_name);
190
191           let r, job' =
192             try job_evaluate job !variables
193             with Invalid_argument err | Failure err ->
194               Syslog.error "error evaluating job %s (at %s): %s"
195                 jobname (Camlp4.PreCast.Ast.Loc.to_string job.job_loc) err;
196               false, job in
197
198           jobs := StringMap.add jobname job' !jobs;
199
200           if !debug then
201             Syslog.notice "evaluate %s -> %b\n" jobname r;
202
203           if r then StringSet.add jobname set else set
204       ) set jobnames in
205     if StringSet.compare set set' <> 0 then
206       loop set' jobnames
207     else
208       set'
209   in
210   let set = loop StringSet.empty jobnames in
211   let jobnames = StringSet.elements set in
212   (* Ensure the jobs always run in predictable (name) order. *)
213   let jobnames = List.sort compare jobnames in
214   List.iter run_job
215     (List.map (fun jobname -> StringMap.find jobname !jobs) jobnames)
216
217 (* Schedule the next every-statement job to run, if there is one.  We
218  * look at the every jobs, work out the time that each must run at,
219  * pick the job(s) which must run soonest, and schedule a timer to run
220  * them.  When the timer fires, it runs those jobs, then calls this
221  * function again.
222  *)
223 and schedule_next_everyjob () =
224   let t = time () in
225
226   (* Get only everyjobs. *)
227   let jobs = StringMap.values !jobs in
228   let jobs = filter_map (
229     function
230     | { job_cond = Every_job period } as job -> Some (job, period)
231     | { job_cond = When_job _ } -> None
232   ) jobs in
233
234   (* Map everyjob to next time it must run. *)
235   let jobs = List.map (
236     fun (job, period) ->
237       let t' = next_periodexpr t period in
238       assert (t' > t); (* serious bug in next_periodexpr if false *)
239       job, t'
240   ) jobs in
241
242   (* Sort, soonest first. *)
243   let jobs = List.sort (fun (_,a) (_,b) -> compare a b) jobs in
244
245   if !debug then (
246     List.iter (
247       fun (job, t) ->
248         Syslog.notice "%s: next scheduled run at %s"
249           job.job_name (string_of_time_t t)
250     ) jobs
251   );
252
253   (* Pick the job(s) which run soonest. *)
254   let rec pick = function
255     | [] -> 0., []
256     | [j, t] -> t, [j]
257     | (j1, t) :: (j2, t') :: _ when t < t' -> t, [j1]
258     | (j1, t) :: (((j2, t') :: _) as rest) -> t, (j1 :: snd (pick rest))
259   in
260   let t, jobs = pick jobs in
261
262   if t > 0. then (
263     if jobs <> [] then (
264       (* Ensure the jobs always run in predictable (name) order. *)
265       let jobs =
266         List.sort (fun { job_name = a } { job_name = b } -> compare a b) jobs in
267
268       if !debug then
269         Syslog.notice "scheduling job(s) %s to run at %s"
270           (String.concat ", " (List.map (fun { job_name = name } -> name) jobs))
271           (string_of_time_t t);
272
273       (* Schedule them to run at time t. *)
274       let g = Unixqueue.new_group esys in
275       let t_diff = t -. Unix.time () in
276       let t_diff = if t_diff < 0. then 0. else t_diff in
277       let run_jobs () =
278         Unixqueue.clear esys g;         (* Delete the timer. *)
279         List.iter run_job jobs;
280         schedule_next_everyjob ()
281       in
282       Unixqueue.weak_once esys g t_diff run_jobs;
283     )
284   )
285
286 and string_of_time_t t =
287   let tm = gmtime t in
288   sprintf "%04d-%02d-%02d %02d:%02d:%02d UTC"
289     (1900+tm.tm_year) (1+tm.tm_mon) tm.tm_mday
290     tm.tm_hour tm.tm_min tm.tm_sec
291
292 and run_job job =
293   Syslog.notice "running %s" job.job_name;
294
295   (* Create a temporary directory.  The current directory of the job
296    * will be in this directory.  The directory is removed when the
297    * child process exits.
298    *)
299   let dir = tmpdir () in
300
301   let pid = fork () in
302   if pid = 0 then ( (* child process running the job *)
303     chdir dir;
304
305     (* Set environment variables corresponding to each variable. *)
306     StringMap.iter
307       (fun name value -> putenv name (string_of_variable value)) !variables;
308
309     (* Set the $JOBNAME environment variable. *)
310     putenv "JOBNAME" job.job_name;
311
312     (* Create a temporary file containing the shell script fragment. *)
313     let script = dir // "script" in
314     let chan = open_out script in
315     output_string chan job.job_script.sh_script;
316     close_out chan;
317     chmod script 0o700;
318
319     (* Execute the shell script. *)
320     (try execvp "bash" [| "bash"; "-c"; script |];
321      with Unix_error (err, fn, _) ->
322        Syslog.error "%s failed: %s: %s" fn script (error_message err)
323     );
324     _exit 1
325   );
326
327   (* Remember this PID, the job and the temporary directory, so we
328    * can clean up when the child exits.
329    *)
330   running := IntMap.add pid (job, dir) !running
331
332 and tmpdir () =
333   let chan = open_in "/dev/urandom" in
334   let data = String.create 16 in
335   really_input chan data 0 (String.length data);
336   close_in chan;
337   let data = Digest.to_hex (Digest.string data) in
338   let dir = Filename.temp_dir_name // sprintf "whenjobs%s" data in
339   mkdir dir 0o700;
340   dir
341
342 (* This is called when a job (child process) exits. *)
343 and handle_sigchld _ =
344   try
345     let pid, status = waitpid [WNOHANG] 0 in
346     if pid > 0 then (
347       (* Look up the PID in the running jobs map. *)
348       let job, dir = IntMap.find pid !running in
349       running := IntMap.remove pid !running;
350       cleanup_job job dir
351     )
352   with Unix_error _ | Not_found -> ()
353
354 and cleanup_job job dir =
355   (* This should be safe because the path cannot contain shell metachars. *)
356   let cmd = sprintf "rm -rf '%s'" dir in
357   ignore (Sys.command cmd)
358
359 let main_loop () =
360   Unixqueue.run esys