Reset the prefault helper's state in forked children (#2391)

packages/runtime/src/scr_alloc.c

@@ -368,6 +368,34 @@ static bool scr_sa_pf_on(void) {
   return s > 0;
 }
 
+/* fork() copies only the calling thread, so a child has no helper. The
+ * prepare handler takes the lock, so no other thread holds it across the
+ * fork. The child handler gives the child fresh helper state: an empty ring,
+ * no busy range or waiters, new conditions (the parent's can still record
+ * the helper as a waiter), and the heap watermarks cleared. Its first
+ * request then starts a helper of its own. */
+static bool scr_sa_pf_atfork;
+
+static void scr_sa_pf_fork_prepare(void) { pthread_mutex_lock(&scr_sa_pf_lock); }
+
+static void scr_sa_pf_fork_parent(void) { pthread_mutex_unlock(&scr_sa_pf_lock); }
+
+static void scr_sa_pf_fork_child(void) {
+  (void)pthread_mutex_init(&scr_sa_pf_lock, NULL);
+  (void)pthread_cond_init(&scr_sa_pf_cond, NULL);
+  (void)pthread_cond_init(&scr_sa_pf_done, NULL);
+  scr_sa_pf_head = scr_sa_pf_tail = 0;
+  scr_sa_pf_busy = (ScrSaPfRange){NULL, NULL};
+  scr_sa_pf_forgetting = 0;
+  scr_sa_pf_waiting = false;
+  scr_sa_pf_started = false;
+#ifdef SCR_SA_PF_HEAP
+  scr_sa_pf_heap_mode = 0;
+  scr_sa_pf_brk0 = scr_sa_pf_brk = NULL;
+  scr_sa_pf_heap_big = false;
+#endif
+}
+
 /* Asks the helper to fault in [lo, hi) (page aligned). Dropped when the
  * ring is full: the owner then faults those pages itself. */
 static void scr_sa_prefault(char *lo, char *hi) {
@@ -375,6 +403,14 @@ static void scr_sa_prefault(char *lo, char *hi) {
   pthread_mutex_lock(&scr_sa_pf_lock);
   if (!scr_sa_pf_started) {
     scr_sa_pf_started = true;
+    /* Once per process: a child inherits the parent's handlers. */
+    if (!scr_sa_pf_atfork &&
+        pthread_atfork(scr_sa_pf_fork_prepare, scr_sa_pf_fork_parent, scr_sa_pf_fork_child) != 0) {
+      pthread_mutex_unlock(&scr_sa_pf_lock);
+      atomic_store_explicit(&scr_sa_pf_state, -1, memory_order_relaxed);
+      return;
+    }
+    scr_sa_pf_atfork = true;
 #ifdef SCR_SA_PF_HEAP
     /* Nobody waits on the condition yet: give it the monotonic clock for
      * the helper's polls (no polls without it). */

packages/runtime/src/scr_alloc.test.c

@@ -1023,6 +1023,65 @@ static void check_discard(void) {
 }
 #endif
 
+/* ThreadSanitizer stops a child that starts a thread after a fork in a
+ * threaded program, so its build skips the fork check. */
+#define FORK_CHECK 1
+#if defined(__SANITIZE_THREAD__)
+#undef FORK_CHECK
+#elif defined(__has_feature)
+#if __has_feature(thread_sanitizer)
+#undef FORK_CHECK
+#endif
+#endif
+
+#if defined(__linux__) && SCR_SMALL_ALLOC && defined(FORK_CHECK)
+#include <signal.h>
+#include <sys/mman.h>
+#include <sys/wait.h>
+#include <time.h>
+#include <unistd.h>
+#define FORK_ROUNDS 64
+#define FORK_SPAN ((size_t)64 << 20)
+#define FORK_OWN ((size_t)1 << 20)
+
+/* A child forked while the prefault helper faults in a range has no helper.
+ * It releases that range (which waits for a busy range to finish), then
+ * requests a fault-in of its own and releases it, and must exit within five
+ * seconds. */
+static void check_fork_while_prefaulting(void) {
+  for (unsigned round = 0; round < FORK_ROUNDS; round++) {
+    char *lo = mmap(NULL, FORK_SPAN, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+    assert(lo != MAP_FAILED);
+    scr_mem_prefault(lo, lo + FORK_SPAN);
+    pid_t pid = fork();
+    assert(pid >= 0);
+    if (pid == 0) {
+      scr_mem_release_ahead(lo, lo + FORK_SPAN);
+      char *own = mmap(NULL, FORK_OWN, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+      if (own == MAP_FAILED) _exit(2);
+      scr_mem_prefault(own, own + FORK_OWN);
+      scr_mem_release_ahead(own, own + FORK_OWN);
+      _exit(0);
+    }
+    int status = 0;
+    pid_t done = 0;
+    for (unsigned ms = 0; ms < 5000 && done == 0; ms++) {
+      done = waitpid(pid, &status, WNOHANG);
+      if (done == 0) nanosleep(&(struct timespec){0, 1000000L}, NULL);
+    }
+    if (done == 0) {
+      kill(pid, SIGKILL);
+      (void)waitpid(pid, &status, 0);
+      fprintf(stderr, "the child forked in round %u did not exit\n", round);
+      abort();
+    }
+    assert(done == pid && WIFEXITED(status) && WEXITSTATUS(status) == 0);
+    scr_mem_release_ahead(lo, lo + FORK_SPAN);
+    assert(munmap(lo, FORK_SPAN) == 0);
+  }
+}
+#endif
+
 #if defined(__linux__) && SCR_SMALL_ALLOC
 /* The program-lifetime region hands out zeroed, 8-byte aligned blocks below
  * 32 GiB, in order from the calling thread's chunk; a block larger than a
@@ -1087,6 +1146,9 @@ int main(void) {
   check_system_pointers();
   check_churn();
   check_free_order();
+#if defined(__linux__) && SCR_SMALL_ALLOC && defined(FORK_CHECK)
+  check_fork_while_prefaulting();
+#endif
 #if SCR_SMALL_ALLOC && defined(__linux__) && defined(__GLIBC__)
   check_heap_huge_pages();
 #endif