Allow get_dependencies to check dependencies of multiple variables.
[whenjobs.git] / daemon / daemon.ml
index 9e972ac..ed5e067 100644 (file)
@@ -32,11 +32,14 @@ let jobsdir = ref ""
 (* The state. *)
 let state = ref Whenstate.empty
 
-(* Jobs that are running; map of PID -> (job, other data).  Note that
- * the job may no longer exist *OR* it may have been renamed,
+(* Jobs that are running: a map of PID -> (job, tmpdir, serial, start_time).
+ * Note that the job may no longer exist *OR* it may have been renamed,
  * eg. if the jobs file was reloaded.
  *)
-let running = ref IntMap.empty
+let runningmap = ref IntMap.empty
+
+(* Serial numbers of running jobs.  Map of serial -> PID (in runningmap). *)
+let serialmap = ref BigIntMap.empty
 
 (* Was debugging requested on the command line? *)
 let debug = ref false
@@ -72,6 +75,10 @@ let rec init j d =
       ~proc_get_variable
       ~proc_get_variable_names
       ~proc_exit_daemon
+      ~proc_get_jobs
+      ~proc_cancel_job
+      ~proc_start_job
+      ~proc_get_job
       (Rpc_server.Unix addr)
       Rpc.Tcp (* not TCP, this is the same as SOCK_STREAM *)
       Rpc.Socket
@@ -100,7 +107,7 @@ and proc_set_variable (name, value) =
     state := Whenstate.set_variable !state name value;
 
     (* Which jobs need to be re-evaluated? *)
-    let jobs = Whenstate.get_dependencies !state name in
+    let jobs = Whenstate.get_dependencies !state [name] in
     reevaluate_whenjobs jobs;
 
     `ok
@@ -133,6 +140,46 @@ and proc_exit_daemon () =
     server := None;
     `ok
 
+and proc_get_jobs () =
+  let running = Array.of_list (IntMap.values !runningmap) in
+  Array.map (
+    fun (job, dir, serial, start_time) ->
+      { Whenproto_aux.job_name = job.job_name;
+        job_serial = string_of_big_int serial;
+        job_tmpdir = dir; job_start_time = Int64.of_float start_time }
+  ) running
+
+and proc_cancel_job serial =
+  try
+    let serial = big_int_of_string serial in
+    let pid = BigIntMap.find serial !serialmap in
+    kill pid 15;
+    `ok
+  with
+  | Not_found -> `error "job not found"
+  | exn -> `error (Printexc.to_string exn)
+
+and proc_start_job jobname =
+  try
+    let job = Whenstate.get_job !state jobname in
+    run_job job;
+    `ok
+  with
+  | Not_found -> `error "job not found"
+  | exn -> `error (Printexc.to_string exn)
+
+and proc_get_job serial =
+  try
+    let serial = big_int_of_string serial in
+    let pid = BigIntMap.find serial !serialmap in
+    let job, dir, serial, start_time = IntMap.find pid !runningmap in
+    { Whenproto_aux.job_name = job.job_name;
+      job_serial = string_of_big_int serial;
+      job_tmpdir = dir; job_start_time = Int64.of_float start_time }
+  with
+  | Not_found -> failwith "job not found"
+  | exn -> failwith (Printexc.to_string exn)
+
 (* Reload the jobs file. *)
 and reload_file () =
   let file = sprintf "%s/jobs.cmo" !jobsdir in
@@ -158,6 +205,7 @@ and reload_file () =
     | exn ->
       failwith (Printexc.to_string exn) in
 
+  let s = Whenstate.copy_prev_state !state s in
   state := s;
 
   (* Re-evaluate all when jobs. *)
@@ -285,73 +333,97 @@ and delete_timer_group () =
     Unixqueue.clear esys g;
     timer_group := None
 
-and string_of_time_t t =
-  let tm = gmtime t in
-  sprintf "%04d-%02d-%02d %02d:%02d:%02d UTC"
-    (1900+tm.tm_year) (1+tm.tm_mon) tm.tm_mday
-    tm.tm_hour tm.tm_min tm.tm_sec
-
 and run_job job =
-  let () =
-    (* Increment JOBSERIAL. *)
-    let serial =
-      match Whenstate.get_variable !state "JOBSERIAL" with
-      | T_int serial ->
-        let serial = succ_big_int serial in
-        state := Whenstate.set_variable !state "JOBSERIAL" (T_int serial);
-        serial
-      | _ -> assert false in
-
-    Syslog.notice "running %s (JOBSERIAL=%s)"
-      job.job_name (string_of_big_int serial) in
-
-  (* Create a temporary directory.  The current directory of the job
-   * will be in this directory.  The directory is removed when the
-   * child process exits.
+  (* Increment JOBSERIAL. *)
+  let serial =
+    match Whenstate.get_variable !state "JOBSERIAL" with
+    | T_int serial ->
+      let serial = succ_big_int serial in
+      state := Whenstate.set_variable !state "JOBSERIAL" (T_int serial);
+      serial
+    | _ -> assert false in
+
+  (* Call the pre-condition script.  Note this may decide not to run
+   * the job by returning false.
    *)
-  let dir = tmpdir () in
-
-  let pid = fork () in
-  if pid = 0 then ( (* child process running the job *)
-    chdir dir;
-
-    (* Set environment variables corresponding to each variable. *)
-    List.iter
-      (fun (name, value) -> putenv name (string_of_variable value))
-      (Whenstate.get_variables !state);
-
-    (* Set the $JOBNAME environment variable. *)
-    putenv "JOBNAME" job.job_name;
-
-    (* Create a temporary file containing the shell script fragment. *)
-    let script = dir // "script.sh" in
-    let chan = open_out script in
-    fprintf chan "set -e\n"; (* So that jobs exit on error. *)
-    output_string chan job.job_script.sh_script;
-    close_out chan;
-    chmod script 0o700;
-
-    let shell = try getenv "SHELL" with Not_found -> "/bin/sh" in
-
-    (* Set output to file. *)
-    let output = dir // "output.txt" in
-    let fd = openfile output [O_WRONLY; O_CREAT; O_TRUNC; O_NOCTTY] 0o600 in
-    dup2 fd stdout;
-    dup2 fd stderr;
-    close fd;
-
-    (* Execute the shell script. *)
-    (try execvp shell [| shell; "-c"; script |];
-     with Unix_error (err, fn, _) ->
-       Syslog.error "%s failed: %s: %s" fn script (error_message err)
+  let pre_condition () =
+    match job.job_pre with
+    | None -> true
+    | Some pre ->
+      let rs = ref [] in
+      IntMap.iter (
+        fun pid (job, _, serial, start_time) ->
+          let r = { pirun_job_name = job.job_name;
+                    pirun_serial = serial;
+                    pirun_start_time = start_time;
+                    pirun_pid = pid } in
+          rs := r :: !rs
+      ) !runningmap;
+      let preinfo = {
+        pi_job_name = job.job_name;
+        pi_serial = serial;
+        pi_variables = Whenstate.get_variables !state;
+        pi_running = !rs;
+      } in
+      pre preinfo
+  in
+  if pre_condition () then (
+    Syslog.notice "running %s (JOBSERIAL=%s)"
+      job.job_name (string_of_big_int serial);
+
+    (* Create a temporary directory.  The current directory of the job
+     * will be in this directory.  The directory is removed when the
+     * child process exits.
+     *)
+    let dir = tmpdir () in
+
+    let pid = fork () in
+    if pid = 0 then ( (* child process running the job *)
+      chdir dir;
+
+      (* Set environment variables corresponding to each variable. *)
+      List.iter
+        (fun (name, value) -> putenv name (string_of_variable value))
+        (Whenstate.get_variables !state);
+
+      (* Set the $JOBNAME environment variable. *)
+      putenv "JOBNAME" job.job_name;
+
+      (* Create a temporary file containing the shell script fragment. *)
+      let script = dir // "script.sh" in
+      let chan = open_out script in
+      fprintf chan "set -e\n"; (* So that jobs exit on error. *)
+      output_string chan job.job_script.sh_script;
+      close_out chan;
+      chmod script 0o700;
+
+      let shell = try getenv "SHELL" with Not_found -> "/bin/sh" in
+
+      (* Set output to file. *)
+      let output = dir // "output.txt" in
+      let fd = openfile output [O_WRONLY; O_CREAT; O_TRUNC; O_NOCTTY] 0o600 in
+      dup2 fd stdout;
+      dup2 fd stderr;
+      close fd;
+
+      (* Execute the shell script. *)
+      (try execvp shell [| shell; "-c"; script |];
+       with Unix_error (err, fn, _) ->
+         Syslog.error "%s failed: %s: %s" fn script (error_message err)
+      );
+      _exit 1
     );
-    _exit 1
-  );
 
-  (* Remember this PID, the job and the temporary directory, so we
-   * can clean up when the child exits.
-   *)
-  running := IntMap.add pid (job, dir) !running
+    (* Remember this PID, the job and the temporary directory, so we
+     * can clean up when the child exits.
+     *)
+    runningmap := IntMap.add pid (job, dir, serial, time ()) !runningmap;
+    serialmap := BigIntMap.add serial pid !serialmap
+  )
+  else (
+    Syslog.notice "not running %s (JOBSERIAL=%s) because pre() condition returned false"
+      job.job_name (string_of_big_int serial);
+  )
 
 and tmpdir () =
   let chan = open_in "/dev/urandom" in
@@ -369,33 +441,36 @@ and handle_sigchld _ =
     let pid, status = waitpid [WNOHANG] 0 in
     if pid > 0 then (
       (* Look up the PID in the running jobs map. *)
-      let job, dir = IntMap.find pid !running in
-      running := IntMap.remove pid !running;
-      cleanup_job job dir status
+      let job, dir, serial, time = IntMap.find pid !runningmap in
+      runningmap := IntMap.remove pid !runningmap;
+      serialmap := BigIntMap.remove serial !serialmap;
+      post_job job dir serial time status
     )
   with Unix_error _ | Not_found -> ()
 
-and cleanup_job job dir status =
-  (* If there is a cleanup function, run it. *)
-  (match job.job_cleanup with
+and post_job job dir serial time status =
+  (* If there is a post function, run it. *)
+  (match job.job_post with
   | None -> ()
-  | Some cleanup ->
+  | Some post ->
     let code =
       match status with
       | WEXITED c -> c
       | WSIGNALED s | WSTOPPED s -> 1 in
     let result = {
       res_job_name = job.job_name;
+      res_serial = serial;
       res_code = code;
       res_tmpdir = dir;
-      res_output = dir // "output.txt"
+      res_output = dir // "output.txt";
+      res_start_time = time
     } in
-    try cleanup result
+    try post result
     with
     | Failure msg ->
-      Syslog.error "job %s cleanup function failed: %s" job.job_name msg
+      Syslog.error "job %s post function failed: %s" job.job_name msg
     | exn ->
-      Syslog.error "job %s cleanup function exception: %s"
+      Syslog.error "job %s post function exception: %s"
         job.job_name (Printexc.to_string exn)
   );