[PATCH bpf-next 1/3] bpf: Add memory usage for arena

[Jiayuan Chen <jiayuan.chen@linux.dev>]

arena is the only map type whose map_mem_usage() still returns 0, so
"bpftool map show" and fdinfo always showed 0 memlock for an arena no
matter how many pages it had.

Count the pages that are actually mapped into the arena: bump a counter in
apply_range_set_cb() when a page goes in and drop it in
apply_range_clear_cb() when a page goes out, both under the arena spinlock.
map_mem_usage() then just returns nr_pages << PAGE_SHIFT.

Only real data pages are counted, not the scratch page.

Signed-off-by: Jiayuan Chen <jiayuan.chen@linux.dev>
---
 kernel/bpf/arena.c | 27 ++++++++++++++++-----------
 1 file changed, 16 insertions(+), 11 deletions(-)

diff --git a/kernel/bpf/arena.c b/kernel/bpf/arena.c
index 80b7b8a69446..bc3460211604 100644
--- a/kernel/bpf/arena.c
+++ b/kernel/bpf/arena.c
@@ -55,8 +55,10 @@ struct bpf_arena {
 	struct vm_struct *kern_vm;
 	struct page *scratch_page;
 	struct range_tree rt;
-	/* protects rt */
+	/* protects rt and nr_pages */
 	rqspinlock_t spinlock;
+	/* number of pages currently populated in the arena */
+	u64 nr_pages;
 	struct list_head vma_list;
 	/* protects vma_list */
 	struct mutex lock;
@@ -143,14 +145,14 @@ static long compute_pgoff(struct bpf_arena *arena, long uaddr)
 }
 
 struct apply_range_data {
+	struct bpf_arena *arena;
 	struct page **pages;
-	struct page *scratch_page;
 	int i;
 };
 
 struct clear_range_data {
+	struct bpf_arena *arena;
 	struct llist_head *free_pages;
-	struct page *scratch_page;
 };
 
 static int apply_range_set_cb(pte_t *pte, unsigned long addr, void *data)
@@ -180,7 +182,7 @@ static int apply_range_set_cb(pte_t *pte, unsigned long addr, void *data)
 
 		if (pte_none(old))
 			continue;
-		if (WARN_ON_ONCE(pte_page(old) != d->scratch_page))
+		if (WARN_ON_ONCE(pte_page(old) != d->arena->scratch_page))
 			return -EBUSY;
 		ptep_get_and_clear(&init_mm, addr, pte);
 		flush_tlb_before_set(addr);
@@ -196,6 +198,7 @@ static int apply_range_set_cb(pte_t *pte, unsigned long addr, void *data)
 	set_pte_at(&init_mm, addr, pte, pteval);
 #endif
 	d->i++;
+	d->arena->nr_pages++;
 	return 0;
 }
 
@@ -227,10 +230,11 @@ static int apply_range_clear_cb(pte_t *pte, unsigned long addr, void *data)
 	 * scratches its PTE. A later bpf_arena_free_pages() over that range walks
 	 * here. Without the skip, scratch_page would be freed.
 	 */
-	if (page == d->scratch_page)
+	if (page == d->arena->scratch_page)
 		return 0;
 
 	__llist_add(&page->pcp_llist, d->free_pages);
+	d->arena->nr_pages--;
 	return 0;
 }
 
@@ -413,7 +417,9 @@ static int arena_map_check_btf(struct bpf_map *map, const struct btf *btf,
 
 static u64 arena_map_mem_usage(const struct bpf_map *map)
 {
-	return 0;
+	struct bpf_arena *arena = container_of(map, struct bpf_arena, map);
+
+	return (u64)READ_ONCE(arena->nr_pages) << PAGE_SHIFT;
 }
 
 struct vma_list {
@@ -506,8 +512,7 @@ static vm_fault_t arena_vm_fault(struct vm_fault *vmf)
 	if (ret)
 		goto out_sigsegv_memcg;
 
-	struct apply_range_data data = { .pages = &page, .i = 0,
-					 .scratch_page = arena->scratch_page };
+	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) {
@@ -696,8 +701,8 @@ static long arena_alloc_pages(struct bpf_arena *arena, long uaddr, long page_cnt
 		bpf_map_memcg_exit(old_memcg, new_memcg);
 		return 0;
 	}
+	data.arena = arena;
 	data.pages = pages;
-	data.scratch_page = arena->scratch_page;
 
 	if (raw_res_spin_lock_irqsave(&arena->spinlock, flags))
 		goto out_free_pages;
@@ -873,8 +878,8 @@ static void arena_free_pages(struct bpf_arena *arena, long uaddr, long page_cnt,
 	range_tree_set(&arena->rt, pgoff, page_cnt);
 
 	init_llist_head(&free_pages);
+	cdata.arena = arena;
 	cdata.free_pages = &free_pages;
-	cdata.scratch_page = arena->scratch_page;
 	/* clear ptes and collect struct pages */
 	apply_to_existing_page_range(&init_mm, kaddr, page_cnt << PAGE_SHIFT,
 				     apply_range_clear_cb, &cdata);
@@ -981,8 +986,8 @@ static void arena_free_worker(struct work_struct *work)
 	bpf_map_memcg_enter(&arena->map, &old_memcg, &new_memcg);
 
 	init_llist_head(&free_pages);
+	cdata.arena = arena;
 	cdata.free_pages = &free_pages;
-	cdata.scratch_page = arena->scratch_page;
 	arena_vm_start = bpf_arena_get_kern_vm_start(arena);
 	user_vm_start = bpf_arena_get_user_vm_start(arena);
 
-- 
2.43.0