1 files changed, 24 insertions, 15 deletions

diff --git a/mm/zsmalloc.c b/mm/zsmalloc.c
index b3b57ef85830..5481c677e465 100644
--- a/mm/zsmalloc.c
+++ b/mm/zsmalloc.c
@@ -881,14 +881,10 @@ static struct notifier_block zs_cpu_nb = {
         .notifier_call = zs_cpu_notifier
 };
 
-static void zs_exit(void)
+static void zs_unregister_cpu_notifier(void)
 {
         int cpu;
 
-#ifdef CONFIG_ZPOOL
-        zpool_unregister_driver(&zs_zpool_driver);
-#endif
-
         cpu_notifier_register_begin();
 
         for_each_online_cpu(cpu)
@@ -898,31 +894,44 @@ static void zs_exit(void)
         cpu_notifier_register_done();
 }
 
-static int zs_init(void)
+static int zs_register_cpu_notifier(void)
 {
-        int cpu, ret;
+        int cpu, uninitialized_var(ret);
 
         cpu_notifier_register_begin();
 
         __register_cpu_notifier(&zs_cpu_nb);
         for_each_online_cpu(cpu) {
                 ret = zs_cpu_notifier(NULL, CPU_UP_PREPARE, (void *)(long)cpu);
-                if (notifier_to_errno(ret)) {
-                        cpu_notifier_register_done();
-                        goto fail;
-                }
+                if (notifier_to_errno(ret))
+                        break;
         }
 
         cpu_notifier_register_done();
+        return notifier_to_errno(ret);
+}
 
+static void __exit zs_exit(void)
+{
 #ifdef CONFIG_ZPOOL
-        zpool_register_driver(&zs_zpool_driver);
+        zpool_unregister_driver(&zs_zpool_driver);
 #endif
+        zs_unregister_cpu_notifier();
+}
+
+static int __init zs_init(void)
+{
+        int ret = zs_register_cpu_notifier();
+
+        if (ret) {
+                zs_unregister_cpu_notifier();
+                return ret;
+        }
 
+#ifdef CONFIG_ZPOOL
+        zpool_register_driver(&zs_zpool_driver);
+#endif
         return 0;
-fail:
-        zs_exit();
-        return notifier_to_errno(ret);
 }
 
 static unsigned int get_maxobj_per_zspage(int size, int pages_per_zspage)