1 files changed, 24 insertions, 12 deletions

diff --git a/drivers/xen/balloon.c b/drivers/xen/balloon.c
index 37d06ea624aa..dd7954922942 100644
--- a/drivers/xen/balloon.c
+++ b/drivers/xen/balloon.c
@@ -592,19 +592,29 @@ static void __init balloon_add_region(unsigned long start_pfn,
         }
 }
 
+static int alloc_balloon_scratch_page(int cpu)
+{
+        if (per_cpu(balloon_scratch_page, cpu) != NULL)
+                return 0;
+
+        per_cpu(balloon_scratch_page, cpu) = alloc_page(GFP_KERNEL);
+        if (per_cpu(balloon_scratch_page, cpu) == NULL) {
+                pr_warn("Failed to allocate balloon_scratch_page for cpu %d\n", cpu);
+                return -ENOMEM;
+        }
+
+        return 0;
+}
+
+
 static int balloon_cpu_notify(struct notifier_block *self,
                                     unsigned long action, void *hcpu)
 {
         int cpu = (long)hcpu;
         switch (action) {
         case CPU_UP_PREPARE:
-                if (per_cpu(balloon_scratch_page, cpu) != NULL)
-                        break;
-                per_cpu(balloon_scratch_page, cpu) = alloc_page(GFP_KERNEL);
-                if (per_cpu(balloon_scratch_page, cpu) == NULL) {
-                        pr_warn("Failed to allocate balloon_scratch_page for cpu %d\n", cpu);
+                if (alloc_balloon_scratch_page(cpu))
                         return NOTIFY_BAD;
-                }
                 break;
         default:
                 break;
@@ -624,15 +634,17 @@ static int __init balloon_init(void)
                 return -ENODEV;
 
         if (!xen_feature(XENFEAT_auto_translated_physmap)) {
-                for_each_online_cpu(cpu)
-                {
-                        per_cpu(balloon_scratch_page, cpu) = alloc_page(GFP_KERNEL);
-                        if (per_cpu(balloon_scratch_page, cpu) == NULL) {
-                                pr_warn("Failed to allocate balloon_scratch_page for cpu %d\n", cpu);
+                register_cpu_notifier(&balloon_cpu_notifier);
+
+                get_online_cpus();
+                for_each_online_cpu(cpu) {
+                        if (alloc_balloon_scratch_page(cpu)) {
+                                put_online_cpus();
+                                unregister_cpu_notifier(&balloon_cpu_notifier);
                                 return -ENOMEM;
                         }
                 }
-                register_cpu_notifier(&balloon_cpu_notifier);
+                put_online_cpus();
         }
 
         pr_info("Initialising balloon driver\n");