Fork actions must not allocate

lib_eio/unix/fork_action.c

@@ -1,3 +1,9 @@
+/* Note: fork actions MUST NOT allocate (either on the OCaml heap or with C malloc).
+ * This is because e.g. we might have forked while another thread in the parent had a lock.
+ * In the child, we inherit a copy of the locked mutex, but no corresponding thread to
+ * release it.
+ */
+
 #include <stdlib.h>
 #include <unistd.h>
 #include <fcntl.h>
@@ -6,6 +12,9 @@
 
 #include <caml/mlvalues.h>
 #include <caml/unixsupport.h>
+#include <caml/memory.h>
+#include <caml/custom.h>
+#include <caml/fail.h>
 
 #include "fork_action.h"
 
@@ -42,24 +51,61 @@ void eio_unix_fork_error(int fd, char *fn, char *buf) {
   try_write_all(fd, buf);
 }
 
-static char **make_string_array(int errors, value v_array) {
-  int n = Wosize_val(v_array);
-  char **c = calloc(sizeof(char *), (n + 1));
-  if (!c) {
-    eio_unix_fork_error(errors, "make_string_array", "out of memory");
-    _exit(1);
-  }
+#define String_array_val(v) *((char ***)Data_custom_val(v))
+
+static void finalize_string_array(value v) {
+  free(String_array_val(v));
+  String_array_val(v) = NULL;
+}
+
+static struct custom_operations string_array_ops = {
+  "string.array",
+  finalize_string_array,
+  custom_compare_default,
+  custom_hash_default,
+  custom_serialize_default,
+  custom_deserialize_default,
+  custom_compare_ext_default,
+  custom_fixed_length_default
+};
+
+CAMLprim value eio_unix_make_string_array(value v_len) {
+  CAMLparam0();
+  CAMLlocal1(v_str_array);
+  int n = Int_val(v_len);
+  uintnat total;
+
+  if (caml_umul_overflow(sizeof(char *), n + 1, &total))
+    caml_raise_out_of_memory();
+
+  v_str_array = caml_alloc_custom_mem(&string_array_ops, sizeof(char ***), total);
+
+  char **c = calloc(sizeof(char *), n + 1);
+  String_array_val(v_str_array) = c;
+  if (!c)
+    caml_raise_out_of_memory();
+
+  CAMLreturn(v_str_array);
+}
+
+static void fill_string_array(char **c, value v_ocaml_array) {
+  int n = Wosize_val(v_ocaml_array);
+
   for (int i = 0; i < n; i++) {
-    c[i] = (char *) String_val(Field(v_array, i));
+    c[i] = (char *) String_val(Field(v_ocaml_array, i));
   }
+
   c[n] = NULL;
-  return c;
 }
 
 static void action_execve(int errors, value v_config) {
   value v_exe = Field(v_config, 1);
-  char **argv = make_string_array(errors, Field(v_config, 2));
-  char **envp = make_string_array(errors, Field(v_config, 3));
+  char **argv = String_array_val(Field(v_config, 2));
+  char **envp = String_array_val(Field(v_config, 4));
+
+  fill_string_array(argv, Field(v_config, 3));
+  fill_string_array(envp, Field(v_config, 5));
+
   execve(String_val(v_exe), argv, envp);
   eio_unix_fork_error(errors, "execve", strerror(errno));
   _exit(1);

lib_eio/unix/fork_action.ml

@@ -17,9 +17,14 @@ let rec with_actions actions fn =
     with_actions xs @@ fun c_actions ->
     fn (c_action :: c_actions)
 
+type c_array
+external make_string_array : int -> c_array = "eio_unix_make_string_array"
 external action_execve : unit -> fork_fn = "eio_unix_fork_execve"
 let action_execve = action_execve ()
-let execve path ~argv ~env = { run = fun k -> k (Obj.repr (action_execve, path, argv, env)) }
+let execve path ~argv ~env =
+  let argv_c_array = make_string_array (Array.length argv) in
+  let env_c_array = make_string_array (Array.length env) in
+  { run = fun k -> k (Obj.repr (action_execve, path, argv_c_array, argv, env_c_array, env)) }
 
 external action_chdir : unit -> fork_fn = "eio_unix_fork_chdir"
 let action_chdir = action_chdir ()

lib_eio/unix/include/fork_action.h

@@ -1,7 +1,8 @@
 #include <caml/mlvalues.h>
 #include <caml/alloc.h>
 
-/* A function that runs in the forked child process. It must not run any OCaml code or invoke the GC.
+/* A function that runs in the forked child process.
+ * It must not run any OCaml code, invoke the GC, or even call [malloc].
  * If the action fails then it writes an error message to the FD [errors] and calls [_exit].
  * v_args is the c_action tuple (where field 0 is the function itself).
  */

lib_eio_linux/eio_stubs.c

@@ -101,7 +101,7 @@ CAMLprim value caml_eio_getrandom(value v_ba, value v_off, value v_len) {
   ssize_t off = (ssize_t)Long_val(v_off);
   ssize_t len = (ssize_t)Long_val(v_len);
   do {
-    void *buf = Caml_ba_data_val(v_ba) + off;
+    void *buf = (char *)Caml_ba_data_val(v_ba) + off;
     caml_enter_blocking_section();
 #if __GLIBC__ > 2 || __GLIBC_MINOR__ > 24
     ret = getrandom(buf, len, 0);

stress/stress_proc.ml

@@ -1,26 +1,35 @@
 open Eio.Std
 
+let n_domains = 4
 let n_rounds = 100
-let n_procs_per_round = 100
+let n_procs_per_round_per_domain = 100 / n_domains
 
-let main mgr =
+let run_in_domain mgr =
   let echo n = Eio.Process.parse_out mgr Eio.Buf_read.line ["sh"; "-c"; "echo " ^ string_of_int n] in
+  Switch.run @@ fun sw ->
+  for j = 1 to n_procs_per_round_per_domain do
+    Fiber.fork ~sw (fun () ->
+        let result = echo j in
+        assert (int_of_string result = j);
+        (* traceln "OK: %d" j *)
+      )
+  done
+
+let main ~dm mgr =
   let t0 = Unix.gettimeofday () in
   for i = 1 to n_rounds do
-    Switch.run @@ fun sw ->
-    for j = 1 to n_procs_per_round do
-      Fiber.fork ~sw (fun () ->
-          let result = echo j in
-          assert (int_of_string result = j);
-          (* traceln "OK: %d" j *)
-        )
-    done;
-    if false then traceln "Finished round %d/%d" i n_rounds
+    Switch.run (fun sw ->
+        for _ = 1 to n_domains - 1 do
+          Fiber.fork ~sw (fun () -> Eio.Domain_manager.run dm (fun () -> run_in_domain mgr))
+        done;
+        Fiber.fork ~sw (fun () -> run_in_domain mgr);
+      );
+    if true then traceln "Finished round %d/%d" i n_rounds
   done;
   let t1 = Unix.gettimeofday () in
-  let n_procs = n_rounds * n_procs_per_round in
-  traceln "Finished process stress test: ran %d processes in %.2fs" n_procs (t1 -. t0)
+  let n_procs = n_rounds * n_procs_per_round_per_domain * n_domains in
+  traceln "Finished process stress test: ran %d processes in %.2fs (using %d domains)" n_procs (t1 -. t0) n_domains
 
 let () =
   Eio_main.run @@ fun env ->
-  main env#process_mgr
+  main ~dm:env#domain_mgr  env#process_mgr