Replace Kanal with async-channel (#11040)

Use async-channel for bounded writer pipelines, runtime dispatch, and
CUDA completion notifications. Consume receivers as streams directly and
remove the unused channel adapter. Preserve queue capacities and use
nonblocking sends for unbounded dispatch and single-shot completion
callbacks.

---------

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Author: gatesn

Cargo.lock

@@ -112,6 +112,7 @@ arrow-ipc = "59.2"
 arrow-row = "59.2"
 arrow-schema = "59.2"
 arrow-select = "59.2"
+async-channel = "2.5.0"
 async-executor = "1.14.0"
 async-fs = "2.2.0"
 async-lock = "3.4"
@@ -184,7 +185,6 @@ inventory = "0.3.20"
 itertools = "0.14.0"
 jiff = "0.2.28"
 jni = { version = "0.22.0" }
-kanal = "0.1.1"
 lasso = { version = "0.7", features = ["multi-threaded", "inline-more"] }
 lending-iterator = "0.1.7"
 libfuzzer-sys = "0.4"

Cargo.toml

@@ -22,10 +22,10 @@ _test-harness = []
 
 [dependencies]
 arrow-schema = { workspace = true, features = ["ffi"] }
+async-channel = { workspace = true }
 async-trait = { workspace = true }
 futures = { workspace = true, features = ["executor"] }
 itertools = { workspace = true }
-kanal = { workspace = true }
 memmap2 = { workspace = true }
 num-traits = { workspace = true }
 object_store = { workspace = true, features = ["fs"] }

vortex-cuda/Cargo.toml

@@ -9,8 +9,9 @@ use std::mem::size_of_val;
 use std::ops::Deref;
 use std::sync::Arc;
 
+use async_channel::Receiver;
+use async_channel::Sender;
 use futures::future::BoxFuture;
-use kanal::Sender;
 use tracing::warn;
 use vortex::array::buffer::BufferHandle;
 use vortex::error::VortexResult;
@@ -221,8 +222,8 @@ pub(crate) async fn await_stream_callback(stream: &CudaStream) -> VortexResult<(
 /// # Errors
 ///
 /// Returns an error if registering the host callback function fails.
-fn register_stream_callback(stream: &CudaStream) -> VortexResult<kanal::AsyncReceiver<()>> {
-    let (tx, rx) = kanal::bounded::<()>(1);
+fn register_stream_callback(stream: &CudaStream) -> VortexResult<Receiver<()>> {
+    let (tx, rx) = async_channel::bounded::<()>(1);
 
     let tx_ptr = Box::into_raw(Box::new(tx));
 
@@ -237,10 +238,10 @@ fn register_stream_callback(stream: &CudaStream) -> VortexResult<kanal::AsyncRec
         // before. We have unique ownership and can therefore free it.
         let tx = unsafe { Box::from_raw(user_data as *mut Sender<()>) };
 
-        // Blocking send as we're in a callback invoked by the CUDA driver.
+        // The callback sends once into an empty capacity-one channel.
         // NOTE: send can fail if the CudaEvent is dropped by the caller, in which case the receiver
         //  is closed and sends will fail.
-        if let Err(_e) = tx.send(()) {
+        if let Err(_e) = tx.try_send(()) {
             warn!(error = ?_e, "register_stream_callback send failed due to error");
         }
     }
@@ -263,7 +264,7 @@ fn register_stream_callback(stream: &CudaStream) -> VortexResult<kanal::AsyncRec
         })?;
     }
 
-    Ok(rx.to_async())
+    Ok(rx)
 }
 
 #[cfg(test)]

vortex-cuda/src/stream.rs

@@ -17,12 +17,12 @@ all-features = true
 [dependencies]
 arrow-array = { workspace = true }
 arrow-schema = { workspace = true }
+async-channel = { workspace = true }
 async-trait = { workspace = true }
 bytes = { workspace = true }
 flatbuffers = { workspace = true }
 futures = { workspace = true, features = ["std", "async-await"] }
 itertools = { workspace = true }
-kanal = { workspace = true }
 moka = { workspace = true, features = ["sync"] }
 object_store = { workspace = true, optional = true }
 parking_lot = { workspace = true }

vortex-file/Cargo.toml

@@ -5,6 +5,7 @@ use std::sync::Arc;
 use std::sync::atomic::AtomicU64;
 use std::sync::atomic::Ordering;
 
+use async_channel::Sender;
 use async_trait::async_trait;
 use parking_lot::Mutex;
 use vortex_buffer::Alignment;
