4 files changed, 49 insertions, 9 deletions

diff --git a/drivers/hv/channel_mgmt.c b/drivers/hv/channel_mgmt.c
index 191a6a3ae6c..2c59f030546 100644
--- a/drivers/hv/channel_mgmt.c
+++ b/drivers/hv/channel_mgmt.c
@@ -216,9 +216,16 @@ static void vmbus_process_rescind_offer(struct work_struct *work)
 	unsigned long flags;
 	struct vmbus_channel *primary_channel;
 	struct vmbus_channel_relid_released msg;
+	struct device *dev;
+
+	if (channel->device_obj) {
+		dev = get_device(&channel->device_obj->device);
+		if (dev) {
+			vmbus_device_unregister(channel->device_obj);
+			put_device(dev);
+		}
+	}
 
-	if (channel->device_obj)
-		vmbus_device_unregister(channel->device_obj);
 	memset(&msg, 0, sizeof(struct vmbus_channel_relid_released));
 	msg.child_relid = channel->offermsg.child_relid;
 	msg.header.msgtype = CHANNELMSG_RELID_RELEASED;
diff --git a/drivers/hv/hv_balloon.c b/drivers/hv/hv_balloon.c
index 5e90c5d771a..b958ded8ac7 100644
--- a/drivers/hv/hv_balloon.c
+++ b/drivers/hv/hv_balloon.c
@@ -1087,10 +1087,12 @@ static void balloon_up(struct work_struct *dummy)
 	struct dm_balloon_response *bl_resp;
 	int alloc_unit;
 	int ret;
-	bool alloc_error = false;
+	bool alloc_error;
 	bool done = false;
 	int i;
 
+	/* The host balloons pages in 2M granularity. */
+	WARN_ON_ONCE(num_pages % PAGES_IN_2M != 0);
 
 	/*
 	 * We will attempt 2M allocations. However, if we fail to
@@ -1107,16 +1109,18 @@ static void balloon_up(struct work_struct *dummy)
 
 
 		num_pages -= num_ballooned;
+		alloc_error = false;
 		num_ballooned = alloc_balloon_pages(&dm_device, num_pages,
 						bl_resp, alloc_unit,
 						 &alloc_error);
 
-		if ((alloc_error) && (alloc_unit != 1)) {
+		if (alloc_unit != 1 && num_ballooned == 0) {
 			alloc_unit = 1;
 			continue;
 		}
 
-		if ((alloc_error) || (num_ballooned == num_pages)) {
+		if ((alloc_unit == 1 && alloc_error) ||
+			(num_ballooned == num_pages)) {
 			bl_resp->more_pages = 0;
 			done = true;
 			dm_device.state = DM_INITIALIZED;
diff --git a/drivers/hv/hv_kvp.c b/drivers/hv/hv_kvp.c
index 521c14625b3..beb8105c0e7 100644
--- a/drivers/hv/hv_kvp.c
+++ b/drivers/hv/hv_kvp.c
@@ -350,6 +350,7 @@ kvp_send_key(struct work_struct *dummy)
 	__u8 pool = kvp_transaction.kvp_msg->kvp_hdr.pool;
 	__u32 val32;
 	__u64 val64;
+	int rc;
 
 	msg = kzalloc(sizeof(*msg) + sizeof(struct hv_kvp_msg) , GFP_ATOMIC);
 	if (!msg)
@@ -446,7 +447,13 @@ kvp_send_key(struct work_struct *dummy)
 	}
 
 	msg->len = sizeof(struct hv_kvp_msg);
-	cn_netlink_send(msg, 0, 0, GFP_ATOMIC);
+	rc = cn_netlink_send(msg, 0, 0, GFP_ATOMIC);
+	if (rc) {
+		pr_debug("KVP: failed to communicate to the daemon: %d\n", rc);
+		if (cancel_delayed_work_sync(&kvp_work))
+			kvp_respond_to_host(message, HV_E_FAIL);
+	}
+
 	kfree(msg);
 
 	return;
diff --git a/drivers/hv/hv_snapshot.c b/drivers/hv/hv_snapshot.c
index 34f14fddb66..9d5e0d1efdb 100644
--- a/drivers/hv/hv_snapshot.c
+++ b/drivers/hv/hv_snapshot.c
@@ -28,7 +28,7 @@
 #define VSS_MINOR  0
 #define VSS_VERSION    (VSS_MAJOR << 16 | VSS_MINOR)
 
-
+#define VSS_USERSPACE_TIMEOUT (msecs_to_jiffies(10 * 1000))
 
 /*
  * Global state maintained for transaction that is being processed.
@@ -55,12 +55,24 @@ static const char vss_name[] = "vss_kernel_module";
 static __u8 *recv_buffer;
 
 static void vss_send_op(struct work_struct *dummy);
+static void vss_timeout_func(struct work_struct *dummy);
+
+static DECLARE_DELAYED_WORK(vss_timeout_work, vss_timeout_func);
 static DECLARE_WORK(vss_send_op_work, vss_send_op);
 
 /*
  * Callback when data is received from user mode.
  */
 
+static void vss_timeout_func(struct work_struct *dummy)
+{
+	/*
+	 * Timeout waiting for userspace component to reply happened.
+	 */
+	pr_warn("VSS: timeout waiting for daemon to reply\n");
+	vss_respond_to_host(HV_E_FAIL);
+}
+
 static void
 vss_cn_callback(struct cn_msg *msg, struct netlink_skb_parms *nsp)
 {
@@ -76,13 +88,15 @@ vss_cn_callback(struct cn_msg *msg, struct netlink_skb_parms *nsp)
 		return;
 
 	}
-	vss_respond_to_host(vss_msg->error);
+	if (cancel_delayed_work_sync(&vss_timeout_work))
+		vss_respond_to_host(vss_msg->error);
 }
 
 
 static void vss_send_op(struct work_struct *dummy)
 {
 	int op = vss_transaction.msg->vss_hdr.operation;
+	int rc;
 	struct cn_msg *msg;
 	struct hv_vss_msg *vss_msg;
 
@@ -98,7 +112,12 @@ static void vss_send_op(struct work_struct *dummy)
 	vss_msg->vss_hdr.operation = op;
 	msg->len = sizeof(struct hv_vss_msg);
 
-	cn_netlink_send(msg, 0, 0, GFP_ATOMIC);
+	rc = cn_netlink_send(msg, 0, 0, GFP_ATOMIC);
+	if (rc) {
+		pr_warn("VSS: failed to communicate to the daemon: %d\n", rc);
+		if (cancel_delayed_work_sync(&vss_timeout_work))
+			vss_respond_to_host(HV_E_FAIL);
+	}
 	kfree(msg);
 
 	return;
@@ -223,6 +242,8 @@ void hv_vss_onchannelcallback(void *context)
 			case VSS_OP_FREEZE:
 			case VSS_OP_THAW:
 				schedule_work(&vss_send_op_work);
+				schedule_delayed_work(&vss_timeout_work,
+						      VSS_USERSPACE_TIMEOUT);
 				return;
 
 			case VSS_OP_HOT_BACKUP:
@@ -277,5 +298,6 @@ hv_vss_init(struct hv_util_service *srv)
 void hv_vss_deinit(void)
 {
 	cn_del_callback(&vss_id);
+	cancel_delayed_work_sync(&vss_timeout_work);
 	cancel_work_sync(&vss_send_op_work);
 }