[PATCH 04/11] memcg: add cma charge/uncharge functions for area counters

[Eric Chanudet]

Introduce mem_cgroup_charge_cma() and mem_cgroup_uncharge_cma() to
account CMA allocations under memcg when memory_cma_accounting is
enabled in cgroupfs.

memcg counters to account for allocation in each CMA area are kcalloc'd
in mem_cgroup_alloc() and added to the page-counter hierarchy in
mem_cgroup_css_alloc(). cma_area_count is used to size the array. It is
calculated during setup_arch(), before mem_cgroup_alloc() is called
(for the first time in cgroup_init()).

On cma_alloc(), pages are charged both to the cgroup's memory controller
(with try_charge_memcg()) and to the per area page counter for the CMA
region used.

CMA ranges can be either pages with a 0-order folio (cma_alloc_frozen,
!__GFP_COMP) or pages under a large folio (cma_alloc_frozen_compound,
__GFP_COMP). In the first case, mark each page's folio with the cgroup's
obj_cgroup. In the second case, only mark the large folio from the first
page.

On cma_release(), get the first page's folio and uncharge both counters
accordingly if it is a large folio or not.

A small cma_area_index() helper is introduced to succinctly get the
counter of a given CMA region in the memcg CMA area counters array.

Signed-off-by: Eric Chanudet <echanude@xxxxxxxxxx>
---
 include/linux/memcontrol.h |  19 ++++++++
 mm/cma.h                   |   5 ++
 mm/memcontrol.c            | 116 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 140 insertions(+)

diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
index 8170bb8066a2..801fbde94137 100644
--- a/include/linux/memcontrol.h
+++ b/include/linux/memcontrol.h
@@ -320,6 +320,10 @@ struct mem_cgroup {
 	spinlock_t event_list_lock;
 #endif /* CONFIG_MEMCG_V1 */
 
+#ifdef CONFIG_CMA
+	struct page_counter *cma_counters;
+#endif /* CONFIG_CMA */
+
 	struct mem_cgroup_per_node *nodeinfo[];
 };
 
@@ -662,6 +666,11 @@ int mem_cgroup_charge_hugetlb(struct folio* folio, gfp_t gfp);
 int mem_cgroup_swapin_charge_folio(struct folio *folio, unsigned short id,
 				   struct mm_struct *mm, gfp_t gfp);
 
