2 files changed, 56 insertions, 18 deletions

diff --git a/sys/dev/hyperv/include/vmbus_xact.h b/sys/dev/hyperv/include/vmbus_xact.h
index 280c5fef5a2d..90711a0be774 100644
--- a/sys/dev/hyperv/include/vmbus_xact.h
+++ b/sys/dev/hyperv/include/vmbus_xact.h
@@ -55,6 +55,8 @@ const void		*vmbus_xact_wait(struct vmbus_xact *xact,
 			    size_t *resp_len);
 const void		*vmbus_xact_busywait(struct vmbus_xact *xact,
 			    size_t *resp_len);
+const void		*vmbus_xact_poll(struct vmbus_xact *xact,
+			    size_t *resp_len);
 void			vmbus_xact_wakeup(struct vmbus_xact *xact,
 			    const void *data, size_t dlen);
 void			vmbus_xact_ctx_wakeup(struct vmbus_xact_ctx *ctx,
diff --git a/sys/dev/hyperv/vmbus/vmbus_xact.c b/sys/dev/hyperv/vmbus/vmbus_xact.c
index c0219e837263..90bdba7e1058 100644
--- a/sys/dev/hyperv/vmbus/vmbus_xact.c
+++ b/sys/dev/hyperv/vmbus/vmbus_xact.c
@@ -71,8 +71,10 @@ static struct vmbus_xact	*vmbus_xact_alloc(struct vmbus_xact_ctx *,
 static void			vmbus_xact_free(struct vmbus_xact *);
 static struct vmbus_xact	*vmbus_xact_get1(struct vmbus_xact_ctx *,
 				    uint32_t);
-const void			*vmbus_xact_wait1(struct vmbus_xact *, size_t *,
+static const void		*vmbus_xact_wait1(struct vmbus_xact *, size_t *,
 				    bool);
+static const void		*vmbus_xact_return(struct vmbus_xact *,
+				    size_t *);
 static void			vmbus_xact_save_resp(struct vmbus_xact *,
 				    const void *, size_t);
 static void			vmbus_xact_ctx_free(struct vmbus_xact_ctx *);
@@ -277,27 +279,13 @@ vmbus_xact_deactivate(struct vmbus_xact *xact)
 	mtx_unlock(&ctx->xc_lock);
 }
 
-const void *
-vmbus_xact_wait1(struct vmbus_xact *xact, size_t *resp_len,
-    bool can_sleep)
+static const void *
+vmbus_xact_return(struct vmbus_xact *xact, size_t *resp_len)
 {
 	struct vmbus_xact_ctx *ctx = xact->x_ctx;
 	const void *resp;
 
-	mtx_lock(&ctx->xc_lock);
-
-	KASSERT(ctx->xc_active == xact, ("xact mismatch"));
-	while (xact->x_resp == NULL &&
-	    (ctx->xc_flags & VMBUS_XACT_CTXF_DESTROY) == 0) {
-		if (can_sleep) {
-			mtx_sleep(&ctx->xc_active, &ctx->xc_lock, 0,
-			    "wxact", 0);
-		} else {
-			mtx_unlock(&ctx->xc_lock);
-			DELAY(1000);
-			mtx_lock(&ctx->xc_lock);
-		}
-	}
+	mtx_assert(&ctx->xc_lock, MA_OWNED);
 	KASSERT(ctx->xc_active == xact, ("xact trashed"));
 
 	if ((ctx->xc_flags & VMBUS_XACT_CTXF_DESTROY) && xact->x_resp == NULL) {
@@ -317,6 +305,32 @@ vmbus_xact_wait1(struct vmbus_xact *xact, size_t *resp_len,
 	resp = xact->x_resp;
 	*resp_len = xact->x_resp_len;
 
+	return (resp);
+}
+
+static const void *
+vmbus_xact_wait1(struct vmbus_xact *xact, size_t *resp_len,
+    bool can_sleep)
+{
+	struct vmbus_xact_ctx *ctx = xact->x_ctx;
+	const void *resp;
+
+	mtx_lock(&ctx->xc_lock);
+
+	KASSERT(ctx->xc_active == xact, ("xact mismatch"));
+	while (xact->x_resp == NULL &&
+	    (ctx->xc_flags & VMBUS_XACT_CTXF_DESTROY) == 0) {
+		if (can_sleep) {
+			mtx_sleep(&ctx->xc_active, &ctx->xc_lock, 0,
+			    "wxact", 0);
+		} else {
+			mtx_unlock(&ctx->xc_lock);
+			DELAY(1000);
+			mtx_lock(&ctx->xc_lock);
+		}
+	}
+	resp = vmbus_xact_return(xact, resp_len);
+
 	mtx_unlock(&ctx->xc_lock);
 
 	return (resp);
@@ -336,6 +350,28 @@ vmbus_xact_busywait(struct vmbus_xact *xact, size_t *resp_len)
 	return (vmbus_xact_wait1(xact, resp_len, false /* can't sleep */));
 }
 
+const void *
+vmbus_xact_poll(struct vmbus_xact *xact, size_t *resp_len)
+{
+	struct vmbus_xact_ctx *ctx = xact->x_ctx;
+	const void *resp;
+
+	mtx_lock(&ctx->xc_lock);
+
+	KASSERT(ctx->xc_active == xact, ("xact mismatch"));
+	if (xact->x_resp == NULL &&
+	    (ctx->xc_flags & VMBUS_XACT_CTXF_DESTROY) == 0) {
+		mtx_unlock(&ctx->xc_lock);
+		*resp_len = 0;
+		return (NULL);
+	}
+	resp = vmbus_xact_return(xact, resp_len);
+
+	mtx_unlock(&ctx->xc_lock);
+
+	return (resp);
+}
+
 static void
 vmbus_xact_save_resp(struct vmbus_xact *xact, const void *data, size_t dlen)
 {