Drivers: hv: balloon: Online memory segments "in context"

index 2b7813e..a7140b3 100644
@@ -541,6 +541,41 @@ struct memory_block *find_memory_block(struct mem_section *section)
 	return find_memory_block_hinted(section, NULL);
 }
 
+/*
+ * Given the start pfn of a memory block; bring the memory
+ * block online. This API would be useful for drivers that may
+ * want to bring "online" the memory that has been hot-added.
+ */
+
+int online_memory_block(unsigned long start_pfn)
+{
+	struct mem_section *cur_section;
+	struct memory_block *cur_memory_block;
+	int ret_val = -EINVAL;
+
+	cur_section = __pfn_to_section(start_pfn);
+
+	if (!valid_section(cur_section))
+		return ret_val;
+
+	mutex_lock(&mem_sysfs_mutex);
+	cur_memory_block = find_memory_block(cur_section);
+
+	if (!cur_memory_block) {
+		mutex_unlock(&mem_sysfs_mutex);
+		return ret_val;
+	}
+
+	lock_device_hotplug();
+	ret_val = memory_block_change_state(cur_memory_block, MEM_ONLINE,
+						MEM_OFFLINE, ONLINE_KEEP);
+	unlock_device_hotplug();
+	mutex_unlock(&mem_sysfs_mutex);
+
+	return ret_val;
+}
+EXPORT_SYMBOL_GPL(online_memory_block);
+
 static struct attribute *memory_memblk_attrs[] = {
 	&dev_attr_phys_index.attr,
 	&dev_attr_end_phys_index.attr,

index 85c31a8..7d0d2a6 100644
@@ -109,12 +109,17 @@ static inline int memory_isolate_notify(unsigned long val, void *v)
 {
 	return 0;
 }
+static inline int online_memory_block(unsigned long start_pfn)
+{
+	return 0;
+}
 #else
 extern int register_memory_notifier(struct notifier_block *nb);
 extern void unregister_memory_notifier(struct notifier_block *nb);
 extern int register_memory_isolate_notifier(struct notifier_block *nb);
 extern void unregister_memory_isolate_notifier(struct notifier_block *nb);
 extern int register_new_memory(int, struct mem_section *);
+extern int online_memory_block(unsigned long start_pfn);
 #ifdef CONFIG_MEMORY_HOTREMOVE
 extern int unregister_memory_section(struct mem_section *);
 #endif

index 2d094cf..c2eec17 100644
@@ -515,11 +515,6 @@ struct hv_dynmem_device {
 	bool host_specified_ha_region;
 
 	/*
-	 * State to synchronize hot-add.
-	 */
-	struct completion  ol_waitevent;
-	bool ha_waiting;
-	/*
 	 * This thread handles hot-add
 	 * requests from the host as well as notifying
 	 * the host with regards to memory pressure in

@@ -581,9 +576,6 @@ static void hv_mem_hot_add(unsigned long start, unsigned long size,
 
 		has->covered_end_pfn +=  processed_pfn;
 
-		init_completion(&dm_device.ol_waitevent);
-		dm_device.ha_waiting = true;
-
 		nid = memory_add_physaddr_to_nid(PFN_PHYS(start_pfn));
 		ret = add_memory(nid, PFN_PHYS((start_pfn)),
 				(HA_CHUNK << PAGE_SHIFT));

@@ -606,12 +598,10 @@ static void hv_mem_hot_add(unsigned long start, unsigned long size,
 		}
 
 		/*
-		 * Wait for the memory block to be onlined.
-		 * Since the hot add has succeeded, it is ok to
-		 * proceed even if the pages in the hot added region
-		 * have not been "onlined" within the allowed time.
+		 * Before proceeding to hot add the next segment,
+		 * online the segment that has been hot added.
 		 */
-		wait_for_completion_timeout(&dm_device.ol_waitevent, 5*HZ);
+		online_memory_block(start_pfn);
 
 	}
 

@@ -625,10 +615,6 @@ static void hv_online_page(struct page *pg)
 	unsigned long cur_start_pgp;
 	unsigned long cur_end_pgp;
 
-	if (dm_device.ha_waiting) {
-		dm_device.ha_waiting = false;
-		complete(&dm_device.ol_waitevent);
-	}
 
 	list_for_each(cur, &dm_device.ha_region_list) {
 		has = list_entry(cur, struct hv_hotadd_state, list);