commit 290de30c0e4a87ed3bd6cdef478376d39a922912
Author: Andres Freund <andres@anarazel.de>
Date:   2025-03-25 16:50:27 -0400

    wip: aio: io_uring wait for in-flight IOs
    
    Author:
    Reviewed-by:
    Discussion: https://postgr.es/m/
    Backpatch:

diff --git a/src/include/storage/aio_internal.h b/src/include/storage/aio_internal.h
index a46d4393ebc..352610f60a8 100644
--- a/src/include/storage/aio_internal.h
+++ b/src/include/storage/aio_internal.h
@@ -287,6 +287,15 @@ typedef struct PgAioCtl
  */
 typedef struct IoMethodOps
 {
+	/* properties */
+
+	/*
+	 * If an FD is about to be closed, do we need to wait for all in-flight
+	 * IOs referencing that FD?
+	 */
+	bool		wait_on_fd_before_close;
+
+
 	/* global initialization */
 
 	/*
@@ -368,6 +377,7 @@ extern PgAioResult pgaio_io_call_complete_local(PgAioHandle *ioh);
 /* aio_io.c */
 extern void pgaio_io_perform_synchronously(PgAioHandle *ioh);
 extern const char *pgaio_io_get_op_name(PgAioHandle *ioh);
+extern bool pgaio_io_uses_fd(PgAioHandle *ioh, int fd);
 
 /* aio_target.c */
 extern bool pgaio_io_can_reopen(PgAioHandle *ioh);
diff --git a/src/backend/storage/aio/aio.c b/src/backend/storage/aio/aio.c
index c512495dce3..8816ccfa0d7 100644
--- a/src/backend/storage/aio/aio.c
+++ b/src/backend/storage/aio/aio.c
@@ -1303,6 +1303,41 @@ pgaio_closing_fd(int fd)
 	 * it's probably not worth it.
 	 */
 	pgaio_submit_staged();
+
+	/*
+	 * If requested by the IO method, wait for all IOs that use the
+	 * to-be-closed FD.
+	 */
+	if (pgaio_method_ops->wait_on_fd_before_close)
+	{
+		/*
+		 * As waiting for one IO to complete may complete multiple IOs, we
+		 * can't just use a mutable list iterator. The maximum number of
+		 * in-flight IOs is fairly small, so just restart the loop after
+		 * waiting for an IO.
+		 */
+		while (!dclist_is_empty(&pgaio_my_backend->in_flight_ios))
+		{
+			dlist_iter	iter;
+			PgAioHandle *ioh = NULL;
+
+			dclist_foreach(iter, &pgaio_my_backend->in_flight_ios)
+			{
+				ioh = dclist_container(PgAioHandle, node, iter.cur);
+
+				if (pgaio_io_uses_fd(ioh, fd))
+					break;
+				else
+					ioh = NULL;
+			}
+
+			if (!ioh)
+				break;
+
+			/* see comment in pgaio_io_wait_for_free() about raciness */
+			pgaio_io_wait(ioh, ioh->generation);
+		}
+	}
 }
 
 /*
diff --git a/src/backend/storage/aio/aio_io.c b/src/backend/storage/aio/aio_io.c
index 97930d687c2..108021fa377 100644
--- a/src/backend/storage/aio/aio_io.c
+++ b/src/backend/storage/aio/aio_io.c
@@ -186,3 +186,25 @@ pgaio_io_get_op_name(PgAioHandle *ioh)
 
 	return NULL;				/* silence compiler */
 }
+
+/*
+ * Used to determine if an IO needs to be waited upon before the file
+ * descriptor can be closed.
+ */
+bool
+pgaio_io_uses_fd(PgAioHandle *ioh, int fd)
+{
+	Assert(ioh->state >= PGAIO_HS_DEFINED);
+
+	switch (ioh->op)
+	{
+		case PGAIO_OP_READV:
+			return ioh->op_data.read.fd == fd;
+		case PGAIO_OP_WRITEV:
+			return ioh->op_data.write.fd == fd;
+		case PGAIO_OP_INVALID:
+			return false;
+	}
+
+	return false;				/* silence compiler */
+}
diff --git a/src/backend/storage/aio/method_io_uring.c b/src/backend/storage/aio/method_io_uring.c
index aab09793684..3b299dcf388 100644
--- a/src/backend/storage/aio/method_io_uring.c
+++ b/src/backend/storage/aio/method_io_uring.c
@@ -58,6 +58,15 @@ static void pgaio_uring_sq_from_io(PgAioHandle *ioh, struct io_uring_sqe *sqe);
 
 
 const IoMethodOps pgaio_uring_ops = {
+	/*
+	 * While io_uring mostly is OK with FDs getting closed while the IO is in
+	 * flight, that is not true for IOs submitted with IOSQE_ASYNC.
+	 *
+	 * See
+	 * https://postgr.es/m/5ons2rtmwarqqhhexb3dnqulw5rjgwgoct57vpdau4rujlrffj%403fls6d2mkiwc
+	 */
+	.wait_on_fd_before_close = true,
+
 	.shmem_size = pgaio_uring_shmem_size,
 	.shmem_init = pgaio_uring_shmem_init,
 	.init_backend = pgaio_uring_init_backend,