[PATCH -next 3/6] nbd: don't clear 'NBD_CMD_INFLIGHT' flag if request is not completed

[Yu Kuai <yukuai3@huawei.com>]

Otherwise io will hung because request will only be completed if the
cmd has the flag 'NBD_CMD_INFLIGHT'.

Fixes: 07175cb1baf4 ("nbd: make sure request completion won't concurrent")
Signed-off-by: Yu Kuai <yukuai3@huawei.com>
---
 drivers/block/nbd.c | 27 ++++++++++++++++++++-------
 1 file changed, 20 insertions(+), 7 deletions(-)

diff --git a/drivers/block/nbd.c b/drivers/block/nbd.c
index c0a787cb5153..4829868706af 100644
--- a/drivers/block/nbd.c
+++ b/drivers/block/nbd.c
@@ -429,6 +429,7 @@ static enum blk_eh_timer_return nbd_xmit_timeout(struct request *req,
 		 * a new connection is reconfigured or util dead timeout.
 		 */
 		if (config->socks) {
+			__set_bit(NBD_CMD_INFLIGHT, &cmd->flags);
 			if (cmd->index < config->num_connections) {
 				struct nbd_sock *nsock =
 					config->socks[cmd->index];
@@ -456,6 +457,8 @@ static enum blk_eh_timer_return nbd_xmit_timeout(struct request *req,
 		 * so just warn and reset the timer.
 		 */
 		struct nbd_sock *nsock = config->socks[cmd->index];
+
+		__set_bit(NBD_CMD_INFLIGHT, &cmd->flags);
 		cmd->retries++;
 		dev_info(nbd_to_dev(nbd), "Possible stuck request %p: control (%s@%llu,%uB). Runtime %u seconds\n",
 			req, nbdcmd_to_ascii(req_to_nbd_cmd_type(req)),
@@ -756,31 +759,31 @@ static struct nbd_cmd *nbd_handle_reply(struct nbd_device *nbd, int index,
 		dev_err(disk_to_dev(nbd->disk), "Unexpected reply %d from different sock %d (expected %d)",
 			tag, index, cmd->index);
 		ret = -ENOENT;
-		goto out;
+		goto out_reset_inflight;
 	}
 	if (cmd->cmd_cookie != nbd_handle_to_cookie(handle)) {
 		dev_err(disk_to_dev(nbd->disk), "Double reply on req %p, cmd_cookie %u, handle cookie %u\n",
 			req, cmd->cmd_cookie, nbd_handle_to_cookie(handle));
 		ret = -ENOENT;
-		goto out;
+		goto out_reset_inflight;
 	}
 	if (cmd->status != BLK_STS_OK) {
 		dev_err(disk_to_dev(nbd->disk), "Command already handled %p\n",
 			req);
 		ret = -ENOENT;
-		goto out;
+		goto out_reset_inflight;
 	}
 	if (test_bit(NBD_CMD_REQUEUED, &cmd->flags)) {
 		dev_err(disk_to_dev(nbd->disk), "Raced with timeout on req %p\n",
 			req);
 		ret = -ENOENT;
-		goto out;
+		goto out_reset_inflight;
 	}
 	if (ntohl(reply->error)) {
 		dev_err(disk_to_dev(nbd->disk), "Other side returned error (%d)\n",
 			ntohl(reply->error));
 		cmd->status = BLK_STS_IOERR;
-		goto out;
+		goto out_reset_inflight;
 	}
 
 	dev_dbg(nbd_to_dev(nbd), "request %p: got reply\n", req);
@@ -803,15 +806,22 @@ static struct nbd_cmd *nbd_handle_reply(struct nbd_device *nbd, int index,
 				 */
 				if (nbd_disconnected(nbd->config)) {
 					cmd->status = BLK_STS_IOERR;
-					goto out;
+					goto out_reset_inflight;
 				}
 				ret = -EIO;
-				goto out;
+				goto out_reset_inflight;
 			}
 			dev_dbg(nbd_to_dev(nbd), "request %p: got %d bytes data\n",
 				req, bvec.bv_len);
 		}
 	}
+out_reset_inflight:
+	if (ret)
+		/*
+		 * Caller will not complete the request, thus set the flag so
+		 * that it can be completed from other context.
+		 */
+		__set_bit(NBD_CMD_INFLIGHT, &cmd->flags);
 out:
 	trace_nbd_payload_received(req, handle);
 	mutex_unlock(&cmd->lock);
@@ -857,6 +867,9 @@ static void recv_work(struct work_struct *work)
 		rq = blk_mq_rq_from_pdu(cmd);
 		if (likely(!blk_should_fake_timeout(rq->q)))
 			blk_mq_complete_request(rq);
+		else
+			/* Timeout rely on this flag to complete request. */
+			__test_and_set_bit(NBD_CMD_INFLIGHT, &cmd->flags);
 		percpu_ref_put(&q->q_usage_counter);
 	}
 
-- 
2.31.1