10 files changed, 318 insertions, 169 deletions

diff --git a/block/nbd-client.c b/block/nbd-client.c
index 25bcaa2346..f0dbea24d3 100644
--- a/block/nbd-client.c
+++ b/block/nbd-client.c
@@ -34,7 +34,7 @@
 #define HANDLE_TO_INDEX(bs, handle) ((handle) ^ ((uint64_t)(intptr_t)bs))
 #define INDEX_TO_HANDLE(bs, index)  ((index)  ^ ((uint64_t)(intptr_t)bs))
 
-static void nbd_recv_coroutines_enter_all(NBDClientSession *s)
+static void nbd_recv_coroutines_wake_all(NBDClientSession *s)
 {
     int i;
 
@@ -112,7 +112,7 @@ static coroutine_fn void nbd_read_reply_entry(void *opaque)
     }
 
     s->quit = true;
-    nbd_recv_coroutines_enter_all(s);
+    nbd_recv_coroutines_wake_all(s);
     s->read_reply_co = NULL;
 }
 
@@ -144,12 +144,12 @@ static int nbd_co_send_request(BlockDriverState *bs,
     request->handle = INDEX_TO_HANDLE(s, i);
 
     if (s->quit) {
-        qemu_co_mutex_unlock(&s->send_mutex);
-        return -EIO;
+        rc = -EIO;
+        goto err;
     }
     if (!s->ioc) {
-        qemu_co_mutex_unlock(&s->send_mutex);
-        return -EPIPE;
+        rc = -EPIPE;
+        goto err;
     }
 
     if (qiov) {
@@ -166,8 +166,13 @@ static int nbd_co_send_request(BlockDriverState *bs,
     } else {
         rc = nbd_send_request(s->ioc, request);
     }
+
+err:
     if (rc < 0) {
         s->quit = true;
+        s->requests[i].coroutine = NULL;
+        s->in_flight--;
+        qemu_co_queue_next(&s->free_sema);
     }
     qemu_co_mutex_unlock(&s->send_mutex);
     return rc;
@@ -201,13 +206,6 @@ static void nbd_co_receive_reply(NBDClientSession *s,
         /* Tell the read handler to read another header.  */
         s->reply.handle = 0;
     }
-}
-
-static void nbd_coroutine_end(BlockDriverState *bs,
-                              NBDRequest *request)
-{
-    NBDClientSession *s = nbd_get_client_session(bs);
-    int i = HANDLE_TO_INDEX(s, request->handle);
 
     s->requests[i].coroutine = NULL;
 
@@ -222,29 +220,40 @@ static void nbd_coroutine_end(BlockDriverState *bs,
     qemu_co_mutex_unlock(&s->send_mutex);
 }
 
+static int nbd_co_request(BlockDriverState *bs,
+                          NBDRequest *request,
+                          QEMUIOVector *qiov)
+{
+    NBDClientSession *client = nbd_get_client_session(bs);
+    NBDReply reply;
+    int ret;
+
+    assert(!qiov || request->type == NBD_CMD_WRITE ||
+           request->type == NBD_CMD_READ);
+    ret = nbd_co_send_request(bs, request,
+                              request->type == NBD_CMD_WRITE ? qiov : NULL);
+    if (ret < 0) {
+        reply.error = -ret;
+    } else {
+        nbd_co_receive_reply(client, request, &reply,
+                             request->type == NBD_CMD_READ ? qiov : NULL);
+    }
+    return -reply.error;
+}
+
 int nbd_client_co_preadv(BlockDriverState *bs, uint64_t offset,
                          uint64_t bytes, QEMUIOVector *qiov, int flags)
 {
-    NBDClientSession *client = nbd_get_client_session(bs);
     NBDRequest request = {
         .type = NBD_CMD_READ,
         .from = offset,
         .len = bytes,
     };
-    NBDReply reply;
-    ssize_t ret;
 
     assert(bytes <= NBD_MAX_BUFFER_SIZE);
     assert(!flags);
 
-    ret = nbd_co_send_request(bs, &request, NULL);
-    if (ret < 0) {
-        reply.error = -ret;
-    } else {
-        nbd_co_receive_reply(client, &request, &reply, qiov);
-    }
-    nbd_coroutine_end(bs, &request);
-    return -reply.error;
+    return nbd_co_request(bs, &request, qiov);
 }
 
 int nbd_client_co_pwritev(BlockDriverState *bs, uint64_t offset,
@@ -256,8 +265,6 @@ int nbd_client_co_pwritev(BlockDriverState *bs, uint64_t offset,
         .from = offset,
         .len = bytes,
     };
-    NBDReply reply;
-    ssize_t ret;
 
     if (flags & BDRV_REQ_FUA) {
         assert(client->info.flags & NBD_FLAG_SEND_FUA);
@@ -266,27 +273,18 @@ int nbd_client_co_pwritev(BlockDriverState *bs, uint64_t offset,
 
     assert(bytes <= NBD_MAX_BUFFER_SIZE);
 
-    ret = nbd_co_send_request(bs, &request, qiov);
-    if (ret < 0) {
-        reply.error = -ret;
-    } else {
-        nbd_co_receive_reply(client, &request, &reply, NULL);
-    }
-    nbd_coroutine_end(bs, &request);
-    return -reply.error;
+    return nbd_co_request(bs, &request, qiov);
 }
 
 int nbd_client_co_pwrite_zeroes(BlockDriverState *bs, int64_t offset,
                                 int bytes, BdrvRequestFlags flags)
 {
-    ssize_t ret;
     NBDClientSession *client = nbd_get_client_session(bs);
     NBDRequest request = {
         .type = NBD_CMD_WRITE_ZEROES,
         .from = offset,
         .len = bytes,
     };
-    NBDReply reply;
 
     if (!(client->info.flags & NBD_FLAG_SEND_WRITE_ZEROES)) {
         return -ENOTSUP;
@@ -300,22 +298,13 @@ int nbd_client_co_pwrite_zeroes(BlockDriverState *bs, int64_t offset,
         request.flags |= NBD_CMD_FLAG_NO_HOLE;
     }
 
-    ret = nbd_co_send_request(bs, &request, NULL);
-    if (ret < 0) {
-        reply.error = -ret;
-    } else {
-        nbd_co_receive_reply(client, &request, &reply, NULL);
-    }
-    nbd_coroutine_end(bs, &request);
-    return -reply.error;
+    return nbd_co_request(bs, &request, NULL);
 }
 
 int nbd_client_co_flush(BlockDriverState *bs)
 {
     NBDClientSession *client = nbd_get_client_session(bs);
     NBDRequest request = { .type = NBD_CMD_FLUSH };
-    NBDReply reply;
-    ssize_t ret;
 
     if (!(client->info.flags & NBD_FLAG_SEND_FLUSH)) {
         return 0;
@@ -324,14 +313,7 @@ int nbd_client_co_flush(BlockDriverState *bs)
     request.from = 0;
     request.len = 0;
 
-    ret = nbd_co_send_request(bs, &request, NULL);
-    if (ret < 0) {
-        reply.error = -ret;
-    } else {
-        nbd_co_receive_reply(client, &request, &reply, NULL);
-    }
-    nbd_coroutine_end(bs, &request);
-    return -reply.error;
+    return nbd_co_request(bs, &request, NULL);
 }
 
 int nbd_client_co_pdiscard(BlockDriverState *bs, int64_t offset, int bytes)
@@ -342,22 +324,12 @@ int nbd_client_co_pdiscard(BlockDriverState *bs, int64_t offset, int bytes)
         .from = offset,
         .len = bytes,
     };
-    NBDReply reply;
-    ssize_t ret;
 
     if (!(client->info.flags & NBD_FLAG_SEND_TRIM)) {
         return 0;
     }
 
-    ret = nbd_co_send_request(bs, &request, NULL);
-    if (ret < 0) {
-        reply.error = -ret;
-    } else {
-        nbd_co_receive_reply(client, &request, &reply, NULL);
-    }
-    nbd_coroutine_end(bs, &request);
-    return -reply.error;
-
+    return nbd_co_request(bs, &request, NULL);
 }
 
 void nbd_client_detach_aio_context(BlockDriverState *bs)
diff --git a/include/block/nbd.h b/include/block/nbd.h
index 9c3d0a5868..040cdd2e60 100644
--- a/include/block/nbd.h
+++ b/include/block/nbd.h
@@ -163,8 +163,8 @@ int nbd_receive_negotiate(QIOChannel *ioc, const char *name,
                           Error **errp);
 int nbd_init(int fd, QIOChannelSocket *sioc, NBDExportInfo *info,
              Error **errp);
-ssize_t nbd_send_request(QIOChannel *ioc, NBDRequest *request);
-ssize_t nbd_receive_reply(QIOChannel *ioc, NBDReply *reply, Error **errp);
+int nbd_send_request(QIOChannel *ioc, NBDRequest *request);
+int nbd_receive_reply(QIOChannel *ioc, NBDReply *reply, Error **errp);
 int nbd_client(int fd);
 int nbd_disconnect(int fd);
 
diff --git a/nbd/client.c b/nbd/client.c
index 0a17de80b5..68a0bc1ffc 100644
--- a/nbd/client.c
+++ b/nbd/client.c
@@ -399,12 +399,10 @@ static int nbd_opt_go(QIOChannel *ioc, const char *wantname,
                phase, but make sure it sent flags */
             if (len) {
                 error_setg(errp, "server sent invalid NBD_REP_ACK");
-                nbd_send_opt_abort(ioc);
                 return -1;
             }
             if (!info->flags) {
                 error_setg(errp, "broken server omitted NBD_INFO_EXPORT");
-                nbd_send_opt_abort(ioc);
                 return -1;
             }
             trace_nbd_opt_go_success();
@@ -898,7 +896,7 @@ int nbd_disconnect(int fd)
 }
 #endif
 
-ssize_t nbd_send_request(QIOChannel *ioc, NBDRequest *request)
+int nbd_send_request(QIOChannel *ioc, NBDRequest *request)
 {
     uint8_t buf[NBD_REQUEST_SIZE];
 
@@ -916,22 +914,22 @@ ssize_t nbd_send_request(QIOChannel *ioc, NBDRequest *request)
     return nbd_write(ioc, buf, sizeof(buf), NULL);
 }
 
-ssize_t nbd_receive_reply(QIOChannel *ioc, NBDReply *reply, Error **errp)
+/* nbd_receive_reply
+ * Returns 1 on success
+ *         0 on eof, when no data was read (errp is not set)
+ *         negative errno on failure (errp is set)
+ */
+int nbd_receive_reply(QIOChannel *ioc, NBDReply *reply, Error **errp)
 {
     uint8_t buf[NBD_REPLY_SIZE];
     uint32_t magic;
-    ssize_t ret;
+    int ret;
 
     ret = nbd_read_eof(ioc, buf, sizeof(buf), errp);
     if (ret <= 0) {
         return ret;
     }
 
-    if (ret != sizeof(buf)) {
-        error_setg(errp, "read failed");
-        return -EINVAL;
-    }
-
     /* Reply
        [ 0 ..  3]    magic   (NBD_REPLY_MAGIC)
        [ 4 ..  7]    error   (0 == no error)
@@ -955,6 +953,7 @@ ssize_t nbd_receive_reply(QIOChannel *ioc, NBDReply *reply, Error **errp)
         error_setg(errp, "invalid magic (got 0x%" PRIx32 ")", magic);
         return -EINVAL;
     }
-    return sizeof(buf);
+
+    return 1;
 }
 
diff --git a/nbd/nbd-internal.h b/nbd/nbd-internal.h
index 396ddb4d3e..03549e3f39 100644
--- a/nbd/nbd-internal.h
+++ b/nbd/nbd-internal.h
@@ -77,21 +77,36 @@
 #define NBD_ESHUTDOWN  108
 
 /* nbd_read_eof
- * Tries to read @size bytes from @ioc. Returns number of bytes actually read.
- * May return a value >= 0 and < size only on EOF, i.e. when iteratively called
- * qio_channel_readv() returns 0. So, there is no need to call nbd_read_eof
- * iteratively.
+ * Tries to read @size bytes from @ioc.
+ * Returns 1 on success
+ *         0 on eof, when no data was read (errp is not set)
+ *         negative errno on failure (errp is set)
  */
-static inline ssize_t nbd_read_eof(QIOChannel *ioc, void *buffer, size_t size,
-                                   Error **errp)
+static inline int nbd_read_eof(QIOChannel *ioc, void *buffer, size_t size,
+                               Error **errp)
 {
     struct iovec iov = { .iov_base = buffer, .iov_len = size };
+    ssize_t ret;
+
     /* Sockets are kept in blocking mode in the negotiation phase.  After
      * that, a non-readable socket simply means that another thread stole
      * our request/reply.  Synchronization is done with recv_coroutine, so
      * that this is coroutine-safe.
      */
-    return nbd_rwv(ioc, &iov, 1, size, true, errp);
+
+    assert(size);
+
+    ret = nbd_rwv(ioc, &iov, 1, size, true, errp);
+    if (ret <= 0) {
+        return ret;
+    }
+
+    if (ret != size) {
+        error_setg(errp, "End of file");
+        return -EINVAL;
+    }
+
+    return 1;
 }
 
 /* nbd_read
@@ -100,9 +115,9 @@ static inline ssize_t nbd_read_eof(QIOChannel *ioc, void *buffer, size_t size,
 static inline int nbd_read(QIOChannel *ioc, void *buffer, size_t size,
                            Error **errp)
 {
-    ssize_t ret = nbd_read_eof(ioc, buffer, size, errp);
+    int ret = nbd_read_eof(ioc, buffer, size, errp);
 
-    if (ret >= 0 && ret != size) {
+    if (ret == 0) {
         ret = -EINVAL;
         error_setg(errp, "End of file");
     }
diff --git a/tests/qemu-iotests/083 b/tests/qemu-iotests/083
index bff9360048..0306f112da 100755
--- a/tests/qemu-iotests/083
+++ b/tests/qemu-iotests/083
@@ -27,6 +27,14 @@ echo "QA output created by $seq"
 here=`pwd`
 status=1	# failure is the default!
 
+_cleanup()
+{
+	rm -f nbd.sock
+	rm -f nbd-fault-injector.out
+	rm -f nbd-fault-injector.conf
+}
+trap "_cleanup; exit \$status" 0 1 2 3 15
+
 # get standard environment, filters and checks
 . ./common.rc
 . ./common.filter
@@ -35,81 +43,105 @@ _supported_fmt generic
 _supported_proto nbd
 _supported_os Linux
 
-# Pick a TCP port based on our pid.  This way multiple instances of this test
-# can run in parallel without conflicting.
-choose_tcp_port() {
-	echo $((($$ % 31744) + 1024)) # 1024 <= port < 32768
-}
-
-wait_for_tcp_port() {
-	while ! (netstat --tcp --listening --numeric | \
-		 grep "$1.*0\\.0\\.0\\.0:\\*.*LISTEN") >/dev/null 2>&1; do
-		sleep 0.1
+check_disconnect() {
+	local event export_name=foo extra_args nbd_addr nbd_url proto when
+
+	while true; do
+		case $1 in
+		--classic-negotiation)
+			shift
+			extra_args=--classic-negotiation
+			export_name=
+			;;
+		--tcp)
+			shift
+			proto=tcp
+			;;
+		--unix)
+			shift
+			proto=unix
+			;;
+		*)
+			break
+			;;
+		esac
 	done
-}
 
-check_disconnect() {
 	event=$1
 	when=$2
-	negotiation=$3
 	echo "=== Check disconnect $when $event ==="
 	echo
 
-	port=$(choose_tcp_port)
-
 	cat > "$TEST_DIR/nbd-fault-injector.conf" <<EOF
 [inject-error]
 event=$event
 when=$when
 EOF
 
-	if [ "$negotiation" = "--classic-negotiation" ]; then
-		extra_args=--classic-negotiation
-		nbd_url="nbd:127.0.0.1:$port"
+	if [ "$proto" = "tcp" ]; then
+		nbd_addr="127.0.0.1:0"
 	else
-		nbd_url="nbd:127.0.0.1:$port:exportname=foo"
+		nbd_addr="$TEST_DIR/nbd.sock"
+	fi
+
+	rm -f "$TEST_DIR/nbd.sock"
+
+	$PYTHON nbd-fault-injector.py $extra_args "$nbd_addr" "$TEST_DIR/nbd-fault-injector.conf" >"$TEST_DIR/nbd-fault-injector.out" 2>&1 &
+
+	# Wait for server to be ready
+	while ! grep -q 'Listening on ' "$TEST_DIR/nbd-fault-injector.out"; do
+		sleep 0.1
+	done
+
+	# Extract the final address (port number has now been assigned in tcp case)
+	nbd_addr=$(sed 's/Listening on \(.*\)$/\1/' "$TEST_DIR/nbd-fault-injector.out")
+
+	if [ "$proto" = "tcp" ]; then
+		nbd_url="nbd+tcp://$nbd_addr/$export_name"
+	else
+		nbd_url="nbd+unix:///$export_name?socket=$nbd_addr"
 	fi
 
-	$PYTHON nbd-fault-injector.py $extra_args "127.0.0.1:$port" "$TEST_DIR/nbd-fault-injector.conf" >/dev/null 2>&1 &
-	wait_for_tcp_port "127\\.0\\.0\\.1:$port"
 	$QEMU_IO -c "read 0 512" "$nbd_url" 2>&1 | _filter_qemu_io | _filter_nbd
 
 	echo
 }
 
-for event in neg1 "export" neg2 request reply data; do
-	for when in before after; do
-		check_disconnect "$event" "$when"
-	done
-
-	# Also inject short replies from the NBD server
-	case "$event" in
-	neg1)
-		for when in 8 16; do
-			check_disconnect "$event" "$when"
-		done
-		;;
-	"export")
-		for when in 4 12 16; do
-			check_disconnect "$event" "$when"
+for proto in tcp unix; do
+	for event in neg1 "export" neg2 request reply data; do
+		for when in before after; do
+			check_disconnect "--$proto" "$event" "$when"
 		done
-		;;
-	neg2)
-		for when in 8 10; do
-			check_disconnect "$event" "$when"
-		done
-		;;
-	reply)
-		for when in 4 8; do
-			check_disconnect "$event" "$when"
-		done
-		;;
-	esac
-done
 
-# Also check classic negotiation without export information
-for when in before 8 16 24 28 after; do
-	check_disconnect "neg-classic" "$when" --classic-negotiation
+		# Also inject short replies from the NBD server
+		case "$event" in
+		neg1)
+			for when in 8 16; do
+				check_disconnect "--$proto" "$event" "$when"
+			done
+			;;
+		"export")
+			for when in 4 12 16; do
+				check_disconnect "--$proto" "$event" "$when"
+			done
+			;;
+		neg2)
+			for when in 8 10; do
+				check_disconnect "--$proto" "$event" "$when"
+			done
+			;;
+		reply)
+			for when in 4 8; do
+				check_disconnect "--$proto" "$event" "$when"
+			done
+			;;
+		esac
+	done
+
+	# Also check classic negotiation without export information
+	for when in before 8 16 24 28 after; do
+		check_disconnect "--$proto" --classic-negotiation "neg-classic" "$when"
+	done
 done
 
 # success, all done
diff --git a/tests/qemu-iotests/083.out b/tests/qemu-iotests/083.out
index a24c6bfece..fb71b6f8ad 100644
--- a/tests/qemu-iotests/083.out
+++ b/tests/qemu-iotests/083.out
@@ -1,43 +1,43 @@
 QA output created by 083
 === Check disconnect before neg1 ===
 
-can't open device nbd:127.0.0.1:PORT:exportname=foo
+can't open device nbd+tcp://127.0.0.1:PORT/foo
 
 === Check disconnect after neg1 ===
 
-can't open device nbd:127.0.0.1:PORT:exportname=foo
+can't open device nbd+tcp://127.0.0.1:PORT/foo
 
 === Check disconnect 8 neg1 ===
 
-can't open device nbd:127.0.0.1:PORT:exportname=foo
+can't open device nbd+tcp://127.0.0.1:PORT/foo
 
 === Check disconnect 16 neg1 ===
 
-can't open device nbd:127.0.0.1:PORT:exportname=foo
+can't open device nbd+tcp://127.0.0.1:PORT/foo
 
 === Check disconnect before export ===
 
-can't open device nbd:127.0.0.1:PORT:exportname=foo
+can't open device nbd+tcp://127.0.0.1:PORT/foo
 
 === Check disconnect after export ===
 
-can't open device nbd:127.0.0.1:PORT:exportname=foo
+can't open device nbd+tcp://127.0.0.1:PORT/foo
 
 === Check disconnect 4 export ===
 
-can't open device nbd:127.0.0.1:PORT:exportname=foo
+can't open device nbd+tcp://127.0.0.1:PORT/foo
 
 === Check disconnect 12 export ===
 
-can't open device nbd:127.0.0.1:PORT:exportname=foo
+can't open device nbd+tcp://127.0.0.1:PORT/foo
 
 === Check disconnect 16 export ===
 
-can't open device nbd:127.0.0.1:PORT:exportname=foo
+can't open device nbd+tcp://127.0.0.1:PORT/foo
 
 === Check disconnect before neg2 ===
 
-can't open device nbd:127.0.0.1:PORT:exportname=foo
+can't open device nbd+tcp://127.0.0.1:PORT/foo
 
 === Check disconnect after neg2 ===
 
@@ -45,11 +45,11 @@ read failed: Input/output error
 
 === Check disconnect 8 neg2 ===
 
-can't open device nbd:127.0.0.1:PORT:exportname=foo
+can't open device nbd+tcp://127.0.0.1:PORT/foo
 
 === Check disconnect 10 neg2 ===
 
-can't open device nbd:127.0.0.1:PORT:exportname=foo
+can't open device nbd+tcp://127.0.0.1:PORT/foo
 
 === Check disconnect before request ===
 
@@ -69,12 +69,12 @@ read failed: Input/output error
 
 === Check disconnect 4 reply ===
 
-read failed
+End of file
 read failed: Input/output error
 
 === Check disconnect 8 reply ===
 
-read failed
+End of file
 read failed: Input/output error
 
 === Check disconnect before data ===
@@ -88,23 +88,134 @@ read 512/512 bytes at offset 0
 
 === Check disconnect before neg-classic ===
 
-can't open device nbd:127.0.0.1:PORT
+can't open device nbd+tcp://127.0.0.1:PORT/
 
 === Check disconnect 8 neg-classic ===
 
-can't open device nbd:127.0.0.1:PORT
+can't open device nbd+tcp://127.0.0.1:PORT/
 
 === Check disconnect 16 neg-classic ===
 
-can't open device nbd:127.0.0.1:PORT
+can't open device nbd+tcp://127.0.0.1:PORT/
 
 === Check disconnect 24 neg-classic ===
 
-can't open device nbd:127.0.0.1:PORT
+can't open device nbd+tcp://127.0.0.1:PORT/
 
 === Check disconnect 28 neg-classic ===
 
-can't open device nbd:127.0.0.1:PORT
+can't open device nbd+tcp://127.0.0.1:PORT/
+
+=== Check disconnect after neg-classic ===
+
+read failed: Input/output error
+
+=== Check disconnect before neg1 ===
+
+can't open device nbd+unix:///foo?socket=TEST_DIR/nbd.sock
+
+=== Check disconnect after neg1 ===
+
+can't open device nbd+unix:///foo?socket=TEST_DIR/nbd.sock
+
+=== Check disconnect 8 neg1 ===
+
+can't open device nbd+unix:///foo?socket=TEST_DIR/nbd.sock
+
+=== Check disconnect 16 neg1 ===
+
+can't open device nbd+unix:///foo?socket=TEST_DIR/nbd.sock
+
+=== Check disconnect before export ===
+
+can't open device nbd+unix:///foo?socket=TEST_DIR/nbd.sock
+
+=== Check disconnect after export ===
+
+can't open device nbd+unix:///foo?socket=TEST_DIR/nbd.sock
+
+=== Check disconnect 4 export ===
+
+can't open device nbd+unix:///foo?socket=TEST_DIR/nbd.sock
+
+=== Check disconnect 12 export ===
+
+can't open device nbd+unix:///foo?socket=TEST_DIR/nbd.sock
+
+=== Check disconnect 16 export ===
+
+can't open device nbd+unix:///foo?socket=TEST_DIR/nbd.sock
+
+=== Check disconnect before neg2 ===
+
+can't open device nbd+unix:///foo?socket=TEST_DIR/nbd.sock
+
+=== Check disconnect after neg2 ===
+
+read failed: Input/output error
+
+=== Check disconnect 8 neg2 ===
+
+can't open device nbd+unix:///foo?socket=TEST_DIR/nbd.sock
+
+=== Check disconnect 10 neg2 ===
+
+can't open device nbd+unix:///foo?socket=TEST_DIR/nbd.sock
+
+=== Check disconnect before request ===
+
+read failed: Input/output error
+
+=== Check disconnect after request ===
+
+read failed: Input/output error
+
+=== Check disconnect before reply ===
+
+read failed: Input/output error
+
+=== Check disconnect after reply ===
+
+read failed: Input/output error
+
+=== Check disconnect 4 reply ===
+
+End of file
+read failed: Input/output error
+
+=== Check disconnect 8 reply ===
+
+End of file
+read failed: Input/output error
+
+=== Check disconnect before data ===
+
+read failed: Input/output error
+
+=== Check disconnect after data ===
+
+read 512/512 bytes at offset 0
+512 bytes, X ops; XX:XX:XX.X (XXX YYY/sec and XXX ops/sec)
+
+=== Check disconnect before neg-classic ===
+
+can't open device nbd+unix:///?socket=TEST_DIR/nbd.sock
+
+=== Check disconnect 8 neg-classic ===
+
+can't open device nbd+unix:///?socket=TEST_DIR/nbd.sock
+
+=== Check disconnect 16 neg-classic ===
+
+can't open device nbd+unix:///?socket=TEST_DIR/nbd.sock
+
+=== Check disconnect 24 neg-classic ===
+
+can't open device nbd+unix:///?socket=TEST_DIR/nbd.sock
+
+=== Check disconnect 28 neg-classic ===
+
+can't open device nbd+unix:///?socket=TEST_DIR/nbd.sock
 
 === Check disconnect after neg-classic ===
 
diff --git a/tests/qemu-iotests/194 b/tests/qemu-iotests/194
index 8028111e21..a3e3bad664 100755
--- a/tests/qemu-iotests/194
+++ b/tests/qemu-iotests/194
@@ -46,16 +46,17 @@ iotests.log('Launching NBD server on destination...')
 iotests.log(dest_vm.qmp('nbd-server-start', addr={'type': 'unix', 'data': {'path': nbd_sock_path}}))
 iotests.log(dest_vm.qmp('nbd-server-add', device='drive0', writable=True))
 
-iotests.log('Starting drive-mirror on source...')
+iotests.log('Starting `drive-mirror` on source...')
 iotests.log(source_vm.qmp(
               'drive-mirror',
               device='drive0',
               target='nbd+unix:///drive0?socket={0}'.format(nbd_sock_path),
               sync='full',
               format='raw', # always raw, the server handles the format
-              mode='existing'))
+              mode='existing',
+              job_id='mirror-job0'))
 
-iotests.log('Waiting for drive-mirror to complete...')
+iotests.log('Waiting for `drive-mirror` to complete...')
 iotests.log(source_vm.event_wait('BLOCK_JOB_READY'),
             filters=[iotests.filter_qmp_event])
 
@@ -67,7 +68,17 @@ dest_vm.qmp('migrate-set-capabilities',
 iotests.log(source_vm.qmp('migrate', uri='unix:{0}'.format(migration_sock_path)))
 
 while True:
-    event = source_vm.event_wait('MIGRATION')
-    iotests.log(event, filters=[iotests.filter_qmp_event])
-    if event['data']['status'] in ('completed', 'failed'):
+    event1 = source_vm.event_wait('MIGRATION')
+    iotests.log(event1, filters=[iotests.filter_qmp_event])
+    if event1['data']['status'] in ('completed', 'failed'):
+        iotests.log('Gracefully ending the `drive-mirror` job on source...')
+        iotests.log(source_vm.qmp('block-job-cancel', device='mirror-job0'))
+        break
+
+while True:
+    event2 = source_vm.event_wait('BLOCK_JOB_COMPLETED')
+    iotests.log(event2, filters=[iotests.filter_qmp_event])
+    if event2['event'] == 'BLOCK_JOB_COMPLETED':
+        iotests.log('Stopping the NBD server on destination...')
+        iotests.log(dest_vm.qmp('nbd-server-stop'))
         break
diff --git a/tests/qemu-iotests/194.out b/tests/qemu-iotests/194.out
index ae501fecac..50ac50da5e 100644
--- a/tests/qemu-iotests/194.out
+++ b/tests/qemu-iotests/194.out
@@ -2,12 +2,17 @@ Launching VMs...
 Launching NBD server on destination...
 {u'return': {}}
 {u'return': {}}
-Starting drive-mirror on source...
+Starting `drive-mirror` on source...
 {u'return': {}}
-Waiting for drive-mirror to complete...
-{u'timestamp': {u'seconds': 'SECS', u'microseconds': 'USECS'}, u'data': {u'device': u'drive0', u'type': u'mirror', u'speed': 0, u'len': 1073741824, u'offset': 1073741824}, u'event': u'BLOCK_JOB_READY'}
+Waiting for `drive-mirror` to complete...
+{u'timestamp': {u'seconds': 'SECS', u'microseconds': 'USECS'}, u'data': {u'device': u'mirror-job0', u'type': u'mirror', u'speed': 0, u'len': 1073741824, u'offset': 1073741824}, u'event': u'BLOCK_JOB_READY'}
 Starting migration...
 {u'return': {}}
 {u'timestamp': {u'seconds': 'SECS', u'microseconds': 'USECS'}, u'data': {u'status': u'setup'}, u'event': u'MIGRATION'}
 {u'timestamp': {u'seconds': 'SECS', u'microseconds': 'USECS'}, u'data': {u'status': u'active'}, u'event': u'MIGRATION'}
 {u'timestamp': {u'seconds': 'SECS', u'microseconds': 'USECS'}, u'data': {u'status': u'completed'}, u'event': u'MIGRATION'}
+Gracefully ending the `drive-mirror` job on source...
+{u'return': {}}
+{u'timestamp': {u'seconds': 'SECS', u'microseconds': 'USECS'}, u'data': {u'device': u'mirror-job0', u'type': u'mirror', u'speed': 0, u'len': 1073741824, u'offset': 1073741824}, u'event': u'BLOCK_JOB_COMPLETED'}
+Stopping the NBD server on destination...
+{u'return': {}}
diff --git a/tests/qemu-iotests/common.filter b/tests/qemu-iotests/common.filter
index 7a58e57317..9d5442ecd9 100644
--- a/tests/qemu-iotests/common.filter
+++ b/tests/qemu-iotests/common.filter
@@ -170,9 +170,9 @@ _filter_nbd()
     #
     # Filter out the TCP port number since this changes between runs.
     sed -e '/nbd\/.*\.c:/d' \
-        -e 's#nbd:\(//\)\?127\.0\.0\.1:[0-9]*#nbd:\1127.0.0.1:PORT#g' \
+        -e 's#127\.0\.0\.1:[0-9]*#127.0.0.1:PORT#g' \
         -e "s#?socket=$TEST_DIR#?socket=TEST_DIR#g" \
-        -e 's#\(exportname=foo\|PORT\): Failed to .*$#\1#'
+        -e 's#\(foo\|PORT/\?\|.sock\): Failed to .*$#\1#'
 }
 
 # make sure this script returns success
diff --git a/tests/qemu-iotests/nbd-fault-injector.py b/tests/qemu-iotests/nbd-fault-injector.py
index 6c07191a5a..1c10dcb51c 100755
--- a/tests/qemu-iotests/nbd-fault-injector.py
+++ b/tests/qemu-iotests/nbd-fault-injector.py
@@ -235,11 +235,15 @@ def open_socket(path):
         sock = socket.socket()
         sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
         sock.bind((host, int(port)))
+
+        # If given port was 0 the final port number is now available
+        path = '%s:%d' % sock.getsockname()
     else:
         sock = socket.socket(socket.AF_UNIX)
         sock.bind(path)
     sock.listen(0)
     print 'Listening on %s' % path
+    sys.stdout.flush() # another process may be waiting, show message now
     return sock
 
 def usage(args):