sh: Add optional ?tmpdir parameter to control tmpdir creation.

[goaljobs.git] / goaljobs.ml

diff --git a/goaljobs.ml b/goaljobs.ml
index e83f00e..94843ac 100644 (file)
--- a/goaljobs.ml
+++ b/goaljobs.ml
@@ -23,6 +23,8 @@ open Printf
 
 open Goaljobs_config
 
+type ('a, 'b) alternative = Either of 'a | Or of 'b
+
 let (//) = Filename.concat
 let quote = Filename.quote
 
@@ -222,7 +224,7 @@ let url_contains_string url str =
   r
 
 (* Create a temporary directory.  It is *not* deleted on exit. *)
-let tmpdir () =
+let make_tmpdir () =
   let chan = open_in "/dev/urandom" in
   let data = String.create 16 in
   really_input chan data 0 (String.length data);
@@ -239,66 +241,79 @@ let rm_rf dir =
 
 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
+(* Used by sh, shout, shlines to handle the script and temporary dir. *)
+let with_script ?(tmpdir = false) script f =
+  let dir = if tmpdir then Some (make_tmpdir ()) else None in
+  let script_file, chan =
+    match dir with
+    | Some dir ->
+      let script_file = dir // "script.sh" in
+      let chan = open_out script_file in
+      script_file, chan
+    | None -> Filename.open_temp_file "goaljobsscript" ".sh" in
+  chmod script_file 0o700;
   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 cmd =
+    match dir with
+    | Some dir -> sprintf "cd %s && exec %s" (quote dir) (quote script_file)
+    | None -> sprintf "exec %s" (quote script_file) in
+  let r = try Either (f cmd) with exn -> Or exn in
+  (match dir with
+  | Some dir -> rm_rf dir
+  | None -> ()
+  );
+  match r with
+  | Either x -> x
+  | Or exn -> raise exn
 
-let sh fs =
+let sh ?tmpdir 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
+    with_script ?tmpdir script (
+      fun cmd ->
+        let r = Sys.command cmd in
+        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_shlines ?tmpdir script =
+  with_script ?tmpdir script (
+    fun cmd ->
+      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 -> ());
+      match close_process_in chan 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 ?tmpdir fs = ksprintf (do_shlines ?tmpdir) fs
 
-let do_shout script =
-  let lines = do_shlines script in
+let do_shout ?tmpdir script =
+  let lines = do_shlines ?tmpdir script in
   String.concat "\n" lines
-let shout fs = ksprintf do_shout fs
+let shout ?tmpdir fs = ksprintf (do_shout ?tmpdir) fs
 
 (*
 val replace_substring : string -> string -> string -> string
@@ -321,7 +336,6 @@ val filter_file_extension : string -> string list -> string
  * 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
@@ -338,7 +352,7 @@ let with_memory_locked ?(write = false) f =
   );
 
   (* Run the function. *)
-  let r = try Either (f fd) with exn -> Or (exn) in
+  let r = try Either (f fd) with exn -> Or exn in
   lockf fd F_ULOCK 0;
   match r with
   | Either x -> x