Re: [RFC PATCH] zram: avoid preemption with CPU-based compression backends

[Sergey Senozhatsky <senozhatsky@chromium.org>]

Hi Barry,

On (26/08/05 15:50), Barry Song wrote:
> BTW, I wonder if compression and decompression could use separate
> mutexes. That way, a sleepable zs_malloc() in the compression path
> would not block decompression, which is the more latency-sensitive
> operation.

quick and dirty patch.  Just curious if this improves anything on your
side.

We also maybe can have more that num_online_cpus() stream, if we
switch to idle streams list instead [1]

[1] https://lore.kernel.org/lkml/20250130111105.2861324-3-senozhatsky@chromium.org/

----

diff --git a/drivers/block/zram/zcomp.c b/drivers/block/zram/zcomp.c
index 974c4691887e..3c523ea0dc27 100644
--- a/drivers/block/zram/zcomp.c
+++ b/drivers/block/zram/zcomp.c
@@ -112,21 +112,28 @@ ssize_t zcomp_available_show(const char *comp, char *buf, ssize_t at)
 	return at;
 }
 
-struct zcomp_strm *zcomp_stream_get(struct zcomp *comp)
+struct zcomp_strm *zcomp_stream_get_write(struct zcomp *comp)
 {
 	for (;;) {
-		struct zcomp_strm *zstrm = raw_cpu_ptr(comp->stream);
-
-		/*
-		 * Inspired by zswap
-		 *
-		 * stream is returned with ->mutex locked which prevents
-		 * cpu_dead() from releasing this stream under us, however
-		 * there is still a race window between raw_cpu_ptr() and
-		 * mutex_lock(), during which we could have been migrated
-		 * from a CPU that has already destroyed its stream.  If
-		 * so then unlock and re-try on the current CPU.
-		 */
+		struct zcomp_strm *zstrm = raw_cpu_ptr(comp->stream_write);
+
+		mutex_lock(&zstrm->lock);
+		if (likely(zstrm->buffer))
+			return zstrm;
+		mutex_unlock(&zstrm->lock);
+	}
+}
+
+void zcomp_stream_put_write(struct zcomp_strm *zstrm)
+{
+	mutex_unlock(&zstrm->lock);
+}
+
+struct zcomp_strm *zcomp_stream_get_read(struct zcomp *comp)
+{
+	for (;;) {
+		struct zcomp_strm *zstrm = raw_cpu_ptr(comp->stream_read);
+
 		mutex_lock(&zstrm->lock);
 		if (likely(zstrm->buffer))
 			return zstrm;
@@ -134,7 +141,7 @@ struct zcomp_strm *zcomp_stream_get(struct zcomp *comp)
 	}
 }
 
-void zcomp_stream_put(struct zcomp_strm *zstrm)
+void zcomp_stream_put_read(struct zcomp_strm *zstrm)
 {
 	mutex_unlock(&zstrm->lock);
 }
