3 files changed, 21 insertions, 10 deletions

diff --git a/mm/vmstat.c b/mm/vmstat.c
index 197b4c4a9587..302dd076b8bf 100644
--- a/mm/vmstat.c
+++ b/mm/vmstat.c
@@ -1298,14 +1298,14 @@ static int __init setup_vmstat(void)
 #ifdef CONFIG_SMP
         int cpu;
 
-        register_cpu_notifier(&vmstat_notifier);
+        cpu_notifier_register_begin();
+        __register_cpu_notifier(&vmstat_notifier);
 
-        get_online_cpus();
         for_each_online_cpu(cpu) {
                 start_cpu_timer(cpu);
                 node_set_state(cpu_to_node(cpu), N_CPU);
         }
-        put_online_cpus();
+        cpu_notifier_register_done();
 #endif
 #ifdef CONFIG_PROC_FS
         proc_create("buddyinfo", S_IRUGO, NULL, &fragmentation_file_operations);
diff --git a/mm/zsmalloc.c b/mm/zsmalloc.c
index c03ca5e9fe15..36b4591a7a2d 100644
--- a/mm/zsmalloc.c
+++ b/mm/zsmalloc.c
@@ -814,21 +814,32 @@ static void zs_exit(void)
 {
         int cpu;
 
+        cpu_notifier_register_begin();
+
         for_each_online_cpu(cpu)
                 zs_cpu_notifier(NULL, CPU_DEAD, (void *)(long)cpu);
-        unregister_cpu_notifier(&zs_cpu_nb);
+        __unregister_cpu_notifier(&zs_cpu_nb);
+
+        cpu_notifier_register_done();
 }
 
 static int zs_init(void)
 {
         int cpu, ret;
 
-        register_cpu_notifier(&zs_cpu_nb);
+        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))
+                if (notifier_to_errno(ret)) {
+                        cpu_notifier_register_done();
                         goto fail;
+                }
         }
+
+        cpu_notifier_register_done();
+
         return 0;
 fail:
         zs_exit();
diff --git a/mm/zswap.c b/mm/zswap.c
index 34b75cc70827..aeaef0fb5624 100644
--- a/mm/zswap.c
+++ b/mm/zswap.c
@@ -388,18 +388,18 @@ static int zswap_cpu_init(void)
 {
         unsigned long cpu;
 
-        get_online_cpus();
+        cpu_notifier_register_begin();
         for_each_online_cpu(cpu)
                 if (__zswap_cpu_notifier(CPU_UP_PREPARE, cpu) != NOTIFY_OK)
                         goto cleanup;
-        register_cpu_notifier(&zswap_cpu_notifier_block);
-        put_online_cpus();
+        __register_cpu_notifier(&zswap_cpu_notifier_block);
+        cpu_notifier_register_done();
         return 0;
 
 cleanup:
         for_each_online_cpu(cpu)
                 __zswap_cpu_notifier(CPU_UP_CANCELED, cpu);
-        put_online_cpus();
+        cpu_notifier_register_done();
         return -ENOMEM;
 }