[RFC PATCH bpf-next 06/10] mm: Add helper to recharge vmalloc'ed address

[Yafang Shao <laoar.shao@xxxxxxxxx>]

This patch introduces a helper to recharge the corresponding pages
of a given vmalloc'ed address. It is similar with how to recharge
a kmalloced'ed address.

Signed-off-by: Yafang Shao <laoar.shao@xxxxxxxxx>
---
 include/linux/slab.h    |  1 +
 include/linux/vmalloc.h |  2 +
 mm/util.c               |  9 +++++
 mm/vmalloc.c            | 87 +++++++++++++++++++++++++++++++++++++++++
 4 files changed, 99 insertions(+)

diff --git a/include/linux/slab.h b/include/linux/slab.h
index 18ab30aa8fe8..e8fb0f6a3660 100644
--- a/include/linux/slab.h
+++ b/include/linux/slab.h
@@ -794,6 +794,7 @@ extern void *kvrealloc(const void *p, size_t oldsize, size_t newsize, gfp_t flag
 		      __alloc_size(3);
 extern void kvfree(const void *addr);
 extern void kvfree_sensitive(const void *addr, size_t len);
+bool kvrecharge(const void *addr, int step);
 
 unsigned int kmem_cache_size(struct kmem_cache *s);
 void __init kmem_cache_init_late(void);
diff --git a/include/linux/vmalloc.h b/include/linux/vmalloc.h
index 096d48aa3437..37c6d0e7b8d5 100644
--- a/include/linux/vmalloc.h
+++ b/include/linux/vmalloc.h
@@ -162,6 +162,8 @@ extern void *vcalloc(size_t n, size_t size) __alloc_size(1, 2);
 
 extern void vfree(const void *addr);
 extern void vfree_atomic(const void *addr);
+bool vrecharge(const void *addr, int step);
+void vuncharge(const void *addr);
 
 extern void *vmap(struct page **pages, unsigned int count,
 			unsigned long flags, pgprot_t prot);
diff --git a/mm/util.c b/mm/util.c
index 0837570c9225..312c05e83132 100644
--- a/mm/util.c
+++ b/mm/util.c
@@ -656,6 +656,15 @@ void kvfree(const void *addr)
 }
 EXPORT_SYMBOL(kvfree);
 
+bool kvrecharge(const void *addr, int step)
+{
+	if (is_vmalloc_addr(addr))
+		return vrecharge(addr, step);
+
+	return krecharge(addr, step);
+}
+EXPORT_SYMBOL(kvrecharge);
+
 /**
  * kvfree_sensitive - Free a data object containing sensitive information.
  * @addr: address of the data object to be freed.
diff --git a/mm/vmalloc.c b/mm/vmalloc.c
index effd1ff6a4b4..7da6e429a45f 100644
--- a/mm/vmalloc.c
+++ b/mm/vmalloc.c
@@ -2745,6 +2745,93 @@ void vfree(const void *addr)
 }
 EXPORT_SYMBOL(vfree);
 
+bool vrecharge(const void *addr, int step)
+{
+	struct obj_cgroup *objcg_new;
+	unsigned int page_order;
+	struct vm_struct *area;
+	struct folio *folio;
+	int i;
+
+	WARN_ON(!in_task());
+
+	if (!addr)
+		return true;
+
+	area = find_vm_area(addr);
+	if (unlikely(!area))
+		return true;
+
+	page_order = vm_area_page_order(area);
+
+	switch (step) {
+	case MEMCG_KMEM_PRE_CHARGE:
+		for (i = 0; i < area->nr_pages; i += 1U << page_order) {
+			struct page *page = area->pages[i];
+
+			WARN_ON(!page);
+			objcg_new = get_obj_cgroup_from_current();
+			WARN_ON(!objcg_new);
+			if (obj_cgroup_charge_pages(objcg_new, GFP_KERNEL,
+						    1 << page_order))
+				goto out_pre;
+			cond_resched();
+		}
+		break;
+	case MEMCG_KMEM_UNCHARGE:
+		for (i = 0; i < area->nr_pages; i += 1U << page_order) {
+			struct page *page = area->pages[i];
+			struct obj_cgroup *objcg_old;
+
+			WARN_ON(!page);
+			folio = page_folio(page);
+			WARN_ON(!folio_memcg_kmem(folio));
+			objcg_old = __folio_objcg(folio);
+
+			obj_cgroup_uncharge_pages(objcg_old, 1 << page_order);
+			/* mod memcg from page */
+			mod_memcg_state(page_memcg(page), MEMCG_VMALLOC,
+					-(1U << page_order));
+			page->memcg_data = 0;
+			obj_cgroup_put(objcg_old);
+			cond_resched();
+		}
+		break;
+	case MEMCG_KMEM_POST_CHARGE:
+		objcg_new = obj_cgroup_from_current();
+		for (i = 0; i < area->nr_pages; i += 1U << page_order) {
+			struct page *page = area->pages[i];
+
+			page->memcg_data = (unsigned long)objcg_new | MEMCG_DATA_KMEM;
+			/* mod memcg from current */
+			mod_memcg_state(page_memcg(page), MEMCG_VMALLOC,
+					1U << page_order);
+
+		}
+		break;
+	case MEMCG_KMEM_CHARGE_ERR:
+		objcg_new = obj_cgroup_from_current();
+		for (i = 0; i < area->nr_pages; i += 1U << page_order) {
+			obj_cgroup_uncharge_pages(objcg_new, 1 << page_order);
+			obj_cgroup_put(objcg_new);
+			cond_resched();
+		}
+		break;
+	}
+
+	return true;
+
+out_pre:
+	for (; i > 0; i -= 1U << page_order) {
+		obj_cgroup_uncharge_pages(objcg_new, 1 << page_order);
+		obj_cgroup_put(objcg_new);
+		cond_resched();
+	}
+
+	return false;
+}
+EXPORT_SYMBOL(vrecharge);
+
 /**
  * vunmap - release virtual mapping obtained by vmap()
  * @addr:   memory base address
-- 
2.17.1