fix: flume bug (#2298)

fix: flume

src/mito2/Cargo.toml

@@ -13,6 +13,7 @@ anymap = "1.0.0-beta.2"
 api.workspace = true
 aquamarine = "0.3"
 arc-swap = "1.0"
+async-channel = "1.9"
 async-compat = "0.2"
 async-stream.workspace = true
 async-trait = "0.1"
@@ -33,7 +34,6 @@ dashmap = "5.4"
 datafusion-common.workspace = true
 datafusion.workspace = true
 datatypes = { workspace = true }
-flume = "0.10"
 futures.workspace = true
 lazy_static = "1.4"
 log-store = { workspace = true }

src/mito2/src/error.rs

@@ -367,8 +367,8 @@ pub enum Error {
         location: Location,
     },
 
-    #[snafu(display("Invalid flume sender, location: {}", location,))]
-    InvalidFlumeSender { location: Location },
+    #[snafu(display("Invalid sender, location: {}", location,))]
+    InvalidSender { location: Location },
 
     #[snafu(display("Invalid scheduler state, location: {}", location))]
     InvalidSchedulerState { location: Location },
@@ -452,7 +452,7 @@ impl ErrorExt for Error {
             ComputeArrow { .. } => StatusCode::Internal,
             ComputeVector { .. } => StatusCode::Internal,
             PrimaryKeyLengthMismatch { .. } => StatusCode::InvalidArguments,
-            InvalidFlumeSender { .. } => StatusCode::InvalidArguments,
+            InvalidSender { .. } => StatusCode::InvalidArguments,
             InvalidSchedulerState { .. } => StatusCode::InvalidArguments,
             StopScheduler { .. } => StatusCode::Internal,
             BuildPredicate { source, .. } => source.status_code(),

src/mito2/src/schedule/scheduler.rs

@@ -23,9 +23,7 @@ use tokio::sync::Mutex;
 use tokio::task::JoinHandle;
 use tokio_util::sync::CancellationToken;
 
-use crate::error::{
-    InvalidFlumeSenderSnafu, InvalidSchedulerStateSnafu, Result, StopSchedulerSnafu,
-};
+use crate::error::{InvalidSchedulerStateSnafu, InvalidSenderSnafu, Result, StopSchedulerSnafu};
 
 pub type Job = Pin<Box<dyn Future<Output = ()> + Send>>;
 
@@ -49,7 +47,7 @@ pub type SchedulerRef = Arc<dyn Scheduler>;
 /// Request scheduler based on local state.
 pub struct LocalScheduler {
     /// Sends jobs to flume bounded channel
-    sender: RwLock<Option<flume::Sender<Job>>>,
+    sender: RwLock<Option<async_channel::Sender<Job>>>,
     /// Task handles
     handles: Mutex<Vec<JoinHandle<()>>>,
     /// Token used to halt the scheduler
@@ -63,7 +61,7 @@ impl LocalScheduler {
     ///
     /// concurrency: the number of bounded receiver
     pub fn new(concurrency: usize) -> Self {
-        let (tx, rx) = flume::unbounded();
+        let (tx, rx) = async_channel::unbounded();
         let token = CancellationToken::new();
         let state = Arc::new(AtomicU8::new(STATE_RUNNING));
 
@@ -79,7 +77,7 @@ impl LocalScheduler {
                         _ = child.cancelled() => {
                             break;
                         }
-                        req_opt = receiver.recv_async() =>{
+                        req_opt = receiver.recv() =>{
                             if let Ok(job) = req_opt {
                                 job.await;
                             }
@@ -89,7 +87,7 @@ impl LocalScheduler {
                 // When task scheduler is cancelled, we will wait all task finished
                 if state_clone.load(Ordering::Relaxed) == STATE_AWAIT_TERMINATION {
                     // recv_async waits until all sender's been dropped.
-                    while let Ok(job) = receiver.recv_async().await {
+                    while let Ok(job) = receiver.recv().await {
                         job.await;
                     }
                     state_clone.store(STATE_STOP, Ordering::Relaxed);
@@ -119,13 +117,14 @@ impl Scheduler for LocalScheduler {
             self.state.load(Ordering::Relaxed) == STATE_RUNNING,
             InvalidSchedulerStateSnafu
         );
+
         self.sender
             .read()
             .unwrap()
             .as_ref()
             .context(InvalidSchedulerStateSnafu)?
-            .send(job)
-            .map_err(|_| InvalidFlumeSenderSnafu {}.build())
+            .try_send(job)
+            .map_err(|_| InvalidSenderSnafu {}.build())
     }
 
     /// if await_termination is true, scheduler will wait all tasks finished before stopping