From 3e30654b509b4c9b9ed6c0df9293113e117832a4 Mon Sep 17 00:00:00 2001
From: "gd.zhou" <gd.zhou@trip.com>
Date: Mon, 14 Jul 2025 21:24:39 +0800
Subject: [PATCH] When async_io is enabled, invoking abortIO may leave pending
 asynchronous requests unprocessed, leading to the inability to handle
 subsequent normal requests. This can ultimately result in thread blocking.

This commit ensures that new requests are properly scheduled and processed even
during or after an abortIO operation, preventing potential deadlocks or stalls.
---
 env/fs_posix.cc | 64 +++++++++++++++++++++++++++++--------------------
 env/io_posix.h  |  2 ++
 2 files changed, 40 insertions(+), 26 deletions(-)

diff --git a/env/fs_posix.cc b/env/fs_posix.cc
index dd2f749350da..90cae77f0cae 100644
--- a/env/fs_posix.cc
+++ b/env/fs_posix.cc
@@ -995,6 +995,31 @@ class PosixFileSystem : public FileSystem {
       return false;
     }
   }
+  void HandleFinishedIO(struct io_uring_cqe* cqe,
+                        Posix_IOHandle* posix_handle) {
+    // Reset cqe data to catch any stray reuse of it
+    static_cast<struct io_uring_cqe*>(cqe)->user_data = 0xd5d5d5d5d5d5d5d5;
+
+    FSReadRequest req;
+    req.scratch = posix_handle->scratch;
+    req.offset = posix_handle->offset;
+    req.len = posix_handle->len;
+
+    size_t finished_len = 0;
+    size_t bytes_read = 0;
+    bool read_again = false;
+    UpdateResult(cqe, "", req.len, posix_handle->iov.iov_len,
+                 true /*async_read*/, posix_handle->use_direct_io,
+                 posix_handle->alignment, finished_len, &req, bytes_read,
+                 read_again);
+    posix_handle->is_finished = true;
+    io_uring_cqe_seen(posix_handle->iu, cqe);
+    posix_handle->cb(req, posix_handle->cb_arg);
+
+    (void)finished_len;
+    (void)bytes_read;
+    (void)read_again;
+  }
 #endif  // ROCKSDB_IOURING_PRESENT
 
   // EXPERIMENTAL
@@ -1043,28 +1068,7 @@ class PosixFileSystem : public FileSystem {
         if (posix_handle->iu != iu) {
           return IOStatus::IOError("");
         }
-        // Reset cqe data to catch any stray reuse of it
-        static_cast<struct io_uring_cqe*>(cqe)->user_data = 0xd5d5d5d5d5d5d5d5;
-
-        FSReadRequest req;
-        req.scratch = posix_handle->scratch;
-        req.offset = posix_handle->offset;
-        req.len = posix_handle->len;
-
-        size_t finished_len = 0;
-        size_t bytes_read = 0;
-        bool read_again = false;
-        UpdateResult(cqe, "", req.len, posix_handle->iov.iov_len,
-                     true /*async_read*/, posix_handle->use_direct_io,
-                     posix_handle->alignment, finished_len, &req, bytes_read,
-                     read_again);
-        posix_handle->is_finished = true;
-        io_uring_cqe_seen(iu, cqe);
-        posix_handle->cb(req, posix_handle->cb_arg);
-
-        (void)finished_len;
-        (void)bytes_read;
-        (void)read_again;
+        HandleFinishedIO(cqe, posix_handle);
 
         if (static_cast<Posix_IOHandle*>(io_handles[i]) == posix_handle) {
           break;
@@ -1121,6 +1125,7 @@ class PosixFileSystem : public FileSystem {
         return IOStatus::IOError("io_uring_submit() requested but returned " +
                                  std::to_string(ret));
       }
+      posix_handle->is_aborted = true;
     }
 
     // After submitting the requests, wait for the requests.
@@ -1146,6 +1151,12 @@ class PosixFileSystem : public FileSystem {
         if (posix_handle->iu != iu) {
           return IOStatus::IOError("");
         }
+        // If the request is not aborted, it means the request is completed
+        // successfully.
+        if (!posix_handle->is_aborted) {
+          HandleFinishedIO(cqe, posix_handle);
+          continue;
+        }
         posix_handle->req_count++;
 
         // Reset cqe data to catch any stray reuse of it
@@ -1163,14 +1174,15 @@ class PosixFileSystem : public FileSystem {
         //
         // Every handle has to wait for 2 requests completion: original one and
         // the cancel request which is tracked by PosixHandle::req_count.
-        if (posix_handle->req_count == 2 &&
-            static_cast<Posix_IOHandle*>(io_handles[i]) == posix_handle) {
+        if (posix_handle->req_count == 2) {
+          // io_uring cancel requests return out of order
           posix_handle->is_finished = true;
           FSReadRequest req;
           req.status = IOStatus::Aborted();
           posix_handle->cb(req, posix_handle->cb_arg);
-
-          break;
+          if (static_cast<Posix_IOHandle*>(io_handles[i]) == posix_handle) {
+            break;
+          }
         }
       }
     }
diff --git a/env/io_posix.h b/env/io_posix.h
index f129668ea546..17a91de82589 100644
--- a/env/io_posix.h
+++ b/env/io_posix.h
@@ -87,6 +87,7 @@ struct Posix_IOHandle {
         use_direct_io(_use_direct_io),
         alignment(_alignment),
         is_finished(false),
+        is_aborted(false),
         req_count(0) {}
 
   struct iovec iov;
@@ -100,6 +101,7 @@ struct Posix_IOHandle {
   size_t alignment;
   bool is_finished;
   // req_count is used by AbortIO API to keep track of number of requests.
+  bool is_aborted;
   uint32_t req_count;
 };