@@ -174,23 +181,39 @@ int zcomp_decompress(struct zcomp *comp, struct zcomp_strm *zstrm,
 int zcomp_cpu_up_prepare(unsigned int cpu, struct hlist_node *node)
 {
 	struct zcomp *comp = hlist_entry(node, struct zcomp, node);
-	struct zcomp_strm *zstrm = per_cpu_ptr(comp->stream, cpu);
+	struct zcomp_strm *zstrm_w = per_cpu_ptr(comp->stream_write, cpu);
+	struct zcomp_strm *zstrm_r = per_cpu_ptr(comp->stream_read, cpu);
 	int ret;
 
-	ret = zcomp_strm_init(comp, zstrm);
-	if (ret)
-		pr_err("Can't allocate a compression stream\n");
-	return ret;
+	ret = zcomp_strm_init(comp, zstrm_w);
+	if (ret) {
+		pr_err("Can't allocate a compression write stream\n");
+		return ret;
+	}
+
+	ret = zcomp_strm_init(comp, zstrm_r);
+	if (ret) {
+		pr_err("Can't allocate a compression read stream\n");
+		zcomp_strm_free(comp, zstrm_w);
+		return ret;
+	}
+
+	return 0;
 }
 
 int zcomp_cpu_dead(unsigned int cpu, struct hlist_node *node)
 {
 	struct zcomp *comp = hlist_entry(node, struct zcomp, node);
-	struct zcomp_strm *zstrm = per_cpu_ptr(comp->stream, cpu);
+	struct zcomp_strm *zstrm_w = per_cpu_ptr(comp->stream_write, cpu);
+	struct zcomp_strm *zstrm_r = per_cpu_ptr(comp->stream_read, cpu);
 
-	mutex_lock(&zstrm->lock);
-	zcomp_strm_free(comp, zstrm);
-	mutex_unlock(&zstrm->lock);
+	mutex_lock(&zstrm_w->lock);
+	zcomp_strm_free(comp, zstrm_w);
+	mutex_unlock(&zstrm_w->lock);
+
+	mutex_lock(&zstrm_r->lock);
+	zcomp_strm_free(comp, zstrm_r);
+	mutex_unlock(&zstrm_r->lock);
 	return 0;
 }
 
@@ -198,17 +221,25 @@ static int zcomp_init(struct zcomp *comp, struct zcomp_params *params)
 {
 	int ret, cpu;
 
-	comp->stream = alloc_percpu(struct zcomp_strm);
-	if (!comp->stream)
+	comp->stream_write = alloc_percpu(struct zcomp_strm);
+	if (!comp->stream_write)
 		return -ENOMEM;
 
+	comp->stream_read = alloc_percpu(struct zcomp_strm);
+	if (!comp->stream_read) {
+		free_percpu(comp->stream_write);
+		return -ENOMEM;
+	}
+
 	comp->params = params;
 	ret = comp->ops->setup_params(comp->params);
 	if (ret)
 		goto cleanup;
 
-	for_each_possible_cpu(cpu)
-		mutex_init(&per_cpu_ptr(comp->stream, cpu)->lock);
+	for_each_possible_cpu(cpu) {
+		mutex_init(&per_cpu_ptr(comp->stream_write, cpu)->lock);
+		mutex_init(&per_cpu_ptr(comp->stream_read, cpu)->lock);
+	}
 
 	ret = cpuhp_state_add_instance(CPUHP_ZCOMP_PREPARE, &comp->node);
 	if (ret < 0)
@@ -218,7 +249,8 @@ static int zcomp_init(struct zcomp *comp, struct zcomp_params *params)
 
 cleanup:
 	comp->ops->release_params(comp->params);
-	free_percpu(comp->stream);
+	free_percpu(comp->stream_read);
+	free_percpu(comp->stream_write);
 	return ret;
 }
 
@@ -226,7 +258,8 @@ void zcomp_destroy(struct zcomp *comp)
 {
 	cpuhp_state_remove_instance(CPUHP_ZCOMP_PREPARE, &comp->node);
 	comp->ops->release_params(comp->params);
-	free_percpu(comp->stream);
+	free_percpu(comp->stream_read);
+	free_percpu(comp->stream_write);
 	kfree(comp);
 }
 
diff --git a/drivers/block/zram/zcomp.h b/drivers/block/zram/zcomp.h
index 81a0f3f6ff48..fd919571d8b7 100644
--- a/drivers/block/zram/zcomp.h
+++ b/drivers/block/zram/zcomp.h
@@ -71,7 +71,8 @@ struct zcomp_ops {
 
 /* dynamic per-device compression frontend */
 struct zcomp {
-	struct zcomp_strm __percpu *stream;
+	struct zcomp_strm __percpu *stream_write;
+	struct zcomp_strm __percpu *stream_read;
 	const struct zcomp_ops *ops;
 	struct zcomp_params *params;
 	struct hlist_node node;
@@ -85,8 +86,11 @@ const char *zcomp_lookup_backend_name(const char *comp);
 struct zcomp *zcomp_create(const char *alg, struct zcomp_params *params);
 void zcomp_destroy(struct zcomp *comp);
 
-struct zcomp_strm *zcomp_stream_get(struct zcomp *comp);
-void zcomp_stream_put(struct zcomp_strm *zstrm);
+struct zcomp_strm *zcomp_stream_get_write(struct zcomp *comp);
+void zcomp_stream_put_write(struct zcomp_strm *zstrm);
+
+struct zcomp_strm *zcomp_stream_get_read(struct zcomp *comp);
+void zcomp_stream_put_read(struct zcomp_strm *zstrm);
 
 int zcomp_compress(struct zcomp *comp, struct zcomp_strm *zstrm,
 		   const void *src, unsigned int *dst_len);
diff --git a/drivers/block/zram/zram_drv.c b/drivers/block/zram/zram_drv.c
index 3b9dfcae9317..6a564a5e0041 100644
--- a/drivers/block/zram/zram_drv.c
+++ b/drivers/block/zram/zram_drv.c
@@ -1358,14 +1358,14 @@ static int decompress_bdev_page(struct zram *zram, struct page *page, u32 index)
 	size = get_slot_size(zram, index);
 	prio = get_slot_comp_priority(zram, index);
 
-	zstrm = zcomp_stream_get(zram->comps[prio]);
+	zstrm = zcomp_stream_get_read(zram->comps[prio]);
 	src = kmap_local_page(page);
 	ret = zcomp_decompress(zram->comps[prio], zstrm, src, size,
 			       zstrm->local_copy);
 	if (!ret)
 		copy_page(src, zstrm->local_copy);
 	kunmap_local(src);
-	zcomp_stream_put(zstrm);
+	zcomp_stream_put_read(zstrm);
 	slot_unlock(zram, index);
 
 	return ret;
@@ -2101,14 +2101,14 @@ static int read_compressed_page(struct zram *zram, struct page *page, u32 index)
 	size = get_slot_size(zram, index);
 	prio = get_slot_comp_priority(zram, index);
 
-	zstrm = zcomp_stream_get(zram->comps[prio]);
+	zstrm = zcomp_stream_get_read(zram->comps[prio]);
 	src = zs_obj_read_begin(zram->mem_pool, handle, size,
 				zstrm->local_copy);
 	dst = kmap_local_page(page);
 	ret = zcomp_decompress(zram->comps[prio], zstrm, src, size, dst);
 	kunmap_local(dst);
 	zs_obj_read_end(zram->mem_pool, handle, size, src);
-	zcomp_stream_put(zstrm);
+	zcomp_stream_put_read(zstrm);
 
 	return ret;
 }
@@ -2129,12 +2129,12 @@ static int read_from_zspool_raw(struct zram *zram, struct page *page, u32 index)
 	 * case if object spans two physical pages. No decompression
 	 * takes place here, as we read raw compressed data.
 	 */
-	zstrm = zcomp_stream_get(zram->comps[ZRAM_PRIMARY_COMP]);
+	zstrm = zcomp_stream_get_read(zram->comps[ZRAM_PRIMARY_COMP]);
 	src = zs_obj_read_begin(zram->mem_pool, handle, size,
 				zstrm->local_copy);
 	memcpy_to_page(page, 0, src, size);
 	zs_obj_read_end(zram->mem_pool, handle, size, src);
-	zcomp_stream_put(zstrm);
+	zcomp_stream_put_read(zstrm);
 
 	memzero_page(page, size, PAGE_SIZE - size);
 
@@ -2285,20 +2285,20 @@ static int zram_write_page(struct zram *zram, struct page *page, u32 index)
 	if (same_filled)
 		return write_same_filled_page(zram, element, index);
 
-	zstrm = zcomp_stream_get(zram->comps[ZRAM_PRIMARY_COMP]);
+	zstrm = zcomp_stream_get_write(zram->comps[ZRAM_PRIMARY_COMP]);
 	mem = kmap_local_page(page);
 	ret = zcomp_compress(zram->comps[ZRAM_PRIMARY_COMP], zstrm,
 			     mem, &comp_len);
 	kunmap_local(mem);
 
 	if (unlikely(ret)) {
-		zcomp_stream_put(zstrm);
+		zcomp_stream_put_write(zstrm);
 		pr_err("Compression failed! err=%d\n", ret);
 		return ret;
 	}
 
 	if (comp_len >= huge_class_size) {
-		zcomp_stream_put(zstrm);
+		zcomp_stream_put_write(zstrm);
 		return write_incompressible_page(zram, page, index);
 	}
 
@@ -2306,18 +2306,18 @@ static int zram_write_page(struct zram *zram, struct page *page, u32 index)
 			   GFP_NOIO | __GFP_NOWARN |
 			   __GFP_HIGHMEM | __GFP_MOVABLE, page_to_nid(page));
 	if (IS_ERR_VALUE(handle)) {
-		zcomp_stream_put(zstrm);
+		zcomp_stream_put_write(zstrm);
 		return PTR_ERR((void *)handle);
 	}
 
 	if (!zram_can_store_page(zram)) {
-		zcomp_stream_put(zstrm);
+		zcomp_stream_put_write(zstrm);
 		zs_free(zram->mem_pool, handle);
 		return -ENOMEM;
 	}
 
 	zs_obj_write(zram->mem_pool, handle, zstrm->buffer, comp_len);
-	zcomp_stream_put(zstrm);
+	zcomp_stream_put_write(zstrm);
 
 	slot_lock(zram, index);
 	slot_free(zram, index);
@@ -2457,7 +2457,7 @@ static int recompress_slot(struct zram *zram, u32 index, struct page *page,
 	 */
 	clear_slot_flag(zram, index, ZRAM_IDLE);
 
-	zstrm = zcomp_stream_get(zram->comps[prio]);
+	zstrm = zcomp_stream_get_write(zram->comps[prio]);
 	src = kmap_local_page(page);
 	ret = zcomp_compress(zram->comps[prio], zstrm, src, &comp_len_new);
 	kunmap_local(src);
@@ -2472,7 +2472,7 @@ static int recompress_slot(struct zram *zram, u32 index, struct page *page,
 		*num_recomp_pages -= 1;
 
 	if (ret) {
-		zcomp_stream_put(zstrm);
+		zcomp_stream_put_write(zstrm);
 		return ret;
 	}
 
@@ -2481,7 +2481,7 @@ static int recompress_slot(struct zram *zram, u32 index, struct page *page,
 
 	if (class_index_new >= class_index_old ||
 	    (threshold && comp_len_new >= threshold)) {
-		zcomp_stream_put(zstrm);
+		zcomp_stream_put_write(zstrm);
 
 		/*
 		 * Secondary algorithms failed to re-compress the page
@@ -2510,12 +2510,12 @@ static int recompress_slot(struct zram *zram, u32 index, struct page *page,
 			       __GFP_HIGHMEM | __GFP_MOVABLE,
 			       page_to_nid(page));
 	if (IS_ERR_VALUE(handle_new)) {
-		zcomp_stream_put(zstrm);
+		zcomp_stream_put_write(zstrm);
 		return PTR_ERR((void *)handle_new);
 	}
 
 	zs_obj_write(zram->mem_pool, handle_new, zstrm->buffer, comp_len_new);
-	zcomp_stream_put(zstrm);
+	zcomp_stream_put_write(zstrm);
 
 	slot_free(zram, index);
 	set_slot_handle(zram, index, handle_new);