[PATCH 4/8] mm/zsmalloc: Store obj_cgroup pointer in zpdesc

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

With each zswap-backing zpdesc now having an array of obj_cgroup pointers,
plumb the obj_cgroup pointer from the zswap / zram layer down to
zsmalloc.

Introduce two helper functions zpdesc_obj_cgroup and zpdesc_set_obj_cgroup,
which abstract the conversion of an object's zspage idx to its zpdesc
idx and the retrieval of the obj_cgroup pointer from the zpdesc.

From the zswap path, store the obj_cgroup pointer after compression when
writing the object and free when the object gets freed. Also handle the
migration of an object across zpdescs.

The lifetime and charging of the obj_cgroup is still handled in the
zswap layer.

Suggested-by: Johannes Weiner <hannes@cmpxchg.org>
Signed-off-by: Joshua Hahn <joshua.hahnjy@gmail.com>
---
 drivers/block/zram/zram_drv.c |  7 ++--
 include/linux/zsmalloc.h      |  3 +-
 mm/zsmalloc.c                 | 71 ++++++++++++++++++++++++++++++++++-
 mm/zswap.c                    |  6 +--
 4 files changed, 79 insertions(+), 8 deletions(-)

diff --git a/drivers/block/zram/zram_drv.c b/drivers/block/zram/zram_drv.c
index 60ee85679730..209668b14428 100644
--- a/drivers/block/zram/zram_drv.c
+++ b/drivers/block/zram/zram_drv.c
@@ -2231,7 +2231,7 @@ static int write_incompressible_page(struct zram *zram, struct page *page,
 	}
 
 	src = kmap_local_page(page);
-	zs_obj_write(zram->mem_pool, handle, src, PAGE_SIZE);
+	zs_obj_write(zram->mem_pool, handle, src, PAGE_SIZE, NULL);
 	kunmap_local(src);
 
 	slot_lock(zram, index);
@@ -2296,7 +2296,7 @@ static int zram_write_page(struct zram *zram, struct page *page, u32 index)
 		return -ENOMEM;
 	}
 
-	zs_obj_write(zram->mem_pool, handle, zstrm->buffer, comp_len);
+	zs_obj_write(zram->mem_pool, handle, zstrm->buffer, comp_len, NULL);
 	zcomp_stream_put(zstrm);
 
 	slot_lock(zram, index);
@@ -2520,7 +2520,8 @@ static int recompress_slot(struct zram *zram, u32 index, struct page *page,
 		return PTR_ERR((void *)handle_new);
 	}
 
-	zs_obj_write(zram->mem_pool, handle_new, zstrm->buffer, comp_len_new);
+	zs_obj_write(zram->mem_pool, handle_new, zstrm->buffer,
+		     comp_len_new, NULL);
 	zcomp_stream_put(zstrm);
 
 	slot_free(zram, index);
diff --git a/include/linux/zsmalloc.h b/include/linux/zsmalloc.h
index 8ef28b964bb0..22f3baa13f24 100644
--- a/include/linux/zsmalloc.h
+++ b/include/linux/zsmalloc.h
@@ -15,6 +15,7 @@
 #define _ZS_MALLOC_H_
 
 #include <linux/types.h>
+#include <linux/memcontrol.h>
 
 struct zs_pool_stats {
 	/* How many pages were migrated (freed) */
@@ -48,7 +49,7 @@ void zs_obj_read_sg_begin(struct zs_pool *pool, unsigned long handle,
 			  struct scatterlist *sg, size_t mem_len);
 void zs_obj_read_sg_end(struct zs_pool *pool, unsigned long handle);
 void zs_obj_write(struct zs_pool *pool, unsigned long handle,
-		  void *handle_mem, size_t mem_len);
+		  void *handle_mem, size_t mem_len, struct obj_cgroup *objcg);
 
 extern const struct movable_operations zsmalloc_mops;
 
diff --git a/mm/zsmalloc.c b/mm/zsmalloc.c
index 7d56bb700e11..e5ae9a0fc78a 100644
--- a/mm/zsmalloc.c
+++ b/mm/zsmalloc.c
@@ -899,6 +899,41 @@ static void init_zspage(struct size_class *class, struct zspage *zspage)
 }
 
 #ifdef CONFIG_MEMCG
