[PATCH bpf-next v2 2/4] 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>

---
To sashiko:
1. __bpf_alloc_page::can_alloc_pages already check whether we are under lock.
2. apply_to_page_range() does not allocate page tables here. arena_map_alloc() already did it.
---
 kernel/bpf/arena.c | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/kernel/bpf/arena.c b/kernel/bpf/arena.c
index 8dbc24460890..f046e878f7ae 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;
@@ -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++;
+	WRITE_ONCE(d->arena->nr_pages, d->arena->nr_pages + 1);
 	return 0;
 }
 
@@ -231,6 +234,7 @@ static int apply_range_clear_cb(pte_t *pte, unsigned long addr, void *data)
 		return 0;
 
 	__llist_add(&page->pcp_llist, d->free_pages);
+	WRITE_ONCE(d->arena->nr_pages, d->arena->nr_pages - 1);
 	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 {
-- 
2.43.0