[PATCH v4 4/5] mm/zswap: Add per-memcg stat for proactive writeback

[Hao Jia <jiahao.kernel@gmail.com>]

From: Hao Jia <jiahao1@lixiang.com>

Add a new stat zswpwb_proactive_b to memory.stat. This counter is
incremented by entry->length during proactive writebacks triggered
via the zswap_writeback_only key in memory.reclaim. It tracks the
compressed size (in bytes) of pages proactively written back from
zswap to swap, allowing users to better monitor and tune the
proactive writeback mechanism.

Signed-off-by: Hao Jia <jiahao1@lixiang.com>
---
 Documentation/admin-guide/cgroup-v2.rst |  4 ++++
 include/linux/memcontrol.h              |  1 +
 mm/memcontrol.c                         |  3 +++
 mm/zswap.c                              | 23 ++++++++++++++++++-----
 4 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/Documentation/admin-guide/cgroup-v2.rst b/Documentation/admin-guide/cgroup-v2.rst
index e52d97e8e9c6..c164bb415002 100644
--- a/Documentation/admin-guide/cgroup-v2.rst
+++ b/Documentation/admin-guide/cgroup-v2.rst
@@ -1748,6 +1748,10 @@ The following nested keys are defined.
 	  zswpwb
 		Number of pages written from zswap to swap.
 
+	  zswpwb_proactive_b
+		Bytes of compressed data proactively written back from
+		zswap to swap via memory.reclaim zswap_writeback_only key.
+
 	  zswap_incomp
 		Number of incompressible pages currently stored in zswap
 		without compression. These pages could not be compressed to
diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
index e1f46a0016fc..56580b264dc4 100644
--- a/include/linux/memcontrol.h
+++ b/include/linux/memcontrol.h
@@ -40,6 +40,7 @@ enum memcg_stat_item {
 	MEMCG_ZSWAP_B,
 	MEMCG_ZSWAPPED,
 	MEMCG_ZSWAP_INCOMP,
+	MEMCG_ZSWPWB_PROACTIVE_B,
 	MEMCG_NR_STAT,
 };
 
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 56cd4af08232..5ffb5095f0ee 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -433,6 +433,7 @@ static const unsigned int memcg_stat_items[] = {
 	MEMCG_ZSWAP_B,
 	MEMCG_ZSWAPPED,
 	MEMCG_ZSWAP_INCOMP,
+	MEMCG_ZSWPWB_PROACTIVE_B,
 };
 
 #define NR_MEMCG_NODE_STAT_ITEMS ARRAY_SIZE(memcg_node_stat_items)
@@ -1558,6 +1559,7 @@ static const struct memory_stat memory_stats[] = {
 	{ "zswap",			MEMCG_ZSWAP_B			},
 	{ "zswapped",			MEMCG_ZSWAPPED			},
 	{ "zswap_incomp",		MEMCG_ZSWAP_INCOMP		},
+	{ "zswpwb_proactive_b",		MEMCG_ZSWPWB_PROACTIVE_B	},
 #endif
 	{ "file_mapped",		NR_FILE_MAPPED			},
 	{ "file_dirty",			NR_FILE_DIRTY			},
@@ -1614,6 +1616,7 @@ static int memcg_page_state_unit(int item)
 	switch (item) {
 	case MEMCG_PERCPU_B:
 	case MEMCG_ZSWAP_B:
+	case MEMCG_ZSWPWB_PROACTIVE_B:
 	case NR_SLAB_RECLAIMABLE_B:
 	case NR_SLAB_UNRECLAIMABLE_B:
 		return 1;
diff --git a/mm/zswap.c b/mm/zswap.c
index 28200552dde3..d78bacf80209 100644
--- a/mm/zswap.c
+++ b/mm/zswap.c
@@ -163,6 +163,7 @@ struct zswap_pool {
 struct zswap_shrink_walk_arg {
 	unsigned long bytes_written;
 	bool encountered_page_in_swapcache;
+	bool proactive;
 };
 
 /* Global LRU lists shared by all zswap pools. */
@@ -990,7 +991,8 @@ static bool zswap_decompress(struct zswap_entry *entry, struct folio *folio)
  * freed.
  */
 static int zswap_writeback_entry(struct zswap_entry *entry,
-				 swp_entry_t swpentry)
+				 swp_entry_t swpentry,
+				 bool proactive)
 {
 	struct xarray *tree;
 	pgoff_t offset = swp_offset(swpentry);
@@ -1045,6 +1047,15 @@ static int zswap_writeback_entry(struct zswap_entry *entry,
 	if (entry->objcg)
 		count_objcg_events(entry->objcg, ZSWPWB, 1);
 
+	if (proactive && entry->objcg) {
+		struct mem_cgroup *memcg;
+
+		rcu_read_lock();
+		memcg = obj_cgroup_memcg(entry->objcg);
+		mod_memcg_state(memcg, MEMCG_ZSWPWB_PROACTIVE_B, entry->length);
+		rcu_read_unlock();
+	}
+
 	zswap_entry_free(entry);
 
 	/* folio is up to date */
@@ -1155,7 +1166,7 @@ static enum lru_status shrink_memcg_cb(struct list_head *item, struct list_lru_o
 	 */
 	spin_unlock(&l->lock);
 
-	writeback_result = zswap_writeback_entry(entry, swpentry);
+	writeback_result = zswap_writeback_entry(entry, swpentry, walk_arg->proactive);
 
 	if (writeback_result) {
 		zswap_reject_reclaim_fail++;
@@ -1184,6 +1195,7 @@ static unsigned long zswap_shrinker_scan(struct shrinker *shrinker,
 	struct zswap_shrink_walk_arg walk_arg = {
 		.bytes_written = 0,
 		.encountered_page_in_swapcache = false,
+		.proactive = false,
 	};
 	unsigned long shrink_ret;
 
@@ -1305,11 +1317,12 @@ static struct shrinker *zswap_alloc_shrinker(void)
  * writeback disabled, is a zombie cgroup, or has empty zswap LRUs.
  */
 static long shrink_memcg(struct mem_cgroup *memcg,
-			 unsigned long nr_to_writeback)
+			 unsigned long nr_to_writeback, bool proactive)
 {
 	struct zswap_shrink_walk_arg walk_arg = {
 		.bytes_written = 0,
 		.encountered_page_in_swapcache = false,
+		.proactive = proactive,
 	};
 	u64 bytes_to_writeback = nr_to_writeback << PAGE_SHIFT;
 	bool memcg_list_is_empty = true;
@@ -1492,7 +1505,7 @@ static int zswap_try_to_writeback(struct mem_cgroup *memcg,
 		}
 
 		batch_size = min(upper_pages - lower_pages, NR_ZSWAP_WB_BATCH);
-		shrunk = shrink_memcg(iter_memcg, batch_size);
+		shrunk = shrink_memcg(iter_memcg, batch_size, proactive);
 		/* drop the extra reference */
 		mem_cgroup_put(iter_memcg);
 
@@ -1642,7 +1655,7 @@ bool zswap_store(struct folio *folio)
 	objcg = get_obj_cgroup_from_folio(folio);
 	if (objcg && !obj_cgroup_may_zswap(objcg)) {
 		memcg = get_mem_cgroup_from_objcg(objcg);
-		if (shrink_memcg(memcg, 1) <= 0) {
+		if (shrink_memcg(memcg, 1, false) <= 0) {
 			mem_cgroup_put(memcg);
 			goto put_objcg;
 		}
-- 
2.34.1