From fd7745186142f81927769a671ed33187367e8a71 Mon Sep 17 00:00:00 2001
From: Kairui Song <kasong@tencent.com>
Date: Tue, 5 Sep 2023 16:01:16 +0800
Subject: [PATCH] emm: memcg, zram: add support for ZRAM memory accounting

Upstream: alternative

Add a CONFIG_MEMCG_ZRAM for ZRAM driver later. This commit only add
basic structures.

Current plan is that we implement the counting in ZRAM block
level for simplicity of design, may more it to zpool level later for a
unified zram / zswap accounting.

Signed-off-by: Kairui Song <kasong@tencent.com>
---
 include/linux/memcontrol.h |  9 +++++++++
 init/Kconfig               |  5 +++++
 mm/memcontrol.c            | 19 +++++++++++++++++--
 3 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
index 3e8931b3a998..f6bdebc72139 100644
--- a/include/linux/memcontrol.h
+++ b/include/linux/memcontrol.h
@@ -37,6 +37,10 @@ enum memcg_stat_item {
 	MEMCG_KMEM,
 	MEMCG_ZSWAP_B,
 	MEMCG_ZSWAPPED,
+#ifdef CONFIG_MEMCG_ZRAM
+	MEMCG_ZRAM_B,
+	MEMCG_ZRAMED,
+#endif
 	MEMCG_NR_STAT,
 };
 
@@ -231,6 +235,11 @@ struct mem_cgroup {
 	unsigned long zswap_max;
 #endif
 
+#ifdef CONFIG_MEMCG_ZRAM
+	unsigned long zram_max;
+	unsigned short zram_prio;
+#endif
+
 	unsigned long soft_limit;
 
 	/* vmpressure notifications */
diff --git a/init/Kconfig b/init/Kconfig
index b299032e6b2c..1f9d63b22fb8 100644
--- a/init/Kconfig
+++ b/init/Kconfig
@@ -969,6 +969,11 @@ config MEMCG_KMEM_DEFAULT_OFF
 
 	  If unsure, say N.
 
+config MEMCG_ZRAM
+	bool
+	depends on MEMCG && SWAP
+	default y
+
 config BLK_CGROUP
 	bool "IO controller"
 	depends on BLOCK
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 9a7995ed2004..fee8797c04ac 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -87,7 +87,8 @@ EXPORT_PER_CPU_SYMBOL_GPL(int_active_memcg);
 static bool cgroup_memory_nosocket __ro_after_init;
 
 /* Kernel memory accounting disabled? */
-static bool cgroup_memory_nokmem __ro_after_init = IS_ENABLED(CONFIG_MEMCG_KMEM_DEFAULT_OFF);
+bool cgroup_memory_nokmem __ro_after_init = IS_ENABLED(CONFIG_MEMCG_KMEM_DEFAULT_OFF);
+EXPORT_SYMBOL(cgroup_memory_nokmem);
 
 /* BPF memory accounting disabled? */
 static bool cgroup_memory_nobpf __ro_after_init;
@@ -1527,6 +1528,10 @@ static const struct memory_stat memory_stats[] = {
 #if defined(CONFIG_MEMCG_KMEM) && defined(CONFIG_ZSWAP)
 	{ "zswap",			MEMCG_ZSWAP_B			},
 	{ "zswapped",			MEMCG_ZSWAPPED			},
+#endif
+#ifdef CONFIG_MEMCG_ZRAM
+	{ "zram",			MEMCG_ZRAM_B },
+	{ "zrammed",			MEMCG_ZRAMED },
 #endif
 	{ "file_mapped",		NR_FILE_MAPPED			},
 	{ "file_dirty",			NR_FILE_DIRTY			},
@@ -1563,6 +1568,7 @@ static int memcg_page_state_unit(int item)
 	switch (item) {
 	case MEMCG_PERCPU_B:
 	case MEMCG_ZSWAP_B:
+	case MEMCG_ZRAM_B:
 	case NR_SLAB_RECLAIMABLE_B:
 	case NR_SLAB_UNRECLAIMABLE_B:
 	case WORKINGSET_REFAULT_ANON:
@@ -3699,6 +3705,10 @@ static unsigned long mem_cgroup_usage(struct mem_cgroup *memcg, bool swap)
 			global_node_page_state(NR_ANON_MAPPED);
 		if (swap)
 			val += total_swap_pages - get_nr_swap_pages();
+#ifdef CONFIG_MEMCG_ZRAM
+		else
+			val += memcg_page_state(memcg, MEMCG_ZRAM_B) / PAGE_SIZE;
+#endif
 	} else {
 		if (!swap)
 			val = page_counter_read(&memcg->memory);
@@ -5793,12 +5803,17 @@ mem_cgroup_css_alloc(struct cgroup_subsys_state *parent_css)
 	WRITE_ONCE(memcg->soft_limit, PAGE_COUNTER_MAX);
 #if defined(CONFIG_MEMCG_KMEM) && defined(CONFIG_ZSWAP)
 	memcg->zswap_max = PAGE_COUNTER_MAX;
+#endif
+#ifdef CONFIG_MEMCG_ZRAM
+	memcg->zram_max = PAGE_COUNTER_MAX;
 #endif
 	page_counter_set_high(&memcg->swap, PAGE_COUNTER_MAX);
 	if (parent) {
 		WRITE_ONCE(memcg->swappiness, mem_cgroup_swappiness(parent));
 		WRITE_ONCE(memcg->oom_kill_disable, READ_ONCE(parent->oom_kill_disable));
-
+#ifdef CONFIG_MEMCG_ZRAM
+		memcg->zram_prio = parent->zram_prio;
+#endif
 		page_counter_init(&memcg->memory, &parent->memory);
 		page_counter_init(&memcg->swap, &parent->swap);
 		page_counter_init(&memcg->kmem, &parent->kmem);