[goaljobs.git] / goaljobs.ml

(* goaljobs
 * Copyright (C) 2013 Red Hat Inc.
 *
 * 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.
 *)

open CalendarLib

open Unix
open Printf

open Goaljobs_config

let (//) = Filename.concat
let quote = Filename.quote

type goal_result_t = Goal_OK | Goal_failed of string
exception Goal_result of goal_result_t

let goal_failed msg = raise (Goal_result (Goal_failed msg))

let target v =
  if v then raise (Goal_result Goal_OK)
let target_all vs = target (List.fold_left (&&) true vs)
let target_exists vs = target (List.fold_left (||) false vs)
let require f = f ()

type period_t = Seconds | Days | Months | Years
let seconds = (1, Seconds)
let sec = seconds and secs = seconds and second = seconds
let minutes = (60, Seconds)
let min = minutes and mins = minutes and minute = minutes
let hours = (3600, Seconds)
let hour = hours
let days = (1, Days)
let day = days
let weeks = (7, Days)
let week = weeks
let months = (1, Months)
let month = months
let years = (1, Years)
let year = years

let periodic_jobs = ref []

(* Register a periodic job. *)
let every ?name i (j, t) f =
  let period = i*j, t in (* 5 minutes -> ((5 * 60), Seconds) *)
  periodic_jobs := (period, (name, f)) :: !periodic_jobs

(* [next_time t period] returns the earliest event of [period]
   strictly after time [t].

   Visualising periods as repeated events on a timeline, this
   returns [t']:

   {v
   events:  ---+---------+---------+---------+---------+---------+-----
   times:          t     t'
   }

   Note that [period_t] events are not necessarily regular.
   eg. The start of a month is not a fixed number of seconds
   after the start of the previous month.  'Epoch' refers
   to the Unix Epoch (ie. 1970-01-01 00:00:00 UTC).

   If [period = i, Seconds i] then events are when
   [t' mod i == 0] when t' is the number of seconds since
   the Epoch.  This returns the next t' > t.

   If [period = i, Days] then events happen at
   midnight UTC every [i] days since the Epoch.
   This returns the next midnight > t.

   If [period = i, Months] then events happen at
   midnight UTC on the 1st day of the month every [i] months
   since the Epoch.  This returns midnight on the
   1st day of the next month > t.

   If [period = i, Years] then events happen at
   midnight UTC on the 1st day of the year when
   [(y - 1970) mod i == 0].  This returns midnight on the
   1st day of the next year > t. *)

let next_time =
  (* Round up 'a' to the next multiple of 'i'. *)
  let round_up_float a i =
    let r = mod_float a i in
    if r = 0. then a +. i else a +. (i -. r)
  and round_up a i =
    let r = a mod i in
    if r = 0 then a + i else a + (i - r)
  in

  fun t -> function
  | (i, Seconds) ->
    let i = float_of_int i in
    round_up_float t i

  | (i, Years) ->
    let tm = gmtime t in

    (* Round 'tm' up to the first day of the next year. *)
    let year = round_up tm.tm_year i in
    let tm = { tm with tm_sec = 0; tm_min = 0; tm_hour = 0;
                       tm_mday = 1; tm_mon = 0; tm_year = year } in
    fst (mktime tm)

  | (i, Days) ->
    let t = Date.from_unixfloat t in
    let t0 = Date.make 1970 1 1 in

    (* Number of whole days since Unix Epoch. *)
    let nb_days = Date.Period.safe_nb_days (Date.sub t t0) in

    let nb_days = round_up nb_days i in
    let t' = Date.add t0 (Date.Period.day nb_days) in
    Date.to_unixfloat t'

  | (i, Months) ->
    (* Calculate number of whole months since Unix Epoch. *)
    let tm = gmtime t in
    let months = 12 * (tm.tm_year - 70) + tm.tm_mon in

    let months = round_up months i in
    let t0 = Date.make 1970 1 1 in
    let t' = Date.add t0 (Date.Period.month months) in
    Date.to_unixfloat t'

let file_exists = Sys.file_exists

let directory_exists path =
  let s =
    try Some (stat path)
    with
    | Unix_error (ENOENT, _, _) -> None
    | Unix_error (err, _, _) ->
      let msg = sprintf "directory_exists: %s: %s" path (error_message err) in
      goal_failed msg in
  match s with
  | Some s -> s.st_kind = S_DIR
  | None -> false

let file_newer_than f1 f2 =
  let stat f =
    try Some (stat f)
    with
    | Unix_error (ENOENT, _, _) -> None
    | Unix_error (err, _, _) ->
      let msg = sprintf "file_newer_than: %s: %s" f (error_message err) in
      goal_failed msg
  in
  let s1 = stat f1 and s2 = stat f2 in
  match s1 with
  | None -> false
  | Some s1 ->
    match s2 with
    | None ->
      let msg = sprintf "file_newer_than: %s: file does not exist" f2 in
      goal_failed msg
    | Some s2 ->
      s1.st_mtime >= s2.st_mtime

let more_recent objs srcs =
  if not (List.for_all file_exists objs) then false
  else (
    List.for_all (
      fun obj -> List.for_all (file_newer_than obj) srcs
    ) objs
  )

let url_exists url =
  (* http://stackoverflow.com/questions/12199059/how-to-check-if-an-url-exists-with-the-shell-and-probably-curl *)
  let cmd =
    sprintf "curl --output /dev/null --silent --head --fail %s" (quote url) in
  match Sys.command cmd with
  | 0 -> true
  | 19|22 -> false
  | r ->
    let msg = sprintf "curl error testing '%s': exit code %d, see curl(1)"
      url r in
    goal_failed msg

let file_contains_string filename str =
  let cmd = sprintf "grep -q %s %s" (quote str) (quote filename) in
  match Sys.command cmd with
  | 0 -> true
  | 1 -> false
  | r ->
    let msg = sprintf "grep error testing for '%s' in '%s' (exit code %d)"
      str filename r in
    goal_failed msg

let url_contains_string url str =
  let tmp = Filename.temp_file "goaljobsurl" "" in
  let cmd =
    sprintf "curl --output %s --silent --fail %s" (quote tmp) (quote url) in
  (match Sys.command cmd with
  | 0 -> ()
  | 19|22 ->
    let msg = sprintf "curl failed to download URL '%s'" url in
    goal_failed msg
  | r ->
    let msg = sprintf "curl error testing '%s': exit code %d, see curl(1)"
      url r in
    goal_failed msg
  );
  let r = file_contains_string tmp str in
  unlink tmp;
  r

(* Create a temporary directory.  It is *not* deleted on exit. *)
let tmpdir () =
  let chan = open_in "/dev/urandom" in
  let data = String.create 16 in
  really_input chan data 0 (String.length data);
  close_in chan;
  let data = Digest.to_hex (Digest.string data) in
  let dir = Filename.temp_dir_name // sprintf "goaljobstmp%s" data in
  mkdir dir 0o700;
  dir

(* Recursively remove directory. *)
let rm_rf dir =
  let cmd = sprintf "rm -rf %s" (quote dir) in
  ignore (Sys.command cmd)

let shell = ref "/bin/sh"

(* Used by sh, shout etc.  Create a temporary directory and a
 * 'script.sh' inside it containing the script to run.  Returns the
 * temporary directory and command to run.
 *)
let create_script script =
  let dir = tmpdir () in
  let script_file = dir // "script.sh" in
  let chan = open_out script_file in
  fprintf chan "#!%s\n" !shell;
  fprintf chan "set -e\n"; (* so that job exits on error *)
  fprintf chan "set -x\n"; (* echo commands (must be last) *)
  fprintf chan "\n";
  output_string chan script;
  close_out chan;
  chmod script_file 0o700;
  let cmd = sprintf "cd %s && exec %s" (quote dir) (quote script_file) in
  dir, cmd

let sh fs =
  let do_sh script =
    let dir, cmd = create_script script in
    let r = Sys.command cmd in
    rm_rf dir;
    if r <> 0 then (
      let msg = sprintf "external command failed with code %d" r in
      goal_failed msg
    )
  in
  ksprintf do_sh fs

let do_shlines script =
  let dir, cmd = create_script script in
  let chan = open_process_in cmd in
  let lines = ref [] in
  let rec loop () =
    let line = input_line chan in
    eprintf "%s\n%!" line;
    lines := line :: !lines;
    loop ()
  in
  (try loop () with End_of_file -> ());
  let r = close_process_in chan in
  rm_rf dir;
  match r with
  | WEXITED 0 -> List.rev !lines
  | WEXITED i ->
    let msg = sprintf "external command failed with code %d" i in
    goal_failed msg
  | WSIGNALED i ->
    let msg = sprintf "external command was killed by signal %d" i in
    goal_failed msg
  | WSTOPPED i ->
    let msg = sprintf "external command was stopped by signal %d" i in
    goal_failed msg
let shlines fs = ksprintf do_shlines fs

let do_shout script =
  let lines = do_shlines script in
  String.concat "\n" lines
let shout fs = ksprintf do_shout fs

(*
val replace_substring : string -> string -> string -> string
*)

let change_file_extension ext filename =
  let i =
    try String.rindex filename '.'
    with Not_found -> String.length filename in
  String.sub filename 0 i ^ "." ^ ext

(*
val filter_file_extension : string -> string list -> string
*)

(* Persistent memory is stored in $HOME/.goaljobs-memory.  We have to
 * lock this file each time we read or write because multiple concurrent
 * jobs may access it at the same time.
 *
 * XXX Replace this with a more efficient and less fragile implementation.
 *)

type ('a, 'b) alternative = Either of 'a | Or of 'b
let with_memory_locked ?(write = false) f =
  let filename = getenv "HOME" // ".goaljobs-memory" in
  let fd = openfile filename [O_RDWR; O_CREAT] 0o644 in
  lockf fd (if write then F_LOCK else F_RLOCK) 0;
  (* If the file is newly created with zero size, write an
   * empty hash table.
   *)
  if (fstat fd).st_size = 0 then (
    let empty : (string, string) Hashtbl.t = Hashtbl.create 13 in
    let chan = out_channel_of_descr fd in
    output_value chan empty;
    Pervasives.flush chan;
    ignore (lseek fd 0 SEEK_SET)
  );

  (* Run the function. *)
  let r = try Either (f fd) with exn -> Or (exn) in
  lockf fd F_ULOCK 0;
  match r with
  | Either x -> x
  | Or exn -> raise exn

let memory_exists key =
  with_memory_locked (
    fun fd ->
      let chan = in_channel_of_descr fd in
      let memory : (string, string) Hashtbl.t = input_value chan in
      Hashtbl.mem memory key
  )

let memory_get key =
  with_memory_locked (
    fun fd ->
      let chan = in_channel_of_descr fd in
      let memory : (string, string) Hashtbl.t = input_value chan in
      try Some (Hashtbl.find memory key) with Not_found -> None
  )

let memory_set key value =
  with_memory_locked ~write:true (
    fun fd ->
      let chan = in_channel_of_descr fd in
      let memory : (string, string) Hashtbl.t = input_value chan in
      Hashtbl.replace memory key value;
      let chan = out_channel_of_descr fd in
      seek_out chan 0;
      output_value chan memory;
      Pervasives.flush chan;
  )

let memory_delete key =
  with_memory_locked ~write:true (
    fun fd ->
      let chan = in_channel_of_descr fd in
      let memory : (string, string) Hashtbl.t = input_value chan in
      Hashtbl.remove memory key;
      let chan = out_channel_of_descr fd in
      seek_out chan 0;
      output_value chan memory;
      Pervasives.flush chan;
  )

let published_goals = ref []
let publish name fn = published_goals := (name, fn) :: !published_goals
let get_goal name =
  try Some (List.assoc name !published_goals) with Not_found -> None

let goal_file_exists filename =
  if not (file_exists filename) then (
    let msg = sprintf "file '%s' required but not found" filename in
    goal_failed msg
  )
let goal_directory_exists path =
  if not (directory_exists path) then (
    let msg = sprintf "directory '%s' required but not found" path in
    goal_failed msg
  )
let goal_file_newer_than f1 f2 =
  if not (file_newer_than f1 f2) then (
    let msg = sprintf "file %s is required to be newer than %s" f1 f2 in
    goal_failed msg
  )
let goal_more_recent objs srcs =
  if not (more_recent objs srcs) then (
    let msg = sprintf "object(s) %s are required to be newer than source(s) %s"
      (String.concat " " objs) (String.concat " " srcs) in
    goal_failed msg
  )
let goal_url_exists url =
  if not (url_exists url) then (
    let msg = sprintf "url_exists: URL '%s' required but does not exist" url in
    goal_failed msg
  )
let goal_file_contains_string filename str =
  if not (file_contains_string filename str) then (
    let msg = sprintf "file_contains_string: file '%s' is required to contain string '%s'" filename str in
    goal_failed msg
  )
let goal_url_contains_string url str =
  if not (url_contains_string url str) then (
    let msg = sprintf "url_contains_string: URL '%s' is required to contain string '%s'" url str in
    goal_failed msg
  )
let goal_memory_exists k =
  if not (memory_exists k) then (
    let msg = sprintf "memory_exists: key '%s' required but does not exist" k in
    goal_failed msg
  )

let guard fn arg =
  try fn arg; true
  with
  | Goal_result (Goal_failed msg) ->
    prerr_endline ("error: " ^ msg);
    false
  | exn ->
    prerr_endline (Printexc.to_string exn);
    false

(* Run the program. *)
let rec init () =
  let prog = Sys.executable_name in
  let prog = Filename.basename prog in

  (* Save the current working directory when the program started. *)
  putenv "builddir" (getcwd ());

  let args = ref [] in

  let display_version () =
    printf "%s %s\n" package_name package_version;
    exit 0
  in

  let list_goals () =
    let names = !published_goals in
    let names = List.map fst names in
    let names = List.sort compare names in
    List.iter print_endline names;
    exit 0
  in

  let argspec = Arg.align [
    "--goals", Arg.Unit list_goals, " List all goals";
    "-l", Arg.Unit list_goals, " List all goals";
    "-V", Arg.Unit display_version, " Display version number and exit";
    "--version", Arg.Unit display_version, " Display version number and exit";
  ] in
  let anon_fun str = args := str :: !args in
  let usage_msg = sprintf "\
%s: a script generated by goaljobs

List all goals:                %s -l
Run a single goal like this:   %s <name-of-goal> [<goal-args ...>]

For more information see the goaljobs(1) man page.

Options:
" prog prog prog in

  Arg.parse argspec anon_fun usage_msg;

  let args = List.rev !args in

  (* Was a goal named on the command line? *)
  (match args with
  | name :: args ->
    (match get_goal name with
    | Some fn ->
      exit (if guard fn args then 0 else 1)
    | None ->
      eprintf "error: no goal called '%s' was found.\n" name;
      eprintf "Use %s -l to list all published goals in this script.\n" name;
      exit 1
    )
  | [] ->
    (* If periodic jobs exist, fall through. *)
    if !periodic_jobs = [] then (
      (* Does a published 'all' goal exist? *)
      match get_goal "all" with
      | Some fn ->
        exit (if guard fn [] then 0 else 1)
      | None ->
        (* No published 'all' goal. *)
        eprintf "error: no goal called 'all' was found.\n";
        exit 1
    )
  );

  assert (!periodic_jobs <> []);

  (* Run the periodic jobs.  Note these run forever, or until killed. *)
  while true do
    (* Find the next job to run. *)
    let now = time () in
    let jobs = List.map (
      fun (period, (_, _ as name_f)) ->
        next_time now period, name_f
    ) !periodic_jobs in
    let jobs = List.sort (fun (t1,_) (t2,_) -> compare t1 t2) jobs in

    (* Find all jobs that have the same next time.
     * XXX When we can handle parallel jobs we can do better here,
     * but until them run all the ones which have the same time
     * in series.
     *)
    let next_t = int_of_float (fst (List.hd jobs)) in
    let jobs = List.filter (fun (t, _) -> int_of_float t = next_t) jobs in

    (* Run next job(s) after waiting for the appropriate amount of time. *)
    let seconds = next_t - int_of_float now in
    eprintf "next job will run in %s\n%!" (printable_seconds seconds);
    sleep seconds;

    List.iter (
      fun (_, (name, f)) ->
        eprintf "running job: %s\n%!"
          (match name with Some name -> name | None -> "[unnamed]");
        ignore (guard f ())
    ) jobs
  done

and printable_seconds s =
  if s < 60 then sprintf "%d seconds" s
  else if s < 6000 then sprintf "%d minutes, %d seconds" (s/60) (s mod 60)
  else if s < 86400 then sprintf "%d hours, %d minutes" (s/3600) (s/60)
  else sprintf "about %d days" (s/86400)