fix invalid response id after closing file in blocking context

Signed-off-by: silver-ymz <yinmingzhuo@gmail.com>

src/auxiliary.rs

@@ -3,7 +3,7 @@ use crate::{lowlevel::Extensions, SftpAuxiliaryData};
 use std::sync::atomic::{AtomicU64, AtomicU8, AtomicUsize, Ordering};
 
 use once_cell::sync::OnceCell;
-use tokio::sync::Notify;
+use tokio::{runtime::Handle, sync::Notify};
 use tokio_util::sync::CancellationToken;
 
 #[derive(Debug, Copy, Clone)]
@@ -55,13 +55,16 @@ pub(super) struct Auxiliary {
     pub(super) auxiliary_data: SftpAuxiliaryData,
 
     pub(super) tokio_compat_file_write_limit: usize,
+
+    pub(super) tokio_handle: Handle,
 }
 
 impl Auxiliary {
     pub(super) fn new(
         max_pending_requests: u16,
         auxiliary_data: SftpAuxiliaryData,
         tokio_compat_file_write_limit: usize,
+        tokio_handle: Handle,
     ) -> Self {
         Self {
             conn_info: OnceCell::new(),
@@ -83,6 +86,8 @@ impl Auxiliary {
             auxiliary_data,
 
             tokio_compat_file_write_limit,
+
+            tokio_handle,
         }
     }
 
@@ -162,4 +167,8 @@ impl Auxiliary {
     pub(super) fn tokio_compat_file_write_limit(&self) -> usize {
         self.tokio_compat_file_write_limit
     }
+
+    pub(super) fn tokio_handle(&self) -> &Handle {
+        &self.tokio_handle
+    }
 }

src/changelog.rs

@@ -1,6 +1,8 @@
 #[allow(unused_imports)]
 use crate::*;
 
+/// # Fixed
+/// Invalid response id after closing file caused by the change in v0.13.8
 #[doc(hidden)]
 pub mod unreleased {}
 

src/file/tokio_compat_file.rs

@@ -22,10 +22,7 @@ use std::{
 use bytes::{Buf, Bytes, BytesMut};
 use derive_destructure2::destructure;
 use pin_project::{pin_project, pinned_drop};
-use tokio::{
-    io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, ReadBuf},
-    runtime,
-};
+use tokio::io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, ReadBuf};
 use tokio_io_utility::ready;
 use tokio_util::sync::WaitForCancellationFutureOwned;
 
@@ -747,31 +744,24 @@ impl TokioCompatFile {
 /// when they should not fail.
 #[pinned_drop]
 impl PinnedDrop for TokioCompatFile {
-    fn drop(self: Pin<&mut Self>) {
-        let this = self.project();
+    fn drop(mut self: Pin<&mut Self>) {
+        let this = self.as_mut().project();
 
         let file = this.inner.clone();
         let read_future = this.read_future.take();
         let write_futures = mem::take(this.write_futures);
 
-        let cancellation_fut = file
-            .inner
-            .get_auxiliary()
-            .cancel_token
-            .clone()
-            .cancelled_owned();
+        let cancellation_fut = self.auxiliary().cancel_token.clone().cancelled_owned();
 
         let do_drop_fut = Self::do_drop(file, read_future, write_futures);
 
-        if let Ok(handle) = runtime::Handle::try_current() {
-            handle.spawn(async move {
-                tokio::select! {
-                    biased;
+        self.auxiliary().tokio_handle().spawn(async move {
+            tokio::select! {
+                biased;
 
-                    _ = cancellation_fut => (),
-                    _ = do_drop_fut => (),
-                }
-            });
-        }
+                _ = cancellation_fut => (),
+                _ = do_drop_fut => (),
+            }
+        });
     }
 }

src/fs/dir.rs

@@ -18,7 +18,6 @@ use std::{
 
 use futures_core::stream::{FusedStream, Stream};
 use pin_project::{pin_project, pinned_drop};
-use tokio::runtime;
 use tokio_util::sync::WaitForCancellationFutureOwned;
 
 type ResponseFuture = crate::lowlevel::AwaitableNameEntriesFuture<crate::Buffer>;
@@ -195,15 +194,13 @@ impl PinnedDrop for ReadDir {
         let cancellation_fut = dir.0.get_auxiliary().cancel_token.clone().cancelled_owned();
         let do_drop_fut = Self::do_drop(dir, future);
 
-        if let Ok(handle) = runtime::Handle::try_current() {
-            handle.spawn(async move {
-                tokio::select! {
-                    biased;
+        this.dir.0.get_auxiliary().tokio_handle().spawn(async move {
+            tokio::select! {
+                biased;
 
-                    _ = cancellation_fut => (),
-                    _ = do_drop_fut => (),
-                }
-            });
-        }
+                _ = cancellation_fut => (),
+                _ = do_drop_fut => (),
+            }
+        });
     }
 }

src/handle.rs

@@ -40,16 +40,14 @@ impl Drop for OwnedHandle {
                     //    size of the Future blows out, becomes double of its size.
                     // 3. the more states the Futures have, the harder it is to optimize and take advantage of the niche.
                     let future = response.wait();
-                    if let Ok(handle) = tokio::runtime::Handle::try_current() {
-                        handle.spawn(async move {
-                            let _res = future.await;
-                            #[cfg(feature = "tracing")]
-                            match _res {
-                                Ok(_) => tracing::debug!("close handle success"),
-                                Err(err) => tracing::error!(?err, "failed to close handle"),
-                            }
-                        });
-                    }
+                    self.get_auxiliary().tokio_handle().spawn(async move {
+                        let _res = future.await;
+                        #[cfg(feature = "tracing")]
+                        match _res {
+                            Ok(_) => tracing::debug!("close handle success"),
+                            Err(err) => tracing::error!(?err, "failed to close handle"),
+                        }
+                    });
                 }
                 Err(_err) => {
                     #[cfg(feature = "tracing")]

src/sftp.rs

@@ -18,6 +18,7 @@ use std::{
 use derive_destructure2::destructure;
 use tokio::{
     io::{AsyncRead, AsyncWrite},
+    runtime::Handle,
     sync::oneshot::Receiver,
     task::JoinHandle,
 };
@@ -146,6 +147,7 @@ impl Sftp {
                 options.get_max_pending_requests(),
                 auxiliary,
                 options.get_tokio_compat_file_write_limit(),
+                Handle::current(),
             ))?;
 
             let flush_task = create_flush_task(
@@ -171,13 +173,15 @@ impl Sftp {
         max_pending_requests: u16,
         auxiliary: SftpAuxiliaryData,
         tokio_compat_file_write_limit: usize,
+        tokio_handle: Handle,
     ) -> Result<WriteEnd, Error> {
         connect(
             MpscQueue::with_capacity(write_end_buffer_size),
             Auxiliary::new(
                 max_pending_requests,
                 auxiliary,
                 tokio_compat_file_write_limit,
+                tokio_handle,
             ),
         )
     }