rt: do not leak fd when cancelling io_uring open operation (#7983)

tokio/src/runtime/io/driver/uring.rs

@@ -2,12 +2,14 @@ use io_uring::{squeue::Entry, IoUring, Probe};
 use mio::unix::SourceFd;
 use slab::Slab;
 
+use crate::runtime::driver::op::CancelData;
+use crate::runtime::driver::op::CqeResult;
 use crate::runtime::driver::op::{Cancellable, Lifecycle};
 use crate::{io::Interest, loom::sync::Mutex};
 
 use super::{Handle, TOKEN_WAKEUP};
 
-use std::os::fd::{AsRawFd, RawFd};
+use std::os::fd::{AsRawFd, FromRawFd, OwnedFd, RawFd};
 use std::{io, mem, task::Waker};
 
 const DEFAULT_RING_SIZE: u32 = 256;
@@ -77,9 +79,16 @@ impl UringContext {
                     waker.wake_by_ref();
                     *ops.get_mut(idx).unwrap() = Lifecycle::Completed(cqe);
                 }
-                Some(Lifecycle::Cancelled(_)) => {
+                Some(Lifecycle::Cancelled(cancel_data)) => {
+                    if let CancelData::Open(_) = cancel_data {
+                        if let Ok(fd) = CqeResult::from(cqe).result {
+                            // SAFETY: the successful CQE result provides
+                            // a non-negative integer, and the event is
+                            // related to an open operation.
+                            unsafe { OwnedFd::from_raw_fd(fd as i32) };
+                        }
+                    }
                     // Op future was cancelled, so we discard the result.
-                    // We just remove the entry from the slab.
                     ops.remove(idx);
                 }
                 Some(other) => {
@@ -147,6 +156,16 @@ impl Drop for UringContext {
 
             for cqe in self.ring_mut().completion() {
                 let idx = cqe.user_data() as usize;
+
+                if let Some(Lifecycle::Cancelled(CancelData::Open(_))) = ops.get_mut(idx) {
+                    if let Ok(fd) = CqeResult::from(cqe).result {
+                        // SAFETY: the successful CQE result provides
+                        // a non-negative integer, and the event is
+                        // related to an open operation.
+                        unsafe { OwnedFd::from_raw_fd(fd as i32) };
+                    }
+                };
+
                 ops.remove(idx);
             }
         }
@@ -272,7 +291,15 @@ impl Handle {
         match mem::replace(lifecycle, Lifecycle::Cancelled(cancel_data)) {
             Lifecycle::Submitted | Lifecycle::Waiting(_) => (),
             // The driver saw the completion, but it was never polled.
-            Lifecycle::Completed(_) => {
+            Lifecycle::Completed(cqe) => {
+                if let Lifecycle::Cancelled(CancelData::Open(_)) = lifecycle {
+                    if let Ok(fd) = CqeResult::from(cqe).result {
+                        // SAFETY: the successful CQE result provides
+                        // a non-negative integer, and the event is
+                        // related to an open operation.
+                        unsafe { OwnedFd::from_raw_fd(fd as i32) };
+                    }
+                }
                 // We can safely remove the entry from the slab, as it has already been completed.
                 ops.remove(index);
             }

tokio/tests/fs_uring_cancel_open.rs

@@ -0,0 +1,71 @@
+//! Uring file operations tests.
+
+#![cfg(all(
+    tokio_unstable,
+    feature = "io-uring",
+    feature = "rt",
+    feature = "fs",
+    target_os = "linux"
+))]
+
+use futures::future::FutureExt;
+use std::fs;
+use std::future::poll_fn;
+use std::task::Poll;
+use tempfile::NamedTempFile;
+
+// see: https://github.com/tokio-rs/tokio/issues/7979
+#[tokio::test]
+async fn file_descriptors_are_closed_when_cancelling_open_op() {
+    let tmp = NamedTempFile::new().unwrap();
+    let path = tmp.path().to_path_buf();
+
+    let fd_count_before_opens = fs::read_dir("/proc/self/fd").unwrap().count();
+
+    for _ in 0..128 {
+        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
+
+        let path = path.clone();
+        let handle = tokio::spawn(async move {
+            poll_fn(|cx| {
+                let opt = {
+                    let mut opt = tokio::fs::OpenOptions::new();
+                    opt.read(true);
+                    opt
+                };
+
+                let fut = opt.open(&path);
+
+                // If io_uring is enabled (and not falling back to the thread pool),
+                // the first poll should return Pending. We don't check if the result
+                // is actually pending because we run some checks on old kernel that
+                // do not support uring.
+                let _pending = Box::pin(fut).poll_unpin(cx);
+
+                tx.send(()).unwrap();
+
+                Poll::<()>::Pending
+            })
+            .await;
+        });
+
+        // Wait for the first poll
+        rx.recv().await.unwrap();
+
+        handle.abort();
+
+        let res = handle.await.unwrap_err();
+        assert!(res.is_cancelled());
+    }
+
+    let fd_count_after_cancel = fs::read_dir("/proc/self/fd").unwrap().count();
+    let leaked = fd_count_after_cancel.saturating_sub(fd_count_before_opens);
+
+    // Since we are opening 128 files, we expect that the related fds
+    // related to this operation will be closed. Since some other fds
+    // can be opened in the meantime, we expect this number to be higher
+    // than the counter before opening the files. This number could be
+    // lower, but to avoid test flakiness we check that this is at most
+    // half the number of the file we opened to check if there's a leak.
+    assert!(leaked <= 64);
+}

tokio/tests/fs_uring_completed_then_dropped_before_repoll.rs

@@ -0,0 +1,130 @@
+#![cfg(all(
+    tokio_unstable,
+    feature = "io-uring",
+    feature = "rt",
+    feature = "fs",
+    target_os = "linux"
+))]
+
+use std::fs;
+use std::future::Future;
+use std::path::PathBuf;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use std::time::Duration;
+
+use tempfile::NamedTempFile;
+use tokio::sync::mpsc::{unbounded_channel, UnboundedSender};
+use tokio::time::timeout;
+
+/// Count currently-open fds in this process.
+fn fd_count() -> usize {
+    fs::read_dir("/proc/self/fd").unwrap().count()
+}
+
+/// First poll:
+/// - polls the inner `tokio::fs::OpenOptions::open()` future once,
+/// - expects `Pending` so we know we took the io_uring path,
+/// - registers the task waker with Tokio's uring machinery.
+///
+/// Second poll:
+/// - happens after the kernel completes the open and Tokio stores the CQE as
+///   `Lifecycle::Completed(cqe)` and wakes the task,
+/// - **intentionally does not poll the inner open future again**,
+/// - stays pending forever so the task can be aborted.
+///
+/// Aborting the task here drops the inner `open()` future while Tokio still has
+/// a completed CQE sitting in the slab.
+struct PollOpenOnceThenNeverRepoll<F> {
+    inner: Pin<Box<F>>,
+    first_poll_tx: Option<UnboundedSender<()>>,
+    second_poll_tx: Option<UnboundedSender<()>>,
+    polled_once: bool,
+}
+
+impl<F: Future> Future for PollOpenOnceThenNeverRepoll<F> {
+    type Output = ();
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        if !self.polled_once {
+            // We don't check if the result is actually pending because
+            // we run some checks on old kernel that do not support uring.
+            let _pending = self.inner.as_mut().poll(cx);
+
+            self.polled_once = true;
+            self.first_poll_tx.take().unwrap().send(()).unwrap();
+            return Poll::Pending;
+        }
+
+        // We were polled again after the inner open completed and woke the task.
+        // Crucially, we do *not* re-poll the inner future here.
+        if let Some(tx) = self.second_poll_tx.take() {
+            tx.send(()).unwrap();
+        }
+
+        Poll::Pending
+    }
+}
+
+async fn completed_then_dropped_before_repoll(path: PathBuf) {
+    let (first_tx, mut first_rx) = unbounded_channel();
+    let (second_tx, mut second_rx) = unbounded_channel();
+
+    let handle = tokio::spawn(async move {
+        let mut opt = tokio::fs::OpenOptions::new();
+        opt.read(true);
+
+        PollOpenOnceThenNeverRepoll {
+            inner: Box::pin(opt.open(&path)),
+            first_poll_tx: Some(first_tx),
+            second_poll_tx: Some(second_tx),
+            polled_once: false,
+        }
+        .await;
+    });
+
+    // Wait until the inner open has been polled once and registered with io_uring.
+    first_rx.recv().await.unwrap();
+
+    // Wait until Tokio wakes the task because the open completed. At this point
+    // the CQE should already be stored as `Lifecycle::Completed(cqe)`.
+    let _ = timeout(Duration::from_secs(2), second_rx.recv()).await;
+
+    // Abort now, before the inner open future gets re-polled and consumes the CQE.
+    handle.abort();
+    let err = handle.await.unwrap_err();
+    assert!(err.is_cancelled(), "task was not cancelled as expected");
+}
+
+#[test]
+fn uring_completed_then_dropped() {
+    let rt = tokio::runtime::Builder::new_multi_thread()
+        .worker_threads(1)
+        .enable_all()
+        .build()
+        .unwrap();
+
+    rt.block_on(async {
+        let before = fd_count();
+        let tmp = NamedTempFile::new().unwrap();
+        let path = tmp.path().to_path_buf();
+
+        for _ in 0..128 {
+            completed_then_dropped_before_repoll(path.clone()).await;
+        }
+
+        // Give completions a moment to settle before counting fds.
+        tokio::time::sleep(Duration::from_millis(250)).await;
+
+        let after = fd_count();
+        let leaked = after.saturating_sub(before);
+
+        // Since we are opening 128 files, we expect that the related fds
+        // related to this operation will be closed. Since some other fds
+        // can be opened in the meantime, we expect this number to be higher
+        // than the counter before opening the files. This number could be
+        // lower, but to avoid test flakiness we check that this is at most
+        // half the number of the file we opened to check if there's a leak.
+        assert!(leaked <= 64);
+    });
+}

