From 2c655f83fc3fe3c6fa2edd340c0f311552a0dfce Mon Sep 17 00:00:00 2001
From: Slobodan Predolac <spredolac@fb.com>
Date: Fri, 25 Jul 2025 10:14:28 -0700
Subject: [PATCH] [process_madvise] Make init lazy so that python tests pass.
 Reset the pidfd on fork

---
 include/jemalloc/internal/atomic.h |  2 ++
 include/jemalloc/internal/pages.h  |  1 +
 src/jemalloc.c                     |  1 +
 src/pages.c                        | 37 ++++++++++++++++++++++++------
 4 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/include/jemalloc/internal/atomic.h b/include/jemalloc/internal/atomic.h
index ddd9341e..f80e5640 100644
--- a/include/jemalloc/internal/atomic.h
+++ b/include/jemalloc/internal/atomic.h
@@ -89,6 +89,8 @@ JEMALLOC_GENERATE_ATOMICS(bool, b, 0)
 
 JEMALLOC_GENERATE_EXPANDED_INT_ATOMICS(unsigned, u, LG_SIZEOF_INT)
 
+JEMALLOC_GENERATE_EXPANDED_INT_ATOMICS(int, i, LG_SIZEOF_INT)
+
 JEMALLOC_GENERATE_EXPANDED_INT_ATOMICS(size_t, zu, LG_SIZEOF_PTR)
 
 JEMALLOC_GENERATE_EXPANDED_INT_ATOMICS(ssize_t, zd, LG_SIZEOF_PTR)
diff --git a/include/jemalloc/internal/pages.h b/include/jemalloc/internal/pages.h
index 31909934..b0cc5bba 100644
--- a/include/jemalloc/internal/pages.h
+++ b/include/jemalloc/internal/pages.h
@@ -127,5 +127,6 @@ bool pages_boot(void);
 void pages_set_thp_state(void *ptr, size_t size);
 void pages_mark_guards(void *head, void *tail);
 void pages_unmark_guards(void *head, void *tail);
+void pages_postfork_child(void);
 
 #endif /* JEMALLOC_INTERNAL_PAGES_EXTERNS_H */
diff --git a/src/jemalloc.c b/src/jemalloc.c
index 876c49e8..4adcbf3c 100644
--- a/src/jemalloc.c
+++ b/src/jemalloc.c
@@ -4549,6 +4549,7 @@ jemalloc_postfork_child(void) {
 	malloc_mutex_postfork_child(tsd_tsdn(tsd), &arenas_lock);
 	tcache_postfork_child(tsd_tsdn(tsd));
 	ctl_postfork_child(tsd_tsdn(tsd));
+	pages_postfork_child();
 }
 
 /******************************************************************************/
diff --git a/src/pages.c b/src/pages.c
index 88301c2b..54678a38 100644
--- a/src/pages.c
+++ b/src/pages.c
@@ -621,7 +621,7 @@ pages_dodump(void *addr, size_t size) {
 #ifdef JEMALLOC_HAVE_PROCESS_MADVISE
 #	include <sys/mman.h>
 #	include <sys/syscall.h>
-static int pidfd;
+static atomic_i_t process_madvise_pidfd = ATOMIC_INIT(-1);
 
 static bool
 init_process_madvise(void) {
@@ -632,11 +632,6 @@ init_process_madvise(void) {
 	if (opt_process_madvise_max_batch > PROCESS_MADVISE_MAX_BATCH_LIMIT) {
 		opt_process_madvise_max_batch = PROCESS_MADVISE_MAX_BATCH_LIMIT;
 	}
-	pid_t pid = getpid();
-	pidfd = syscall(SYS_pidfd_open, pid, 0);
-	if (pidfd == -1) {
-		return true;
-	}
 
 	return false;
 }
@@ -651,12 +646,38 @@ init_process_madvise(void) {
 static bool
 pages_purge_process_madvise_impl(
     void *vec, size_t vec_len, size_t total_bytes) {
-	size_t purged_bytes = (size_t)syscall(JE_SYS_PROCESS_MADVISE_NR, pidfd,
+	int pid_fd = atomic_load_i(&process_madvise_pidfd, ATOMIC_SEQ_CST);
+	while (pid_fd == -1) {
+		int newfd = syscall(SYS_pidfd_open, getpid(), 0);
+		if (newfd == -1) {
+			return true;
+		}
+		if (!atomic_compare_exchange_strong_i(&process_madvise_pidfd,
+						      &pid_fd, newfd,
+						      ATOMIC_SEQ_CST,
+						      ATOMIC_SEQ_CST)) {
+			/* Someone else set the fd, so we close ours */
+			assert(pid_fd != -1);
+			close(newfd);
+		} else {
+			pid_fd = newfd;
+		}
+	}
+	size_t purged_bytes = (size_t)syscall(JE_SYS_PROCESS_MADVISE_NR, pid_fd,
 	    (struct iovec *)vec, vec_len, MADV_DONTNEED, 0);
 
 	return purged_bytes != total_bytes;
 }
 
+void pages_postfork_child(void) {
+	/* Reset the file descriptor we inherited from parent process */
+	int pid_fd = atomic_load_i(&process_madvise_pidfd, ATOMIC_SEQ_CST);
+	if (pid_fd != -1) {
+		atomic_store_i(&process_madvise_pidfd, -1, ATOMIC_SEQ_CST);
+		close(pid_fd);
+	}
+}
+
 #else
 
 static bool
@@ -671,6 +692,8 @@ pages_purge_process_madvise_impl(
 	return true;
 }
 
+void pages_postfork_child(void) {}
+
 #endif
 
 bool