@@ -18,13 +19,13 @@ use vortex_layout::sequence::SequenceId;
 use crate::footer::SegmentSpec;
 
 pub struct BufferedSegmentSink {
-    buffers: kanal::AsyncSender<ByteBuffer>,
+    buffers: Sender<ByteBuffer>,
     byte_offset: AtomicU64,
     segment_specs: Mutex<Vec<SegmentSpec>>,
 }
 
 impl BufferedSegmentSink {
-    pub fn new(send: kanal::AsyncSender<ByteBuffer>, byte_offset: u64) -> Self {
+    pub fn new(send: Sender<ByteBuffer>, byte_offset: u64) -> Self {
         Self {
             buffers: send,
             byte_offset: AtomicU64::new(byte_offset),
@@ -90,10 +91,10 @@ impl SegmentSink for BufferedSegmentSink {
         };
 
         if let Some(padding) = padding_buffer {
-            let _ = self.buffers.send(padding).await;
+            drop(self.buffers.send(padding).await);
         }
         for buffer in buffers {
-            let _ = self.buffers.send(buffer).await;
+            drop(self.buffers.send(buffer).await);
         }
 
         Ok(segment_id)

vortex-file/src/segments/writer.rs

@@ -7,6 +7,9 @@ use std::sync::Arc;
 use std::sync::atomic::AtomicU64;
 use std::sync::atomic::Ordering;
 
+use async_channel::SendError;
+use async_channel::Sender;
+use async_channel::bounded;
 use futures::FutureExt;
 use futures::StreamExt;
 use futures::TryStreamExt;
@@ -40,7 +43,6 @@ use vortex_error::vortex_bail;
 use vortex_error::vortex_err;
 use vortex_io::IoBuf;
 use vortex_io::VortexWrite;
-use vortex_io::kanal_ext::KanalExt;
 use vortex_io::runtime::BlockingRuntime;
 use vortex_io::session::RuntimeSessionExt;
 use vortex_layout::BufferedBytesTracker;
@@ -332,7 +334,7 @@ impl VortexWriteOptions {
         let mut position = MAGIC_BYTES.len() as u64;
 
         // Create a channel to send buffers from the segment sink to the output stream.
-        let (send, recv) = kanal::bounded_async(1);
+        let (send, recv_stream) = bounded(1);
 
         let segments = Arc::new(BufferedSegmentSink::new(send, position));
 
@@ -354,7 +356,6 @@ impl VortexWriteOptions {
         });
 
         // Flush buffers as they arrive
-        let recv_stream = recv.into_stream();
         pin_mut!(recv_stream);
         while let Some(buffer) = recv_stream.next().await {
             if buffer.is_empty() {
@@ -420,10 +421,9 @@ impl VortexWriteOptions {
     /// flush remaining buffers, and receive the [`WriteSummary`].
     pub fn writer<'w, W: VortexWrite + Unpin + 'w>(self, write: W, dtype: DType) -> Writer<'w> {
         // Create a channel for sending arrays to the layout task.
-        let (arrays_send, arrays_recv) = kanal::bounded_async(1);
+        let (arrays_send, arrays_recv) = bounded(1);
 
-        let arrays =
-            ArrayStreamExt::boxed(ArrayStreamAdapter::new(dtype, arrays_recv.into_stream()));
+        let arrays = ArrayStreamExt::boxed(ArrayStreamAdapter::new(dtype, arrays_recv));
 
         let write = CountingVortexWrite::new(write);
         let bytes_written = write.counter();
@@ -571,7 +571,7 @@ fn validate_metadata_segments(metadata: &HashMap<String, ByteBuffer>) -> VortexR
 /// An async API for writing Vortex files.
 pub struct Writer<'w> {
     // The input channel for sending arrays to the writer.
-    arrays: Option<kanal::AsyncSender<VortexResult<ArrayRef>>>,
+    arrays: Option<Sender<VortexResult<ArrayRef>>>,
     // The writer task that ultimately produces the footer.
     future: Fuse<LocalBoxFuture<'w, VortexResult<WriteSummary>>>,
     // The bytes written so far.
@@ -620,7 +620,7 @@ impl Writer<'_> {
             while let Some(chunk) = stream.next().await {
                 arrays.send(chunk).await?;
             }
-            Ok::<_, kanal::SendError>(())
+            Ok::<_, SendError<VortexResult<ArrayRef>>>(())
         }
         .fuse();
         pin_mut!(stream_fut);

vortex-file/src/writer.rs

@@ -16,6 +16,7 @@ version = { workspace = true }
 all-features = true
 
 [dependencies]
+async-channel = { workspace = true }
 async-executor = { workspace = true }
 async-fs = { workspace = true }
 async-stream = { workspace = true }
@@ -24,7 +25,6 @@ bytes = { workspace = true }
 crossbeam-channel = { workspace = true }
 futures = { workspace = true, features = ["std", "executor"] }
 glob = { workspace = true }
-kanal = { workspace = true }
 object_store = { workspace = true, optional = true, features = ["fs"] }
 parking_lot = { workspace = true }
 pin-project-lite = { workspace = true }

vortex-io/Cargo.toml

@@ -18,7 +18,6 @@ pub use write::*;
 pub mod compat;
 pub mod filesystem;
 mod io_buf;
-pub mod kanal_ext;
 mod limit;
 #[cfg(feature = "object_store")]
 pub mod object_store;