diff --git a/drivers/hv/hv_snapshot.c b/drivers/hv/hv_snapshot.c
index 815405f..a548ae4 100644
--- a/drivers/hv/hv_snapshot.c
+++ b/drivers/hv/hv_snapshot.c
@@ -53,7 +53,6 @@
 	struct vmbus_channel *recv_channel; /* chn we got the request */
 	u64 recv_req_id; /* request ID. */
 	struct hv_vss_msg  *msg; /* current message */
-	void *vss_context; /* for the channel callback */
 } vss_transaction;
 
 
@@ -74,6 +73,13 @@
 static DECLARE_DELAYED_WORK(vss_timeout_work, vss_timeout_func);
 static DECLARE_WORK(vss_send_op_work, vss_send_op);
 
+static void vss_poll_wrapper(void *channel)
+{
+	/* Transaction is finished, reset the state here to avoid races. */
+	vss_transaction.state = HVUTIL_READY;
+	hv_vss_onchannelcallback(channel);
+}
+
 /*
  * Callback when data is received from user mode.
  */
@@ -86,12 +92,7 @@
 	pr_warn("VSS: timeout waiting for daemon to reply\n");
 	vss_respond_to_host(HV_E_FAIL);
 
-	/* Transaction is finished, reset the state. */
-	if (vss_transaction.state > HVUTIL_READY)
-		vss_transaction.state = HVUTIL_READY;
-
-	hv_poll_channel(vss_transaction.vss_context,
-			hv_vss_onchannelcallback);
+	hv_poll_channel(vss_transaction.recv_channel, vss_poll_wrapper);
 }
 
 static int vss_handle_handshake(struct hv_vss_msg *vss_msg)
@@ -138,9 +139,8 @@
 		if (cancel_delayed_work_sync(&vss_timeout_work)) {
 			vss_respond_to_host(vss_msg->error);
 			/* Transaction is finished, reset the state. */
-			vss_transaction.state = HVUTIL_READY;
-			hv_poll_channel(vss_transaction.vss_context,
-					hv_vss_onchannelcallback);
+			hv_poll_channel(vss_transaction.recv_channel,
+					vss_poll_wrapper);
 		}
 	} else {
 		/* This is a spurious call! */
@@ -238,15 +238,8 @@
 	struct icmsg_hdr *icmsghdrp;
 	struct icmsg_negotiate *negop = NULL;
 
-	if (vss_transaction.state > HVUTIL_READY) {
-		/*
-		 * We will defer processing this callback once
-		 * the current transaction is complete.
-		 */
-		vss_transaction.vss_context = context;
+	if (vss_transaction.state > HVUTIL_READY)
 		return;
-	}
-	vss_transaction.vss_context = NULL;
 
 	vmbus_recvpacket(channel, recv_buffer, PAGE_SIZE * 2, &recvlen,
 			 &requestid);