3 files changed, 37 insertions, 53 deletions

diff --git a/drivers/block/zram/zcomp.c b/drivers/block/zram/zcomp.c
index 4b5cd3a7b2b6..12046f4f00e4 100644
--- a/drivers/block/zram/zcomp.c
+++ b/drivers/block/zram/zcomp.c
@@ -160,82 +160,56 @@ int zcomp_decompress(struct zcomp_strm *zstrm,
 			dst, &dst_len);
 }
 
-static int __zcomp_cpu_notifier(struct zcomp *comp,
-		unsigned long action, unsigned long cpu)
+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;
 
-	switch (action) {
-	case CPU_UP_PREPARE:
-		if (WARN_ON(*per_cpu_ptr(comp->stream, cpu)))
-			break;
-		zstrm = zcomp_strm_alloc(comp);
-		if (IS_ERR_OR_NULL(zstrm)) {
-			pr_err("Can't allocate a compression stream\n");
-			return NOTIFY_BAD;
-		}
-		*per_cpu_ptr(comp->stream, cpu) = zstrm;
-		break;
-	case CPU_DEAD:
-	case CPU_UP_CANCELED:
-		zstrm = *per_cpu_ptr(comp->stream, cpu);
-		if (!IS_ERR_OR_NULL(zstrm))
-			zcomp_strm_free(zstrm);
-		*per_cpu_ptr(comp->stream, cpu) = NULL;
-		break;
-	default:
-		break;
+	if (WARN_ON(*per_cpu_ptr(comp->stream, cpu)))
+		return 0;
+
+	zstrm = zcomp_strm_alloc(comp);
+	if (IS_ERR_OR_NULL(zstrm)) {
+		pr_err("Can't allocate a compression stream\n");
+		return -ENOMEM;
 	}
-	return NOTIFY_OK;
+	*per_cpu_ptr(comp->stream, cpu) = zstrm;
+	return 0;
 }
 
-static int zcomp_cpu_notifier(struct notifier_block *nb,
-		unsigned long action, void *pcpu)
+int zcomp_cpu_dead(unsigned int cpu, struct hlist_node *node)
 {
-	unsigned long cpu = (unsigned long)pcpu;
-	struct zcomp *comp = container_of(nb, typeof(*comp), notifier);
+	struct zcomp *comp = hlist_entry(node, struct zcomp, node);
+	struct zcomp_strm *zstrm;
 
-	return __zcomp_cpu_notifier(comp, action, cpu);
+	zstrm = *per_cpu_ptr(comp->stream, cpu);
+	if (!IS_ERR_OR_NULL(zstrm))
+		zcomp_strm_free(zstrm);
+	*per_cpu_ptr(comp->stream, cpu) = NULL;
+	return 0;
 }
 
 static int zcomp_init(struct zcomp *comp)
 {
-	unsigned long cpu;
 	int ret;
 
-	comp->notifier.notifier_call = zcomp_cpu_notifier;
-
 	comp->stream = alloc_percpu(struct zcomp_strm *);
 	if (!comp->stream)
 		return -ENOMEM;
 
-	cpu_notifier_register_begin();
-	for_each_online_cpu(cpu) {
-		ret = __zcomp_cpu_notifier(comp, CPU_UP_PREPARE, cpu);
-		if (ret == NOTIFY_BAD)
-			goto cleanup;
-	}
-	__register_cpu_notifier(&comp->notifier);
-	cpu_notifier_register_done();
+	ret = cpuhp_state_add_instance(CPUHP_ZCOMP_PREPARE, &comp->node);
+	if (ret < 0)
+		goto cleanup;
 	return 0;
 
 cleanup:
-	for_each_online_cpu(cpu)
-		__zcomp_cpu_notifier(comp, CPU_UP_CANCELED, cpu);
-	cpu_notifier_register_done();
-	return -ENOMEM;
+	free_percpu(comp->stream);
+	return ret;
 }
 
 void zcomp_destroy(struct zcomp *comp)
 {
-	unsigned long cpu;
-
-	cpu_notifier_register_begin();
-	for_each_online_cpu(cpu)
-		__zcomp_cpu_notifier(comp, CPU_UP_CANCELED, cpu);
-	__unregister_cpu_notifier(&comp->notifier);
-	cpu_notifier_register_done();
-
+	cpuhp_state_remove_instance(CPUHP_ZCOMP_PREPARE, &comp->node);
 	free_percpu(comp->stream);
 	kfree(comp);
 }
diff --git a/drivers/block/zram/zcomp.h b/drivers/block/zram/zcomp.h
index 478cac2ed465..41c1002a7d7d 100644
--- a/drivers/block/zram/zcomp.h
+++ b/drivers/block/zram/zcomp.h
@@ -19,11 +19,12 @@ struct zcomp_strm {
 /* dynamic per-device compression frontend */
 struct zcomp {
 	struct zcomp_strm * __percpu *stream;
-	struct notifier_block notifier;
-
 	const char *name;
+	struct hlist_node node;
 };
 
+int zcomp_cpu_up_prepare(unsigned int cpu, struct hlist_node *node);
+int zcomp_cpu_dead(unsigned int cpu, struct hlist_node *node);
 ssize_t zcomp_available_show(const char *comp, char *buf);
 bool zcomp_available_algorithm(const char *comp);
 
diff --git a/drivers/block/zram/zram_drv.c b/drivers/block/zram/zram_drv.c
index 5497f7fc44d0..15f58ab44d0b 100644
--- a/drivers/block/zram/zram_drv.c
+++ b/drivers/block/zram/zram_drv.c
@@ -30,6 +30,7 @@
 #include <linux/err.h>
 #include <linux/idr.h>
 #include <linux/sysfs.h>
+#include <linux/cpuhotplug.h>
 
 #include "zram_drv.h"
 
@@ -1443,15 +1444,22 @@ static void destroy_devices(void)
 	idr_for_each(&zram_index_idr, &zram_remove_cb, NULL);
 	idr_destroy(&zram_index_idr);
 	unregister_blkdev(zram_major, "zram");
+	cpuhp_remove_multi_state(CPUHP_ZCOMP_PREPARE);
 }
 
 static int __init zram_init(void)
 {
 	int ret;
 
+	ret = cpuhp_setup_state_multi(CPUHP_ZCOMP_PREPARE, "block/zram:prepare",
+				      zcomp_cpu_up_prepare, zcomp_cpu_dead);
+	if (ret < 0)
+		return ret;
+
 	ret = class_register(&zram_control_class);
 	if (ret) {
 		pr_err("Unable to register zram-control class\n");
+		cpuhp_remove_multi_state(CPUHP_ZCOMP_PREPARE);
 		return ret;
 	}
 
@@ -1459,6 +1467,7 @@ static int __init zram_init(void)
 	if (zram_major <= 0) {
 		pr_err("Unable to get major number\n");
 		class_unregister(&zram_control_class);
+		cpuhp_remove_multi_state(CPUHP_ZCOMP_PREPARE);
 		return -EBUSY;
 	}