[PATCH 1/1] block: Manage bio references so the bio persists until necessary

[Bijan Mottahedeh <bijan.mottahedeh@xxxxxxxxxx>]

Get a reference to a bio, so it won't be freed if end_io() gets to
it before submit_io() returns.  Defer the release of the first bio
in a mult-bio request until the last end_io() since the first bio is
embedded in the dio structure and must therefore persist through an
entire multi-bio request.

Signed-off-by: Bijan Mottahedeh <bijan.mottahedeh@xxxxxxxxxx>
---
 fs/block_dev.c | 78 ++++++++++++++++++++++++++++++----------------------------
 1 file changed, 40 insertions(+), 38 deletions(-)

diff --git a/fs/block_dev.c b/fs/block_dev.c
index 69bf2fb..19fff6b 100644
--- a/fs/block_dev.c
+++ b/fs/block_dev.c
@@ -303,7 +303,9 @@ static void blkdev_bio_end_io(struct bio *bio)
 	if (bio->bi_status && !dio->bio.bi_status)
 		dio->bio.bi_status = bio->bi_status;
 
-	if (!dio->multi_bio || atomic_dec_and_test(&dio->ref)) {
+	if (atomic_dec_and_test(&dio->ref)) {
+		if (dio->multi_bio)
+			bio_put(&dio->bio);
 		if (!dio->is_sync) {
 			struct kiocb *iocb = dio->iocb;
 			ssize_t ret;
@@ -316,8 +318,6 @@ static void blkdev_bio_end_io(struct bio *bio)
 			}
 
 			dio->iocb->ki_complete(iocb, ret, 0);
-			if (dio->multi_bio)
-				bio_put(&dio->bio);
 		} else {
 			struct task_struct *waiter = dio->waiter;
 
@@ -348,6 +348,7 @@ static void blkdev_bio_end_io(struct bio *bio)
 	loff_t pos = iocb->ki_pos;
 	blk_qc_t qc = BLK_QC_T_NONE;
 	int ret = 0;
+	int nr_bios = 1;
 
 	if ((pos | iov_iter_alignment(iter)) &
 	    (bdev_logical_block_size(bdev) - 1))
@@ -357,16 +358,15 @@ static void blkdev_bio_end_io(struct bio *bio)
 
 	dio = container_of(bio, struct blkdev_dio, bio);
 	dio->is_sync = is_sync = is_sync_kiocb(iocb);
-	if (dio->is_sync) {
+	if (dio->is_sync)
 		dio->waiter = current;
-		bio_get(bio);
-	} else {
+	else
 		dio->iocb = iocb;
-	}
 
 	dio->size = 0;
 	dio->multi_bio = false;
 	dio->should_dirty = is_read && iter_is_iovec(iter);
+	atomic_set(&dio->ref, 1);
 
 	/*
 	 * Don't plug for HIPRI/polled IO, as those should go straight
@@ -375,7 +375,7 @@ static void blkdev_bio_end_io(struct bio *bio)
 	if (!is_poll)
 		blk_start_plug(&plug);
 
-	for (;;) {
+	do {
 		bio_set_dev(bio, bdev);
 		bio->bi_iter.bi_sector = pos >> 9;
 		bio->bi_write_hint = iocb->ki_hint;
@@ -403,62 +403,64 @@ static void blkdev_bio_end_io(struct bio *bio)
 		pos += bio->bi_iter.bi_size;
 
 		nr_pages = iov_iter_npages(iter, BIO_MAX_PAGES);
-		if (!nr_pages) {
-			bool polled = false;
+		if (!nr_pages && is_poll)
+			bio_set_polled(bio, iocb);
 
-			if (iocb->ki_flags & IOCB_HIPRI) {
-				bio_set_polled(bio, iocb);
-				polled = true;
-			}
+		/*
+		 * Get a reference to a bio, so it won't be freed
+		 * if end_io() gets to it before submit_io() returns.
+		 * Defer the release of the first bio in a mult-bio
+		 * request until the last end_io() since the first bio
+		 * is embedded in the dio structure and must therefore
+		 * persist through an entire multi-bio request.
+		 */
+		bio_get(bio);
 
-			qc = submit_bio(bio);
+		qc = submit_bio(bio);
 
-			if (polled)
-				WRITE_ONCE(iocb->ki_cookie, qc);
-			break;
-		}
+		if (bio->bi_opf & REQ_HIPRI)
+			WRITE_ONCE(iocb->ki_cookie, qc);
 
-		if (!dio->multi_bio) {
-			/*
-			 * AIO needs an extra reference to ensure the dio
-			 * structure which is embedded into the first bio
-			 * stays around.
-			 */
-			if (!is_sync)
-				bio_get(bio);
-			dio->multi_bio = true;
-			atomic_set(&dio->ref, 2);
-		} else {
+		if (nr_pages) {
+			if (++nr_bios == 2)
+				dio->multi_bio = true;
+			else
+				bio_put(bio);
+			bio = bio_alloc(GFP_KERNEL, nr_pages);
 			atomic_inc(&dio->ref);
 		}
-
-		submit_bio(bio);
-		bio = bio_alloc(GFP_KERNEL, nr_pages);
-	}
+	} while (nr_pages);
 
 	if (!is_poll)
 		blk_finish_plug(&plug);
 
-	if (!is_sync)
-		return -EIOCBQUEUED;
+	if (!is_sync) {
+		ret = -EIOCBQUEUED;
+		goto done;
+	}
+
+	if (!blk_qc_t_valid(qc))
+		goto done;
 
 	for (;;) {
 		set_current_state(TASK_UNINTERRUPTIBLE);
 		if (!READ_ONCE(dio->waiter))
 			break;
 
-		if (!(iocb->ki_flags & IOCB_HIPRI) ||
+		if (!is_poll ||
 		    !blk_poll(bdev_get_queue(bdev), qc, true))
 			io_schedule();
 	}
 	__set_current_state(TASK_RUNNING);
 
+done:
 	if (!ret)
 		ret = blk_status_to_errno(dio->bio.bi_status);
 	if (likely(!ret))
 		ret = dio->size;
 
-	bio_put(&dio->bio);
+	if (!dio->multi_bio)
+		bio_put(&dio->bio);
 	return ret;
 }
 
-- 
1.8.3.1