Merge branch 'bpf-arena-handle-memory-max-on-fault-in-with-reclaim'

Jiayuan Chen says:

====================
bpf: arena: handle memory.max on fault-in with reclaim

Since commit e66fe1bc6d ("bpf: arena: Reintroduce memcg accounting"),
arena pages are charged to the memcg of the process that created the arena.
That exposes a problem in the arena user page fault path: the fault-in
allocation runs under arena->spinlock, so it can only use the non-blocking
allocator, which never reclaims. Once memory.current is at memory.max the
allocation simply fails, even when the memcg is full of page cache that
could be dropped right away. Reaching memory.max is completely normal for a
healthy application - e.g. reading a large file fills memory.current with
page cache - and the process then gets SIGSEGV on a perfectly valid arena
address.

Preallocate the page outside the lock (patch 1), the way do_anonymous_page()
does, so the allocation can sleep and reclaim. It never invokes the OOM
killer: the page is charged to the map's memcg, which need not be the
faulting task's, so an OOM there could kill unrelated tasks in the map's
cgroup. On a genuine failure the fault returns VM_FAULT_SIGBUS. patch 2&3
adds a selftest that fills a memcg with reclaimable page cache and faults an
arena in under memory.max: without the fix the child gets SIGSEGV on a valid
address, with it the fault-in reclaims and succeeds.

v8 -> v9:
   - Trim the two comments in arena_vm_fault() to a sentence or two; the
     details are in the commit log. (Alexei Starovoitov)

v7 -> v8:
   - Merge patch 1 into patch 2. (Alexei Starovoitov)

v4 -> v7:
   - Simplify the implementation: do not aim for OOM anymore. An arena is
     shared between processes and can be shared across cgroups, so the OOM
     killer is the wrong tool here - it would act on the map's memcg, which
     need not be the faulting task's. Only try to reclaim now, via
     __GFP_RETRY_MAYFAIL. Everything else is kept as before.
   - selftest: changed accordingly - fill the memcg with reclaimable page
     cache and check that the arena fault-in succeeds by reclaiming it,
     instead of relying on an OOM kill.

v3 -> v4:
   - rebase bpf-next and fix conflict
   - add Reviewed-by tag from Emil Tsalapatis

v2 -> v3:
   - selftest: check the memcg OOM via memory.events "oom_kill" instead of
     the exit signal; it only aims to pass on the fixed kernel, since the
     unfixed SIGSEGV is racy.

v1 -> v2:
   - Rebase on the separate deadlock fix (found by the Sashiko AI review),
     now applied to bpf-next.
   - Honor the map's NUMA node on fault-in.
   - Return VM_FAULT_SIGBUS for the non-recoverable faults (lock, range-tree
     and page-table failures); a scratch-page hole stays VM_FAULT_SIGSEGV
     only under BPF_F_SEGV_ON_FAULT. (Kumar Kartikeya Dwivedi)
   - Add read_cgroup_file() to cgroup_helpers instead of open-coding the
     /mnt/... path in the test. (Emil Tsalapatis)
   - Dump the cgroup memory stats on test failure to ease debugging.

v8:
https://lore.kernel.org/bpf/20260918084123.39477-1-jiayuan.chen@linux.dev/
v7:
https://lore.kernel.org/bpf/20260917015618.7488-1-jiayuan.chen@linux.dev/#r
v4:
https://lore.kernel.org/bpf/20260821050250.35112-1-jiayuan.chen@linux.dev/T/#t
v2:
https://lore.kernel.org/bpf/20260805091720.139924-1-jiayuan.chen@linux.dev/
v1:
https://lore.kernel.org/bpf/20260727062521.376231-1-jiayuan.chen@linux.dev/
====================

Link: https://patch.msgid.link/20260922101831.192102-1-jiayuan.chen@linux.dev
Signed-off-by: Alexei Starovoitov <ast@kernel.org>

kernel/bpf/arena.c

@@ -481,7 +481,8 @@ static vm_fault_t arena_vm_fault(struct vm_fault *vmf)
 	struct bpf_map *map = vmf->vma->vm_file->private_data;
 	struct bpf_arena *arena = container_of(map, struct bpf_arena, map);
 	struct mem_cgroup *new_memcg, *old_memcg;
-	struct page *page;
+	struct page *page, *new_page = NULL;
+	vm_fault_t fault_ret;
 	long kbase, kaddr;
 	unsigned long flags;
 	int ret;
