[PATCH 4/6] mm/page_owner: add bpf_page_owner_get_nid() kfunc

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

Extract the NUMA node lookup from read_page_owner() into a shared
page_owner_get_nid() helper, and expose it to BPF programs as
bpf_page_owner_get_nid().

The nid bitfield layout is build-time configuration, so BPF programs
cannot accurately derive it from page->flags.

Signed-off-by: Zhen Ni <zhen.ni@easystack.cn>
---
 kernel/bpf/helpers.c |  1 +
 mm/page_owner.c      | 26 +++++++++++++++-----------
 mm/page_owner.h      |  1 +
 mm/page_owner_iter.c |  5 +++++
 4 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
index 23a2d2d71f4d..3e8afd19a1f4 100644
--- a/kernel/bpf/helpers.c
+++ b/kernel/bpf/helpers.c
@@ -4951,6 +4951,7 @@ BTF_ID_FLAGS(func, bpf_iter_page_owner_new, KF_ITER_NEW | KF_SLEEPABLE)
 BTF_ID_FLAGS(func, bpf_iter_page_owner_next, KF_ITER_NEXT | KF_RET_NULL | KF_SLEEPABLE)
 BTF_ID_FLAGS(func, bpf_iter_page_owner_destroy, KF_ITER_DESTROY | KF_SLEEPABLE)
 BTF_ID_FLAGS(func, bpf_page_owner_stack_snprint, KF_SLEEPABLE)
+BTF_ID_FLAGS(func, bpf_page_owner_get_nid, KF_SLEEPABLE)
 #endif
 BTF_ID_FLAGS(func, __bpf_trap)
 BTF_ID_FLAGS(func, bpf_strcmp, KF_PERFMON);
diff --git a/mm/page_owner.c b/mm/page_owner.c
index eb420803a434..43a722678acc 100644
--- a/mm/page_owner.c
+++ b/mm/page_owner.c
@@ -711,6 +711,19 @@ void __dump_page_owner(const struct page *page)
 	page_ext_put(page_ext);
 }
 
+/*
+ * NUMA node of @page, bypassing PF_POISONED_CHECK() in page_to_nid() to
+ * avoid VM_BUG_ON on poisoned pages; -1 for poisoned pages.
+ */
+int page_owner_get_nid(struct page *page)
+{
+	memdesc_flags_t page_flags = READ_ONCE(page->flags);
+
+	if (page_flags.f == PAGE_POISON_PATTERN)
+		return -1;
+	return memdesc_nid(&page_flags);
+}
+
 /*
  * Advance @scan to the next page eligible for page_owner output.
  * On a hit, fill scan->pfn/page/po_snap and return true; the caller
@@ -810,18 +823,9 @@ read_page_owner(struct file *file, char __user *buf, size_t count, loff_t *ppos)
 	/* Find an allocated page */
 	while (page_owner_next_eligible(&scan)) {
 		if (state->nid_filter_enabled) {
-			int nid;
-			memdesc_flags_t page_flags =
-				READ_ONCE(scan.page->flags);
+			int nid = page_owner_get_nid(scan.page);
 
-			/*
-			 * Bypass PF_POISONED_CHECK() in page_to_nid() to avoid
-			 * VM_BUG_ON when accessing poisoned pages.
-			 */
-			if (page_flags.f == PAGE_POISON_PATTERN)
-				goto skip_continue;
-			nid = memdesc_nid(&page_flags);
-			if (!node_isset(nid, state->nid_filter))
+			if (nid < 0 || !node_isset(nid, state->nid_filter))
 				goto skip_continue;
 		}
 
diff --git a/mm/page_owner.h b/mm/page_owner.h
index fda83be69275..c57859d18dc3 100644
--- a/mm/page_owner.h
+++ b/mm/page_owner.h
@@ -39,6 +39,7 @@ struct page_owner_scan {
 };
 
 bool page_owner_next_eligible(struct page_owner_scan *scan);
+int page_owner_get_nid(struct page *page);
 
 /*
  * mm/page_owner_iter.c
diff --git a/mm/page_owner_iter.c b/mm/page_owner_iter.c
index 80693d417407..907517a29144 100644
--- a/mm/page_owner_iter.c
+++ b/mm/page_owner_iter.c
@@ -165,4 +165,9 @@ __bpf_kfunc int bpf_page_owner_stack_snprint(depot_stack_handle_t handle,
 	return stack_depot_snprint(handle, buf, buf_size, 0);
 }
 
+__bpf_kfunc int bpf_page_owner_get_nid(struct page_owner_scan *scan)
+{
+	return page_owner_get_nid(scan->page);
+}
+
 __bpf_kfunc_end_defs();
-- 
2.20.1