[PATCH v3 4/7] mm/page_owner: add open-coded page_owner iterator

[Zhen Ni <zhen.ni@easystack.cn>]

Add an open-coded bpf_iter over the pages eligible for page_owner
output, reusing page_owner_next_eligible().

A sleepable BPF program drives it with
bpf_iter_page_owner_new()/next()/destroy(); each next() returns the
page_owner_scan. po_snap is copied inside the page_ext RCU window, so
the program needs no RCU section when it reads it.

Also add bpf_page_owner_stack_snprint(), which renders the stack trace
of the current page into a caller-provided buffer.

Signed-off-by: Zhen Ni <zhen.ni@easystack.cn>
---
Changes in v3:
- Gate the kfuncs with KF_PERFMON.
- Check the page_owner_inited static key in the iterator.
- Without scan->page, start at min_low_pfn - 1 and let next() pre-increment.
- bpf_page_owner_stack_snprint() rejects a NULL or zero-size buffer.

Changes in v2:
- bpf_page_owner_stack_snprint() now takes the scan cursor.
- Buffer args renamed to buf__uninit/buf__sz.

v2: https://lore.kernel.org/linux-mm/20261013122125.843451-4-zhen.ni@easystack.cn/
v1: https://lore.kernel.org/linux-mm/20261009112308.240769-4-zhen.ni@easystack.cn/
---
 kernel/bpf/helpers.c |  9 +++++++
 mm/page_owner.h      |  3 +++
 mm/page_owner_iter.c | 64 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 76 insertions(+)

diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
index 712dca5a2c5b..0ba836f77957 100644
--- a/kernel/bpf/helpers.c
+++ b/kernel/bpf/helpers.c
@@ -4946,6 +4946,15 @@ BTF_ID_FLAGS(func, bpf_iter_dmabuf_new, KF_ITER_NEW | KF_SLEEPABLE)
 BTF_ID_FLAGS(func, bpf_iter_dmabuf_next, KF_ITER_NEXT | KF_RET_NULL | KF_SLEEPABLE)
 BTF_ID_FLAGS(func, bpf_iter_dmabuf_destroy, KF_ITER_DESTROY | KF_SLEEPABLE)
 #endif
+#ifdef CONFIG_PAGE_OWNER
+BTF_ID_FLAGS(func, bpf_iter_page_owner_new,
+	     KF_ITER_NEW | KF_SLEEPABLE | KF_PERFMON)
+BTF_ID_FLAGS(func, bpf_iter_page_owner_next,
+	     KF_ITER_NEXT | KF_RET_NULL | KF_SLEEPABLE | KF_PERFMON)
+BTF_ID_FLAGS(func, bpf_iter_page_owner_destroy,
+	     KF_ITER_DESTROY | KF_SLEEPABLE | KF_PERFMON)
+BTF_ID_FLAGS(func, bpf_page_owner_stack_snprint, KF_SLEEPABLE | KF_PERFMON)
+#endif
 BTF_ID_FLAGS(func, __bpf_trap)
 BTF_ID_FLAGS(func, bpf_strcmp, KF_PERFMON);
 BTF_ID_FLAGS(func, bpf_strcasecmp, KF_PERFMON);
diff --git a/mm/page_owner.h b/mm/page_owner.h
index 0aa7695bec82..4ab513f17fc9 100644
--- a/mm/page_owner.h
+++ b/mm/page_owner.h
@@ -6,6 +6,7 @@
 #ifndef __MM_PAGE_OWNER_H
 #define __MM_PAGE_OWNER_H
 
+#include <linux/jump_label.h>
 #include <linux/page_ext.h>
 #include <linux/sched.h>
 #include <linux/stackdepot.h>
@@ -13,6 +14,8 @@
 /*
  * mm/page_owner.c
  */
+DECLARE_STATIC_KEY_FALSE(page_owner_inited);
+
 struct page_owner {
 	unsigned short order;
 	short last_migrate_reason;
diff --git a/mm/page_owner_iter.c b/mm/page_owner_iter.c
index 529b2c4a589a..4b9741a2de53 100644
--- a/mm/page_owner_iter.c
+++ b/mm/page_owner_iter.c
@@ -4,6 +4,7 @@
 #include <linux/memblock.h>
 #include <linux/mm.h>
 #include <linux/seq_file.h>
+#include <linux/stackdepot.h>
 
 #include "page_owner.h"
 
@@ -116,3 +117,66 @@ int page_owner_iter_register(void)
 
 	return bpf_iter_reg_target(&page_owner_iter_reg_info);
 }
+
+/* opaque to BPF programs */
+struct bpf_iter_page_owner {
+	__u64 __opaque[16];
+} __attribute__((aligned(8)));
+
+struct bpf_iter_page_owner_kern {
+	struct page_owner_scan scan;
+} __attribute__((aligned(8)));
+
+__bpf_kfunc_start_defs();
+
+__bpf_kfunc int bpf_iter_page_owner_new(struct bpf_iter_page_owner *it)
+{
+	struct bpf_iter_page_owner_kern *kit = (void *)it;
+
+	BUILD_BUG_ON(sizeof(*kit) > sizeof(*it));
+	BUILD_BUG_ON(__alignof__(*kit) != __alignof__(*it));
+
+	kit->scan.pfn = min_low_pfn - 1;
+	return 0;
+}
+
+__bpf_kfunc struct page_owner_scan *
+bpf_iter_page_owner_next(struct bpf_iter_page_owner *it)
+{
+	struct bpf_iter_page_owner_kern *kit = (void *)it;
+
+	if (!static_branch_unlikely(&page_owner_inited))
+		return NULL;
+
+	cond_resched();
+
+	kit->scan.pfn++;
+	if (!page_owner_next_eligible(&kit->scan))
+		return NULL;
+
+	return &kit->scan;
+}
+
+__bpf_kfunc void bpf_iter_page_owner_destroy(struct bpf_iter_page_owner *it)
+{
+}
+
+__bpf_kfunc int bpf_page_owner_stack_snprint(struct page_owner_scan *scan,
+					     char *buf__uninit, u32 buf__sz)
+{
+	int len = 0;
+
+	if (!buf__uninit || !buf__sz)
+		return 0;
+
+	if (static_branch_unlikely(&page_owner_inited))
+		len = stack_depot_snprint(scan->po_snap.handle, buf__uninit,
+					  buf__sz, 0);
+
+	if (!len)
+		buf__uninit[0] = '\0';
+
+	return len;
+}
+
+__bpf_kfunc_end_defs();
-- 
2.20.1