tokio/tests/fs_uring_runtime_shutdown.rs

@@ -0,0 +1,78 @@
+//! Uring file operations tests.
+
+#![cfg(all(
+    tokio_unstable,
+    feature = "io-uring",
+    feature = "rt",
+    feature = "fs",
+    target_os = "linux"
+))]
+
+use futures::FutureExt;
+use std::fs;
+use std::future::poll_fn;
+use std::task::Poll;
+use tempfile::NamedTempFile;
+use tokio::runtime::Builder;
+
+// see: https://github.com/tokio-rs/tokio/issues/7979
+#[test]
+fn shutdown_runtime_while_performing_io_uring_ops() {
+    let tmp = NamedTempFile::new().unwrap();
+    let path = tmp.path().to_path_buf();
+
+    let fd_count_before_opens = fs::read_dir("/proc/self/fd").unwrap().count();
+
+    let rt = Builder::new_multi_thread().enable_all().build().unwrap();
+
+    rt.block_on(async {
+        for _ in 0..128 {
+            let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
+
+            let path = path.clone();
+            let handle = tokio::spawn(async move {
+                poll_fn(|cx| {
+                    let opt = {
+                        let mut opt = tokio::fs::OpenOptions::new();
+                        opt.read(true);
+                        opt
+                    };
+
+                    let fut = opt.open(&path);
+
+                    // If io_uring is enabled (and not falling back to the thread pool),
+                    // the first poll should return Pending. We don't check if the result
+                    // is actually a pending because we run some CI checks based on old
+                    // kernels that don't support uring.
+                    let _pending = Box::pin(fut).poll_unpin(cx);
+
+                    tx.send(()).unwrap();
+
+                    Poll::<()>::Pending
+                })
+                .await;
+            });
+
+            // Wait for the first poll
+            rx.recv().await.unwrap();
+
+            handle.abort();
+
+            let res = handle.await.unwrap_err();
+            assert!(res.is_cancelled());
+        }
+    });
+
+    rt.shutdown_background();
+
+    let fd_count_after_cancel = fs::read_dir("/proc/self/fd").unwrap().count();
+    let leaked = fd_count_after_cancel.saturating_sub(fd_count_before_opens);
+
+    // Since we are opening 128 files, we expect that the related fds
+    // related to this operation will be closed. Since some other fds
+    // can be opened in the meantime, we expect this number to be higher
+    // than the counter before opening the files. This number could be
+    // lower, but to avoid test flakiness we check that this is at most
+    // half the number of the file we opened to check if there's a leak.
+    assert!(leaked <= 64);
+}