[PATCH 03/11] mm/zsmalloc: Introduce conditional memcg awareness to zs_pool

[Joshua Hahn <joshua.hahnjy@gmail.com>]

Introduce 3 new fields to struct zs_pool to allow individual zpools to
be "memcg-aware": memcg_aware, compressed_stat, and uncompressed_stat.

memcg_aware is used in later patches to determine whether memory
should be allocated to keep track of per-compresed object objgs.
compressed_stat and uncompressed_stat are enum indices that point into
memcg (node) stats that zsmalloc will account towards.

In reality, these fields help distinguish between the two users of
zsmalloc, zswap and zram. The enum indices compressed_stat and
uncompressed_stat are parametrized to minimize zswap-specific hardcoding
in zsmalloc.

Suggested-by: Yosry Ahmed <yosry@kernel.org>
Signed-off-by: Joshua Hahn <joshua.hahnjy@gmail.com>
---
 drivers/block/zram/zram_drv.c |  3 ++-
 include/linux/zsmalloc.h      |  5 ++++-
 mm/zsmalloc.c                 | 13 ++++++++++++-
 mm/zswap.c                    |  3 ++-
 4 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/drivers/block/zram/zram_drv.c b/drivers/block/zram/zram_drv.c
index bca33403fc8b..d1eae5c20df7 100644
--- a/drivers/block/zram/zram_drv.c
+++ b/drivers/block/zram/zram_drv.c
@@ -1980,7 +1980,8 @@ static bool zram_meta_alloc(struct zram *zram, u64 disksize)
 	if (!zram->table)
 		return false;
 
-	zram->mem_pool = zs_create_pool(zram->disk->disk_name);
+	/* zram does not support memcg accounting */
+	zram->mem_pool = zs_create_pool(zram->disk->disk_name, false, 0, 0);
 	if (!zram->mem_pool) {
 		vfree(zram->table);
 		zram->table = NULL;
diff --git a/include/linux/zsmalloc.h b/include/linux/zsmalloc.h
index 478410c880b1..24fb2e0fdf67 100644
--- a/include/linux/zsmalloc.h
+++ b/include/linux/zsmalloc.h
@@ -23,8 +23,11 @@ struct zs_pool_stats {
 
 struct zs_pool;
 struct scatterlist;
+enum memcg_stat_item;
 
-struct zs_pool *zs_create_pool(const char *name);
+struct zs_pool *zs_create_pool(const char *name, bool memcg_aware,
+			       enum memcg_stat_item compressed_stat,
+			       enum memcg_stat_item uncompressed_stat);
 void zs_destroy_pool(struct zs_pool *pool);
 
 unsigned long zs_malloc(struct zs_pool *pool, size_t size, gfp_t flags,
diff --git a/mm/zsmalloc.c b/mm/zsmalloc.c
index 7758486e1d06..3f0f42b78314 100644
--- a/mm/zsmalloc.c
+++ b/mm/zsmalloc.c
@@ -214,6 +214,9 @@ struct zs_pool {
 #ifdef CONFIG_COMPACTION
 	struct work_struct free_work;
 #endif
+	bool memcg_aware;
+	enum memcg_stat_item compressed_stat;
+	enum memcg_stat_item uncompressed_stat;
 	/* protect zspage migration/compaction */
 	rwlock_t lock;
 	atomic_t compaction_in_progress;
@@ -2050,6 +2053,9 @@ static int calculate_zspage_chain_size(int class_size)
 /**
  * zs_create_pool - Creates an allocation pool to work from.
  * @name: pool name to be created
+ * @memcg_aware: whether the consumer of this pool will account memcg stats
+ * @compressed_stat: compressed memcontrol stat item to account
+ * @uncompressed_stat: uncompressed memcontrol stat item to account
  *
  * This function must be called before anything when using
  * the zsmalloc allocator.
@@ -2057,7 +2063,9 @@ static int calculate_zspage_chain_size(int class_size)
  * On success, a pointer to the newly created pool is returned,
  * otherwise NULL.
  */
-struct zs_pool *zs_create_pool(const char *name)
+struct zs_pool *zs_create_pool(const char *name, bool memcg_aware,
+			       enum memcg_stat_item compressed_stat,
+			       enum memcg_stat_item uncompressed_stat)
 {
 	int i;
 	struct zs_pool *pool;
@@ -2071,6 +2079,9 @@ struct zs_pool *zs_create_pool(const char *name)
 	rwlock_init(&pool->lock);
 	atomic_set(&pool->compaction_in_progress, 0);
 
+	pool->memcg_aware = memcg_aware;
+	pool->compressed_stat = compressed_stat;
+	pool->uncompressed_stat = uncompressed_stat;
 	pool->name = kstrdup(name, GFP_KERNEL);
 	if (!pool->name)
 		goto err;
diff --git a/mm/zswap.c b/mm/zswap.c
index e6ec3295bdb0..ff9abaa8aa38 100644
--- a/mm/zswap.c
+++ b/mm/zswap.c
@@ -257,7 +257,8 @@ static struct zswap_pool *zswap_pool_create(char *compressor)
 
 	/* unique name for each pool specifically required by zsmalloc */
 	snprintf(name, 38, "zswap%x", atomic_inc_return(&zswap_pools_count));
-	pool->zs_pool = zs_create_pool(name);
+	pool->zs_pool = zs_create_pool(name, true, MEMCG_ZSWAP_B,
+				       MEMCG_ZSWAPPED);
 	if (!pool->zs_pool)
 		goto error;
 
-- 
2.52.0