+/* idx is indexed per-zspage, not per-zpdesc. */
+static inline struct obj_cgroup *zpdesc_obj_cgroup(struct zpdesc *zpdesc,
+						   unsigned int idx,
+						   int size)
+{
+	struct obj_cgroup **objcgs = zpdesc_objcgs(zpdesc);
+	unsigned int off = offset_in_page(size * idx);
+	unsigned int zpdesc_idx = DIV_ROUND_UP(off, size);
+
+	if (!objcgs)
+		return NULL;
+
+	return objcgs[zpdesc_idx];
+}
+
+/* idx is indexed per-zspage, not per-zpdesc. */
+static inline void zpdesc_set_obj_cgroup(struct zpdesc *zpdesc,
+					 unsigned int idx, int size,
+					 struct obj_cgroup *objcg)
+{
+	struct obj_cgroup **objcgs = zpdesc_objcgs(zpdesc);
+	unsigned int off = offset_in_page(size * idx);
+	unsigned int zpdesc_idx = DIV_ROUND_UP(off, size);
+
+	if (!objcgs)
+		return;
+
+	objcgs[zpdesc_idx] = objcg;
+	if (off + size > PAGE_SIZE) {
+		/* object spans two pages */
+		objcgs = zpdesc_objcgs(get_next_zpdesc(zpdesc));
+		objcgs[0] = objcg;
+	}
+}
+
 static bool alloc_zspage_objcgs(struct size_class *class, gfp_t gfp,
 				struct zpdesc *zpdescs[])
 {
@@ -927,12 +962,40 @@ static bool alloc_zspage_objcgs(struct size_class *class, gfp_t gfp,
 
 	return true;
 }
+
+static void migrate_obj_objcg(unsigned long used_obj, unsigned long free_obj,
+			      int size)
+{
+	unsigned int s_obj_idx, d_obj_idx;
+	struct zpdesc *s_zpdesc, *d_zpdesc;
+	struct obj_cgroup *objcg;
+
+	obj_to_location(used_obj, &s_zpdesc, &s_obj_idx);
+	obj_to_location(free_obj, &d_zpdesc, &d_obj_idx);
+	objcg = zpdesc_obj_cgroup(s_zpdesc, s_obj_idx, size);
+
+	zpdesc_set_obj_cgroup(d_zpdesc, d_obj_idx, size, objcg);
+	zpdesc_set_obj_cgroup(s_zpdesc, s_obj_idx, size, NULL);
+}
 #else
+static inline struct obj_cgroup *zpdesc_obj_cgroup(struct zpdesc *zpdesc,
+						   unsigned int offset,
+						   int size)
+{
+	return NULL;
+}
+
+static inline void zpdesc_set_obj_cgroup(struct zpdesc *zpdesc,
+					 unsigned int offset, int size,
+					 struct obj_cgroup *objcg) {}
 static bool alloc_zspage_objcgs(struct size_class *class, gfp_t gfp,
 				struct zpdesc *zpdescs[])
 {
 	return true;
 }
+
+static void migrate_obj_objcg(unsigned long used_obj, unsigned long free_obj,
+			      int size) {}
 #endif
 
 static void create_page_chain(struct size_class *class, struct zspage *zspage,
@@ -1221,7 +1284,7 @@ void zs_obj_read_sg_end(struct zs_pool *pool, unsigned long handle)
 EXPORT_SYMBOL_GPL(zs_obj_read_sg_end);
 
 void zs_obj_write(struct zs_pool *pool, unsigned long handle,
-		  void *handle_mem, size_t mem_len)
+		  void *handle_mem, size_t mem_len, struct obj_cgroup *objcg)
 {
 	struct zspage *zspage;
 	struct zpdesc *zpdesc;
@@ -1242,6 +1305,9 @@ void zs_obj_write(struct zs_pool *pool, unsigned long handle,
 	class = zspage_class(pool, zspage);
 	off = offset_in_page(class->size * obj_idx);
 
+	if (objcg)
+		zpdesc_set_obj_cgroup(zpdesc, obj_idx, class->size, objcg);
+
 	if (!ZsHugePage(zspage))
 		off += ZS_HANDLE_SIZE;
 
@@ -1415,6 +1481,8 @@ static void obj_free(int class_size, unsigned long obj)
 	f_offset = offset_in_page(class_size * f_objidx);
 	zspage = get_zspage(f_zpdesc);
 
+	zpdesc_set_obj_cgroup(f_zpdesc, f_objidx, class_size, NULL);
+
 	vaddr = kmap_local_zpdesc(f_zpdesc);
 	link = (struct link_free *)(vaddr + f_offset);
 
@@ -1587,6 +1655,7 @@ static void migrate_zspage(struct zs_pool *pool, struct zspage *src_zspage,
 		used_obj = handle_to_obj(handle);
 		free_obj = obj_malloc(pool, dst_zspage, handle);
 		zs_obj_copy(class, free_obj, used_obj);
+		migrate_obj_objcg(used_obj, free_obj, class->size);
 		obj_idx++;
 		obj_free(class->size, used_obj);
 
diff --git a/mm/zswap.c b/mm/zswap.c
index dd083110bfa0..1e2d60f47919 100644
--- a/mm/zswap.c
+++ b/mm/zswap.c
@@ -851,7 +851,7 @@ static void acomp_ctx_put_unlock(struct crypto_acomp_ctx *acomp_ctx)
 }
 
 static bool zswap_compress(struct page *page, struct zswap_entry *entry,
-			   struct zswap_pool *pool)
+			   struct zswap_pool *pool, struct obj_cgroup *objcg)
 {
 	struct crypto_acomp_ctx *acomp_ctx;
 	struct scatterlist input, output;
@@ -911,7 +911,7 @@ static bool zswap_compress(struct page *page, struct zswap_entry *entry,
 		goto unlock;
 	}
 
-	zs_obj_write(pool->zs_pool, handle, dst, dlen);
+	zs_obj_write(pool->zs_pool, handle, dst, dlen, objcg);
 	entry->handle = handle;
 	entry->length = dlen;
 
@@ -1413,7 +1413,7 @@ static bool zswap_store_page(struct page *page,
 		return false;
 	}
 
-	if (!zswap_compress(page, entry, pool))
+	if (!zswap_compress(page, entry, pool, objcg))
 		goto compress_failed;
 
 	old = xa_store(swap_zswap_tree(page_swpentry),
-- 
2.47.3