Use tokio::sync::oneshot to prevent FuturesUnordered reentrant drop crash (#9)

Author: sgrebnov

vortex-io/src/file/driver.rs

@@ -5,7 +5,12 @@ use std::pin::Pin;
 use std::task::{Context, Poll};
 
 use futures::Stream;
+#[cfg(all(test, not(feature = "tokio")))]
+use oneshot;
 use pin_project_lite::pin_project;
+// Prefer tokio::sync::oneshot when tokio feature is enabled
+#[cfg(all(test, feature = "tokio"))]
+use tokio::sync::oneshot;
 use vortex_error::VortexExpect;
 use vortex_metrics::VortexMetrics;
 

vortex-io/src/file/read/mod.rs

@@ -15,8 +15,13 @@ use async_trait::async_trait;
 use futures::channel::mpsc;
 use futures::future::{BoxFuture, Shared};
 use futures::{FutureExt, TryFutureExt};
+#[cfg(not(feature = "tokio"))]
+use oneshot;
 pub use request::*;
 pub use source::*;
+// Prefer tokio::sync::oneshot when tokio feature is enabled
+#[cfg(feature = "tokio")]
+use tokio::sync::oneshot;
 use vortex_buffer::{Alignment, ByteBuffer};
 use vortex_error::{SharedVortexResult, VortexError, VortexResult, vortex_err};
 

vortex-io/src/file/read/request.rs

@@ -6,6 +6,11 @@ use std::fmt::{Debug, Formatter};
 use std::ops::Range;
 use std::sync::Arc;
 
+#[cfg(not(feature = "tokio"))]
+use oneshot;
+// Prefer tokio::sync::oneshot when tokio feature is enabled
+#[cfg(feature = "tokio")]
+use tokio::sync::oneshot;
 use vortex_buffer::{Alignment, ByteBuffer};
 use vortex_error::{VortexError, VortexExpect, VortexResult};
 
@@ -118,6 +123,12 @@ impl Debug for ReadRequest {
 
 impl ReadRequest {
     pub(crate) fn resolve(self, result: VortexResult<ByteBuffer>) {
+        // tokio::sync::oneshot::Sender::send returns Err with the value if receiver dropped
+        #[cfg(feature = "tokio")]
+        if self.callback.send(result).is_err() {
+            log::debug!("ReadRequest {} dropped before resolving", self.id);
+        }
+        #[cfg(not(feature = "tokio"))]
         if let Err(e) = self.callback.send(result) {
             log::debug!("ReadRequest {} dropped before resolving: {e}", self.id);
         }

vortex-io/src/runtime/handle.rs

@@ -7,6 +7,11 @@ use std::task::{Context, Poll, ready};
 
 use futures::channel::mpsc;
 use futures::{FutureExt, StreamExt};
+#[cfg(not(feature = "tokio"))]
+use oneshot;
+// Prefer tokio::sync::oneshot when tokio feature is enabled
+#[cfg(feature = "tokio")]
+use tokio::sync::oneshot;
 use vortex_error::{VortexResult, vortex_panic};
 use vortex_metrics::VortexMetrics;
 

vortex-io/src/runtime/single.rs

@@ -7,8 +7,13 @@ use std::sync::Arc;
 use futures::future::BoxFuture;
 use futures::stream::LocalBoxStream;
 use futures::{Stream, StreamExt};
+#[cfg(not(feature = "tokio"))]
+use oneshot;
 use parking_lot::Mutex;
 use smol::LocalExecutor;
+// Prefer tokio::sync::oneshot when tokio feature is enabled
+#[cfg(feature = "tokio")]
+use tokio::sync::oneshot;
 use vortex_error::vortex_panic;
 
 use crate::runtime::smol::SmolAbortHandle;