@@ -489,59 +490,101 @@ static vm_fault_t arena_vm_fault(struct vm_fault *vmf)
 	kbase = bpf_arena_get_kern_vm_start(arena);
 	kaddr = kbase + (u32)(vmf->address);
 
-	if (raw_res_spin_lock_irqsave(&arena->spinlock, flags))
+	page = vmalloc_to_page((void *)kaddr);
+	if (!page && !(arena->map.map_flags & BPF_F_SEGV_ON_FAULT)) {
+		/*
+		 * Preallocate outside the lock so the allocation can reclaim;
+		 * __GFP_RETRY_MAYFAIL keeps the OOM killer out of it.
+		 */
+		bpf_map_memcg_enter(&arena->map, &old_memcg, &new_memcg);
+		new_page = alloc_pages_node(map->numa_node,
+					    GFP_KERNEL | __GFP_ZERO |
+					    __GFP_ACCOUNT | __GFP_NOWARN |
+					    __GFP_RETRY_MAYFAIL, 0);
+		bpf_map_memcg_exit(old_memcg, new_memcg);
+	}
+
+	if (raw_res_spin_lock_irqsave(&arena->spinlock, flags)) {
 		/*
 		 * A failed lock means a possible deadlock was detected. Don't
 		 * return VM_FAULT_RETRY: this handler never took mmap_lock, but
 		 * the fault path would re-take it on retry and deadlock. Fail.
 		 */
+		if (new_page)
+			free_pages_nolock(new_page, 0);
 		return VM_FAULT_SIGBUS;
+	}
 
 	page = vmalloc_to_page((void *)kaddr);
 	if (page) {
-		if (page == arena->scratch_page)
+		if (page == arena->scratch_page) {
 			/* BPF triggered scratch here; don't lazy-alloc over it */
-			goto out_sigsegv;
+			fault_ret = (arena->map.map_flags & BPF_F_SEGV_ON_FAULT)
+				    ? VM_FAULT_SIGSEGV : VM_FAULT_SIGBUS;
+			goto out_err_locked;
+		}
 		/* already have a page vmap-ed */
 		goto out;
 	}
 
+	if (arena->map.map_flags & BPF_F_SEGV_ON_FAULT) {
+		/*
+		 * User space requested to segfault when page is not allocated
+		 * by bpf prog
+		 */
+		fault_ret = VM_FAULT_SIGSEGV;
+		goto out_err_locked;
+	}
+
 	bpf_map_memcg_enter(&arena->map, &old_memcg, &new_memcg);
 
-	if (arena->map.map_flags & BPF_F_SEGV_ON_FAULT)
-		/* User space requested to segfault when page is not allocated by bpf prog */
-		goto out_sigsegv_memcg;
+	if (!new_page) {
+		/*
+		 * The probed page was freed meanwhile or preallocation failed;
+		 * try the non-blocking allocator, we cannot sleep here.
+		 */
+		ret = bpf_map_alloc_pages(map, map->numa_node, 1, &new_page);
+		if (ret) {
+			fault_ret = VM_FAULT_SIGBUS;
+			goto out_err_locked_memcg;
+		}
+	}
 
 	ret = range_tree_clear(&arena->rt, vmf->pgoff, 1);
-	if (ret)
-		goto out_sigsegv_memcg;
-
-	struct apply_range_data data = { .arena = arena, .pages = &page, .i = 0 };
-	/* Account into memcg of the process that created bpf_arena */
-	ret = bpf_map_alloc_pages(map, NUMA_NO_NODE, 1, &page);
 	if (ret) {
-		range_tree_set(&arena->rt, vmf->pgoff, 1);
-		goto out_sigsegv_memcg;
+		fault_ret = VM_FAULT_SIGBUS;
+		goto out_err_locked_memcg;
 	}
+	struct apply_range_data data = {
+		.arena = arena, .pages = &new_page, .i = 0
+	};
 
 	ret = apply_to_page_range(&init_mm, kaddr, PAGE_SIZE, apply_range_set_cb, &data);
 	if (ret) {
 		range_tree_set(&arena->rt, vmf->pgoff, 1);
-		free_pages_nolock(page, 0);
-		goto out_sigsegv_memcg;
+		fault_ret = VM_FAULT_SIGBUS;
+		goto out_err_locked_memcg;
 	}
 	flush_vmap_cache(kaddr, PAGE_SIZE);
 	bpf_map_memcg_exit(old_memcg, new_memcg);
