[RFC PATCH 1/2] memcg: Add support for zram object charge

[Zhongkun He <hezhongkun.hzk@bytedance.com>]

The compressed RAM is currently charged to kernel, not to
any memory cgroup, which is not satisfy our usage scenario.
if the memory of a task is limited by memcgroup, it will
swap out the memory to zram swap device when the memory
is insufficient. In that case, the memory limit will have
no effect. So, it should makes sense to charge the
compressed RAM to the page's memory cgroup.

As we know, zram can be used in two ways, direct and
indirect, this patchset can charge memory in both cases.
Direct zram usage by process within a cgroup will fail
to charge if there is no memory. Indirect zram usage by
process within a cgroup via swap in PF_MEMALLOC context,
will charge successfully.

Signed-off-by: Zhongkun He <hezhongkun.hzk@bytedance.com>
---
 include/linux/memcontrol.h | 10 ++++++++++
 mm/memcontrol.c            | 23 +++++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
index 5818af8eca5a..716cc08a6527 100644
--- a/include/linux/memcontrol.h
+++ b/include/linux/memcontrol.h
@@ -1819,6 +1819,8 @@ static inline void count_objcg_event(struct obj_cgroup *objcg,
 	rcu_read_unlock();
 }
 
+int obj_cgroup_charge_zram(struct obj_cgroup *objcg, size_t size);
+void obj_cgroup_uncharge_zram(struct obj_cgroup *objcg, size_t size);
 #else
 static inline bool mem_cgroup_kmem_disabled(void)
 {
@@ -1880,6 +1882,14 @@ static inline void count_objcg_event(struct obj_cgroup *objcg,
 {
 }
 
+int obj_cgroup_charge_zram(struct obj_cgroup *objcg, size_t size)
+{
+	return 0;
+}
+
+void obj_cgroup_uncharge_zram(struct obj_cgroup *objcg, size_t size)
+{
+}
 #endif /* CONFIG_MEMCG_KMEM */
 
 #if defined(CONFIG_MEMCG_KMEM) && defined(CONFIG_ZSWAP)
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index e8ca4bdcb03c..e1d2c084c132 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -3059,6 +3059,7 @@ struct obj_cgroup *get_obj_cgroup_from_page(struct page *page)
 	}
 	return objcg;
 }
+EXPORT_SYMBOL(get_obj_cgroup_from_page);
 
 static void memcg_account_kmem(struct mem_cgroup *memcg, int nr_pages)
 {
@@ -3409,6 +3410,28 @@ void obj_cgroup_uncharge(struct obj_cgroup *objcg, size_t size)
 	refill_obj_stock(objcg, size, true);
 }
 
+int obj_cgroup_charge_zram(struct obj_cgroup *objcg, size_t size)
+{
+	if (!cgroup_subsys_on_dfl(memory_cgrp_subsys))
+		return 0;
+
+	/*
+	 * Indirect zram usage in PF_MEMALLOC, charging must succeed.
+	 * Direct zram usage, charging  may failed.
+	 */
+	return obj_cgroup_charge(objcg, GFP_KERNEL, size);
+}
+EXPORT_SYMBOL(obj_cgroup_charge_zram);
+
+void obj_cgroup_uncharge_zram(struct obj_cgroup *objcg, size_t size)
+{
+	if (!cgroup_subsys_on_dfl(memory_cgrp_subsys))
+		return;
+
+	obj_cgroup_uncharge(objcg, size);
+}
+EXPORT_SYMBOL(obj_cgroup_uncharge_zram);
+
 #endif /* CONFIG_MEMCG_KMEM */
 
 /*
-- 
2.25.1