+int mem_cgroup_charge_cma(struct page *page, unsigned long count,
+			  struct cma *cma);
+void mem_cgroup_uncharge_cma(struct page *page, unsigned long count,
+			     struct cma *cma);
+
 void __mem_cgroup_uncharge(struct folio *folio);
 
 /**
@@ -1160,6 +1169,16 @@ static inline int mem_cgroup_swapin_charge_folio(struct folio *folio,
 	return 0;
 }
 
+static inline int mem_cgroup_charge_cma(struct page *page, unsigned long count,
+					struct cma *cma)
+{
+	return 0;
+}
+static inline void mem_cgroup_uncharge_cma(struct page *page,
+					   unsigned long count, struct cma *cma)
+{
+}
+
 static inline void mem_cgroup_uncharge(struct folio *folio)
 {
 }
diff --git a/mm/cma.h b/mm/cma.h
index 0c9de38d6bff..52b7d4be5170 100644
--- a/mm/cma.h
+++ b/mm/cma.h
@@ -81,6 +81,11 @@ static inline unsigned long cma_bitmap_maxno(struct cma *cma,
 	return cmr->count >> cma->order_per_bit;
 }
 
+static inline int cma_area_index(const struct cma *cma)
+{
+	return cma - cma_areas;
+}
+
 #ifdef CONFIG_CMA_SYSFS
 void cma_sysfs_account_success_pages(struct cma *cma, unsigned long nr_pages);
 void cma_sysfs_account_fail_pages(struct cma *cma, unsigned long nr_pages);
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 6dc4888a90f3..8eca00d30fc5 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -68,6 +68,7 @@
 #include <net/ip.h>
 #include "slab.h"
 #include "memcontrol-v1.h"
+#include "cma.h"
 
 #include <linux/uaccess.h>
 
@@ -4062,6 +4063,9 @@ static void __mem_cgroup_free(struct mem_cgroup *memcg)
 {
 	int node;
 
+#ifdef CONFIG_CMA
+	kfree(memcg->cma_counters);
+#endif
 	for_each_node(node) {
 		struct mem_cgroup_per_node *pn = memcg->nodeinfo[node];
 		if (!pn)
@@ -4144,6 +4148,15 @@ static struct mem_cgroup *mem_cgroup_alloc(struct mem_cgroup *parent)
 	for (i = 0; i < MEMCG_CGWB_FRN_CNT; i++)
 		memcg->cgwb_frn[i].done =
 			__WB_COMPLETION_INIT(&memcg_cgwb_frn_waitq);
+#endif
+#ifdef CONFIG_CMA
+	if (cma_area_count) {
+		memcg->cma_counters = kcalloc(cma_area_count,
+					      sizeof(struct page_counter),
+					      GFP_KERNEL);
+		if (!memcg->cma_counters)
+			goto fail;
+	}
 #endif
 	lru_gen_init_memcg(memcg);
 	return memcg;
@@ -4159,6 +4172,7 @@ mem_cgroup_css_alloc(struct cgroup_subsys_state *parent_css)
 	struct mem_cgroup *parent = mem_cgroup_from_css(parent_css);
 	struct mem_cgroup *memcg, *old_memcg;
 	bool memcg_on_dfl = cgroup_subsys_on_dfl(memory_cgrp_subsys);
+	unsigned int __maybe_unused i;
 
 	old_memcg = set_active_memcg(parent);
 	memcg = mem_cgroup_alloc(parent);
@@ -4171,6 +4185,12 @@ mem_cgroup_css_alloc(struct cgroup_subsys_state *parent_css)
 #ifdef CONFIG_ZSWAP
 	memcg->zswap_max = PAGE_COUNTER_MAX;
 	WRITE_ONCE(memcg->zswap_writeback, true);
+#endif
+#ifdef CONFIG_CMA
+	for (i = 0; i < cma_area_count; ++i)
+		page_counter_init(&memcg->cma_counters[i],
+				  parent ? &parent->cma_counters[i] : NULL,
+				  false);
 #endif
 	page_counter_set_high(&memcg->swap, PAGE_COUNTER_MAX);
 	if (parent) {
@@ -5213,6 +5233,102 @@ int mem_cgroup_swapin_charge_folio(struct folio *folio, unsigned short id,
 	return ret;
 }
 
+#ifdef CONFIG_CMA
+static bool memcg_accounts_cma(void)
+{
+	return cgrp_dfl_root.flags & CGRP_ROOT_MEMORY_CMA_ACCOUNTING;
+}
+
+int mem_cgroup_charge_cma(struct page *page, unsigned long count,
+			  struct cma *cma)
+{
+	struct page_counter *counter, *fail;
+	struct obj_cgroup *objcg;
+	struct mem_cgroup *memcg;
+	unsigned int i;
+	int rc = 0;
+
+	if (mem_cgroup_disabled() || !memcg_accounts_cma())
+		return 0;
+
+	memcg = get_mem_cgroup_from_current();
+	if (!memcg)
+		return 0;
+
+	if (mem_cgroup_is_root(memcg))
+		goto cgroup_put;
+
+	rc = try_charge_memcg(memcg, GFP_KERNEL, count);
+	if (rc)
+		goto cgroup_put;
+
+	counter = &memcg->cma_counters[cma_area_index(cma)];
+	if (!page_counter_try_charge(counter, count, &fail)) {
+		refill_stock(memcg, count);
+		rc = -ENOMEM;
+		goto cgroup_put;
+	}
+
+	objcg = get_obj_cgroup_from_memcg(memcg);
+	if (folio_test_large(page_folio(page))) {
+		commit_charge(page_folio(page), objcg);
+	} else {
+		obj_cgroup_get_many(objcg, count - 1);
+		for (i = 0; i < count; i++)
+			commit_charge(page_folio(page + i), objcg);
+	}
+
+cgroup_put:
+	mem_cgroup_put(memcg);
+	return rc;
+}
+
+void mem_cgroup_uncharge_cma(struct page *page, unsigned long count,
+			     struct cma *cma)
+{
+	struct page_counter *counter;
+	struct obj_cgroup *objcg;
+	struct mem_cgroup *memcg;
+	struct folio *folio;
+	unsigned int i;
+
+	if (mem_cgroup_disabled() || !memcg_accounts_cma())
+		return;
+
+	/*
+	 * Get the objcg from the first page.
+	 * page_objcg() expects MEMCG_DATA_KMEM, but for CMA we used
+	 * commit_charge() which sets folio->memcg_data = objcg
+	 * without flags, so we cannot use it.
+	 */
+	objcg = folio_objcg(page_folio(page));
+	if (!objcg)
+		return;
+
+	rcu_read_lock();
+	memcg = obj_cgroup_memcg(objcg);
+
+	counter = &memcg->cma_counters[cma_area_index(cma)];
+	page_counter_uncharge(counter, count);
+
+	memcg_uncharge(memcg, count);
+
+	rcu_read_unlock();
+
+	folio = page_folio(page);
+	if (folio_test_large(folio)) {
+		folio->memcg_data = 0;
+		obj_cgroup_put(objcg);
+	} else {
+		for (i = 0; i < count; ++i) {
+			folio = page_folio(page + i);
+			folio->memcg_data = 0;
+			obj_cgroup_put(objcg);
+		}
+	}
+}
+#endif /* CONFIG_CMA */
+
 struct uncharge_gather {
 	struct obj_cgroup *objcg;
 	unsigned long nr_memory;

-- 
2.53.0