+	/* new_page was consumed */
+	page = new_page;
+	new_page = NULL;
 out:
 	page_ref_add(page, 1);
 	raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
+	if (new_page)
+		free_pages_nolock(new_page, 0);
 	vmf->page = page;
 	return 0;
-out_sigsegv_memcg:
+
+out_err_locked_memcg:
 	bpf_map_memcg_exit(old_memcg, new_memcg);
-out_sigsegv:
+out_err_locked:
 	raw_res_spin_unlock_irqrestore(&arena->spinlock, flags);
-	return VM_FAULT_SIGSEGV;
+	if (new_page)
+		free_pages_nolock(new_page, 0);
+	return fault_ret;
 }
 
 static const struct vm_operations_struct arena_vm_ops = {

tools/testing/selftests/bpf/cgroup_helpers.c

@@ -188,6 +188,73 @@ int write_cgroup_file_parent(const char *relative_path, const char *file,
 	return __write_cgroup_file(cgroup_path, file, buf);
 }
 
+static int __read_cgroup_file(const char *cgroup_path, const char *file,
+			      char *buf, size_t len)
+{
+	char file_path[PATH_MAX + 1];
+	ssize_t got;
+	int fd;
+
+	snprintf(file_path, sizeof(file_path), "%s/%s", cgroup_path, file);
+	fd = open(file_path, O_RDONLY);
+	if (fd < 0) {
+		log_err("Opening %s", file_path);
+		return 1;
+	}
+
+	got = read(fd, buf, len - 1);
+	if (got < 0) {
+		log_err("Reading %s", file_path);
+		close(fd);
+		return 1;
+	}
+	buf[got] = '\0';
+	close(fd);
+	return 0;
+}
+
+/**
+ * read_cgroup_file() - Read from a cgroup file
+ * @relative_path: The cgroup path, relative to the workdir
+ * @file: The name of the file in cgroupfs to read from
+ * @buf: Buffer to read into, NUL-terminated on success
+ * @len: Size of @buf
+ *
+ * Read from a file in the given cgroup's directory.
+ *
+ * If successful, 0 is returned.
+ */
+int read_cgroup_file(const char *relative_path, const char *file,
+		     char *buf, size_t len)
+{
+	char cgroup_path[PATH_MAX - 24];
+
+	format_cgroup_path(cgroup_path, relative_path);
+	return __read_cgroup_file(cgroup_path, file, buf, len);
+}
+
+/**
+ * read_cgroup_file_parent() - Read from a cgroup file in the parent process
+ *                             workdir
+ * @relative_path: The cgroup path, relative to the parent process workdir
+ * @file: The name of the file in cgroupfs to read from
+ * @buf: Buffer to read into, NUL-terminated on success
+ * @len: Size of @buf
+ *
+ * Read from a file in the given cgroup's directory under the parent process
+ * workdir.
+ *
+ * If successful, 0 is returned.
+ */
+int read_cgroup_file_parent(const char *relative_path, const char *file,
+			    char *buf, size_t len)
+{
+	char cgroup_path[PATH_MAX - 24];
+
+	format_parent_cgroup_path(cgroup_path, relative_path);
+	return __read_cgroup_file(cgroup_path, file, buf, len);
+}
+
 /**
  * setup_cgroup_environment() - Setup the cgroup environment
  *

tools/testing/selftests/bpf/cgroup_helpers.h

@@ -15,6 +15,10 @@ int write_cgroup_file(const char *relative_path, const char *file,
 		      const char *buf);
 int write_cgroup_file_parent(const char *relative_path, const char *file,
 			     const char *buf);
+int read_cgroup_file(const char *relative_path, const char *file,
+		     char *buf, size_t len);
+int read_cgroup_file_parent(const char *relative_path, const char *file,
+			    char *buf, size_t len);
 int cgroup_setup_and_join(const char *relative_path);
 int get_root_cgroup(void);
 int create_and_get_cgroup(const char *relative_path);

tools/testing/selftests/bpf/prog_tests/arena_memcg.c

@@ -0,0 +1,196 @@
+// SPDX-License-Identifier: GPL-2.0
+#define _GNU_SOURCE
+#include <test_progs.h>
+#include <fcntl.h>
+#include <linux/magic.h>
+#include <sys/vfs.h>
+#include <sys/wait.h>
+#include <unistd.h>
+#include <sys/user.h>
+#ifndef PAGE_SIZE /* on some archs it comes in sys/user.h */
+#define PAGE_SIZE getpagesize()
+#endif
+
+#include "cgroup_helpers.h"
+#include "arena_memcg.skel.h"
+
+#define CG_PATH		"/arena_memcg"
+
+/* Reclaimable page cache the child builds up before it gets capped. */
+#define RECLAIMABLE	(128 * 1024 * 1024)
+/* Headroom left under memory.max, far less than the arena we fault in. */
+#define HEADROOM	(8 * 1024 * 1024)
+/* Arena to fault in; it only fits by reclaiming the page cache. */
+#define ARENA_FAULT	(64 * 1024 * 1024)
+/* Child exit code for "this environment cannot host the test". */
+#define CHILD_UNSUPPORTED	9
+
+static void dump_memcg(void)
+{
+	char buf[512];
+
+	if (!read_cgroup_file(CG_PATH, "memory.max", buf, sizeof(buf)))
+		fprintf(stderr, "memory.max: %s", buf);
+	if (!read_cgroup_file(CG_PATH, "memory.peak", buf, sizeof(buf)))
+		fprintf(stderr, "memory.peak: %s", buf);
+	if (!read_cgroup_file(CG_PATH, "memory.events", buf, sizeof(buf)))
+		fprintf(stderr, "memory.events:\n%s", buf);
+}
+
+/*
+ * Fill the page cache with @size bytes of clean, reclaimable pages by
+ * reading a sparse temp file, the way the cgroup selftests do. Returns the
+ * fd, which must stay open: closing it drops the cache. Returns -EOPNOTSUPP
+ * if the working directory cannot back such a file, -1 on error.
+ */
+static int alloc_pagecache(size_t size)
+{
+	struct statfs stfs;
+	char buf[4096];
+	size_t off;
+	int fd;
+
+	fd = open(".", O_TMPFILE | O_RDWR | O_EXCL, 0600);
+	if (fd < 0)
+		return errno == EOPNOTSUPP ? -EOPNOTSUPP : -1;
+	/* tmpfs hands out shmem pages, which are not reclaimable without swap */
+	if (fstatfs(fd, &stfs) || stfs.f_type == TMPFS_MAGIC) {
+		close(fd);
+		return -EOPNOTSUPP;
+	}
+	if (ftruncate(fd, size))
+		goto err;
+	for (off = 0; off < size; off += sizeof(buf))
+		if (read(fd, buf, sizeof(buf)) < 0)
+			goto err;
+	return fd;
+err:
+	close(fd);
+	return -1;
+}
+
+void serial_test_arena_memcg(void)
+{
+	int cgroup_fd = -1, status, err;
+	const long ps = PAGE_SIZE;
+	char buf[64];
+	pid_t pid;
+
+	err = setup_cgroup_environment();
+	if (!ASSERT_OK(err, "setup_cgroup_environment"))
+		goto out;
+
+	cgroup_fd = create_and_get_cgroup(CG_PATH);
+	if (!ASSERT_OK_FD(cgroup_fd, "create_and_get_cgroup"))
+		goto out;
+
+	/* No memory controller -> nothing to test. */
+	if (read_cgroup_file(CG_PATH, "memory.current", buf, sizeof(buf))) {
+		fprintf(stderr, "%s:SKIP:no memory controller or other env error\n",
+			__func__);
+		test__skip();
+		goto out;
+	}
+
+	pid = fork();
+	if (!ASSERT_GE(pid, 0, "fork"))
+		goto out;
+	if (pid == 0) {
+		struct arena_memcg *cskel;
+		__u32 i, npages;
+		char *base;
+		size_t sz;
+		long cur;
+		int fd;
+
+		/*
+		 * Everything runs in the child: the arena vma is VM_DONTCOPY so
+		 * it does not survive fork(), and only the child should be under
+		 * the limit. The work dir belongs to the parent, so use the
+		 * _parent() helpers; errors come back as an exit code, ASSERT_*
+		 * does not reach the parent from here.
+		 */
+
+		/* Step 1: join the memcg, so what follows is charged to it. */
+		snprintf(buf, sizeof(buf), "%d", getpid());
+		if (write_cgroup_file_parent(CG_PATH, "cgroup.procs", buf))
+			_exit(2);
+
+		/*
+		 * Step 2: load the arena. A map is charged to whoever creates
+		 * it, hence joining first.
+		 */
+		cskel = arena_memcg__open_and_load();
+		if (!cskel)
+			_exit(3);
+		base = bpf_map__initial_value(cskel->maps.arena, &sz);
+		if (!base)
+			_exit(4);
+		npages = ARENA_FAULT / ps;
+		if (npages > bpf_map__max_entries(cskel->maps.arena))
+			_exit(5);
+
+		/* Step 3: make RECLAIMABLE bytes of clean page cache. */
+		fd = alloc_pagecache(RECLAIMABLE);
+		if (fd == -EOPNOTSUPP)
+			_exit(CHILD_UNSUPPORTED);
+		if (fd < 0)
+			_exit(6);
+
+		/*
+		 * Step 4: set memory.max to what we use now plus HEADROOM. The
+		 * page cache is already inside the limit, so only HEADROOM is
+		 * left.
+		 */
+		if (read_cgroup_file_parent(CG_PATH, "memory.current", buf, sizeof(buf)))
+			_exit(7);
+		cur = strtol(buf, NULL, 10);
+		snprintf(buf, sizeof(buf), "%ld", cur + HEADROOM);
+		if (write_cgroup_file_parent(CG_PATH, "memory.max", buf))
+			_exit(8);
+
+		/*
+		 * Step 5: fault ARENA_FAULT of arena in, much more than
+		 * HEADROOM. Once it hits memory.max every further page has to
+		 * come from reclaiming the page cache. With the fix the
+		 * fault-in reclaims and all of it succeeds; without it the
+		 * allocation cannot reclaim and we die on a valid address.
+		 */
+		for (i = 0; i < npages; i++)
+			base[(size_t)i * ps] = 1;
+		_exit(0); /* fd deliberately kept open until here */
+	}
+
+	if (!ASSERT_EQ(waitpid(pid, &status, 0), pid, "waitpid"))
+		goto out;
+
+	/* The working directory cannot hold a reclaimable page cache. */
+	if (WIFEXITED(status) && WEXITSTATUS(status) == CHILD_UNSUPPORTED) {
+		fprintf(stderr, "%s:SKIP:no disk-backed O_TMPFILE in cwd\n", __func__);
+		test__skip();
+		goto out;
+	}
+
+	/* A non-zero exit means the child failed to set up; the code says where. */
+	if (WIFEXITED(status) && WEXITSTATUS(status)) {
+		ASSERT_OK(WEXITSTATUS(status), "child setup");
+		goto out;
+	}
+
+	/*
+	 * With the fix the arena fault-in reclaims the page cache and every
+	 * fault succeeds, so the child exits 0. Without it the allocation
+	 * cannot reclaim, fails once the headroom is used up, and the child
+	 * dies with SIGSEGV on a valid arena address.
+	 */
+	if (!ASSERT_TRUE(WIFEXITED(status) && !WEXITSTATUS(status),
+			 "child faulted the arena in")) {
+		if (WIFSIGNALED(status))
+			fprintf(stderr, "child killed by signal %d\n", WTERMSIG(status));
+		dump_memcg();
+	}
+out:
+	if (cgroup_fd >= 0)
+		close(cgroup_fd);
+	cleanup_cgroup_environment();
+}

tools/testing/selftests/bpf/progs/arena_memcg.c

@@ -0,0 +1,23 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include <vmlinux.h>
+#include <bpf/bpf_helpers.h>
+#include "bpf_arena_common.h"
+
+struct {
+	__uint(type, BPF_MAP_TYPE_ARENA);
+	__uint(map_flags, BPF_F_MMAPABLE);
+	/*
+	 * Number of pages. Must cover ARENA_FAULT on the smallest page size
+	 * (64M/4K = 16384) yet stay under the 4G arena limit on 64K pages
+	 * (50000*64K = 3.2G).
+	 */
+	__uint(max_entries, 50000);
+#ifdef __TARGET_ARCH_arm64
+	__ulong(map_extra, 0x1ull << 32); /* start of mmap() region */
+#else
+	__ulong(map_extra, 0x1ull << 44); /* start of mmap() region */
+#endif
+} arena SEC(".maps");
+
+char _license[] SEC("license") = "GPL";