fix: fix spill flaky tests (#13294)

Author: xudong963

src/query/service/src/pipelines/processors/transforms/hash_join/build_spill/build_spill_coordinator.rs

@@ -26,6 +26,7 @@ use common_exception::Result;
 use common_expression::DataBlock;
 use log::info;
 use parking_lot::Mutex;
+use parking_lot::RwLock;
 
 /// Coordinate all hash join build processors to spill.
 /// It's shared by all hash join build processors.
@@ -41,7 +42,7 @@ pub struct BuildSpillCoordinator {
     /// Spill tasks, the size is the same as the total active processor count.
     pub(crate) spill_tasks: Mutex<VecDeque<Vec<(u8, DataBlock)>>>,
     /// When a build processor won't trigger spill, the field will plus one
-    pub(crate) non_spill_processors: AtomicUsize,
+    pub(crate) non_spill_processors: RwLock<usize>,
     /// If there is the last active processor, send true to watcher channel
     pub(crate) ready_spill_watcher: Sender<bool>,
     pub(crate) dummy_ready_spill_receiver: Receiver<bool>,
@@ -74,15 +75,18 @@ impl BuildSpillCoordinator {
                 .send(false)
                 .map_err(|_| ErrorCode::TokioError("ready_spill_watcher channel is closed"))?;
         }
-        let old_val = self.waiting_spill_count.fetch_add(1, Ordering::Relaxed);
+        let non_spill_processors = self.non_spill_processors.read();
+        let old_val = self.waiting_spill_count.fetch_add(1, Ordering::Release);
         let waiting_spill_count = old_val + 1;
-        let non_spill_processors = self.non_spill_processors.load(Ordering::Relaxed);
         info!(
             "waiting_spill_count: {:?}, non_spill_processors: {:?}, total_builder_count: {:?}",
-            waiting_spill_count, non_spill_processors, self.total_builder_count
+            waiting_spill_count, *non_spill_processors, self.total_builder_count
         );
 
-        if waiting_spill_count + non_spill_processors == self.total_builder_count {
+        if (waiting_spill_count + *non_spill_processors == self.total_builder_count)
+            && self.get_need_spill()
+        {
+            self.no_need_spill();
             // Reset waiting_spill_count
             self.waiting_spill_count.store(0, Ordering::Relaxed);
             // No need to wait spill, the processor is the last one
@@ -116,13 +120,6 @@ impl BuildSpillCoordinator {
 
     // Get active processor count
     pub fn active_processor_num(&self) -> usize {
-        self.total_builder_count - self.non_spill_processors.load(Ordering::Relaxed)
-    }
-
-    // Add one to `non_spill_processors`
-    // Return value after adding
-    pub fn increase_non_spill_processors(&self) -> usize {
-        let old = self.non_spill_processors.fetch_add(1, Ordering::Relaxed);
-        old + 1
+        self.total_builder_count - *self.non_spill_processors.read()
     }
 }

src/query/service/src/pipelines/processors/transforms/hash_join/hash_join_build_state.rs

@@ -258,7 +258,7 @@ impl HashJoinBuildState {
                     .store(true, Ordering::Relaxed);
                 self.hash_join_state
                     .build_done_watcher
-                    .send(self.send_val.load(Ordering::Relaxed))
+                    .send(self.send_val.load(Ordering::Acquire))
                     .map_err(|_| ErrorCode::TokioError("build_done_watcher channel is closed"))?;
                 return Ok(());
             }
@@ -615,7 +615,7 @@ impl HashJoinBuildState {
             }
             self.hash_join_state
                 .build_done_watcher
-                .send(self.send_val.load(Ordering::Relaxed))
+                .send(self.send_val.load(Ordering::Acquire))
                 .map_err(|_| ErrorCode::TokioError("build_done_watcher channel is closed"))?;
         }
         Ok(())

src/query/service/src/pipelines/processors/transforms/hash_join/transform_hash_join_build.rs

@@ -92,8 +92,6 @@ impl TransformHashJoinBuild {
         if wait {
             self.step = HashJoinBuildStep::WaitSpill;
         } else {
-            // Make `need_spill` to false for `SpillCoordinator`
-            spill_state.spill_coordinator.no_need_spill();
             // Before notify all processors to spill, we need to collect all buffered data in `RowSpace` and `Chunks`
             // Partition all rows and stat how many partitions and rows in each partition.
             // Then choose the largest partitions(which contain rows that can avoid oom exactly) to spill.
@@ -116,12 +114,17 @@ impl TransformHashJoinBuild {
 
     // Called after processor read spilled data
     // It means next round build will start, need to reset some variables.
-    fn reset(&mut self) -> Result<()> {
+    async fn reset(&mut self) -> Result<()> {
         self.finalize_finished = false;
         self.from_spill = true;
         // Only need to reset the following variables once
-        if self.build_state.row_space_builders.load(Ordering::Relaxed) == 0 {
-            self.build_state.send_val.store(2, Ordering::Relaxed);
+        if self
+            .build_state
+            .row_space_builders
+            .fetch_add(1, Ordering::Acquire)
+            == 0
+        {
+            self.build_state.send_val.store(2, Ordering::Release);
             // Before build processors into `WaitProbe` state, set the channel message to false.
             // Then after all probe processors are ready, the last one will send true to channel and wake up all build processors.
             self.build_state
@@ -130,16 +133,14 @@ impl TransformHashJoinBuild {
                 .send(false)
                 .map_err(|_| ErrorCode::TokioError("continue_build_watcher channel is closed"))?;
             let worker_num = self.build_state.build_worker_num.load(Ordering::Relaxed) as usize;
-            self.build_state
-                .row_space_builders
-                .store(worker_num, Ordering::Relaxed);
             self.build_state
                 .hash_join_state
                 .hash_table_builders
                 .store(worker_num, Ordering::Relaxed);
             self.build_state.hash_join_state.reset();
         }
         self.step = HashJoinBuildStep::Running;
+        self.build_state.restore_barrier.wait().await;
         Ok(())
     }
 }
@@ -166,11 +167,13 @@ impl Processor for TransformHashJoinBuild {
                         // The processor won't be triggered spill, because there won't be data from input port
                         // Add the processor to `non_spill_processors`
                         let spill_coordinator = &spill_state.spill_coordinator;
-                        let waiting_spill_count = spill_coordinator.waiting_spill_count.load(Ordering::Relaxed);
-                        let non_spill_processors = spill_coordinator.increase_non_spill_processors();
-                        info!("waiting_spill_count: {:?}, non_spill_processors: {:?}, total_builder_count: {:?}", waiting_spill_count, non_spill_processors, spill_state.spill_coordinator.total_builder_count);
-                        if waiting_spill_count != 0 && non_spill_processors + waiting_spill_count == spill_state.spill_coordinator.total_builder_count {
+                        let mut non_spill_processors = spill_coordinator.non_spill_processors.write();
+                        *non_spill_processors += 1;
+                        let waiting_spill_count = spill_coordinator.waiting_spill_count.load(Ordering::Acquire);
+                        info!("waiting_spill_count: {:?}, non_spill_processors: {:?}, total_builder_count: {:?}", waiting_spill_count, *non_spill_processors, spill_state.spill_coordinator.total_builder_count);
+                        if (waiting_spill_count != 0 && *non_spill_processors + waiting_spill_count == spill_state.spill_coordinator.total_builder_count) && spill_coordinator.get_need_spill() {
                             spill_coordinator.no_need_spill();
+                            drop(non_spill_processors);
                             let mut spill_task = spill_coordinator.spill_tasks.lock();
                             spill_state.split_spill_tasks(spill_coordinator.active_processor_num(), &mut spill_task)?;
                             spill_coordinator.waiting_spill_count.store(0, Ordering::Relaxed);
@@ -201,15 +204,7 @@ impl Processor for TransformHashJoinBuild {
                 true => {
                     // If join spill is enabled, we should wait probe to spill.
                     // Then restore data from disk and build hash table, util all spilled data are processed.
-                    if let Some(spill_state) = &mut self.spill_state {
-                        // Send spilled partition to `HashJoinState`, used by probe spill.
-                        // The method should be called only once.
-                        if !self.send_partition_set {
-                            self.build_state
-                                .hash_join_state
-                                .set_spilled_partition(&spill_state.spiller.spilled_partition_set);
-                            self.send_partition_set = true;
-                        }
+                    if self.spill_state.is_some() {
                         self.step = HashJoinBuildStep::WaitProbe;
                         Ok(Event::Async)
                     } else {
@@ -273,6 +268,16 @@ impl Processor for TransformHashJoinBuild {
                     self.build_state.finalize(task)
                 } else {
                     self.finalize_finished = true;
+                    if let Some(spill_state) = &mut self.spill_state {
+                        // Send spilled partition to `HashJoinState`, used by probe spill.
+                        // The method should be called only once.
+                        if !self.send_partition_set {
+                            self.build_state
+                                .hash_join_state
+                                .set_spilled_partition(&spill_state.spiller.spilled_partition_set);
+                            self.send_partition_set = true;
+                        }
+                    }
                     self.build_state.build_done()
                 }
             }
@@ -364,7 +369,7 @@ impl Processor for TransformHashJoinBuild {
                     self.input_data = Some(DataBlock::concat(&spilled_data)?);
                 }
                 self.build_state.restore_barrier.wait().await;
-                self.reset()?;
+                self.reset().await?;
             }
             _ => unreachable!(),
         }

src/query/service/src/pipelines/processors/transforms/hash_join/transform_hash_join_probe.rs

@@ -223,7 +223,7 @@ impl TransformHashJoinProbe {
         Ok(Event::NeedData)
     }
 
-    fn reset(&mut self) -> Result<()> {
+    async fn reset(&mut self) -> Result<()> {
         self.step = HashJoinProbeStep::Running;
         // self.probe_state.reset();
         if (self.join_probe_state.hash_join_state.need_outer_scan()
@@ -238,7 +238,7 @@ impl TransformHashJoinProbe {
         if self
             .join_probe_state
             .final_probe_workers
-            .load(Ordering::Relaxed)
+            .fetch_add(1, Ordering::Acquire)
             == 0
         {
             // Before probe processor into `WaitBuild` state, send `1` to channel
@@ -248,11 +248,9 @@ impl TransformHashJoinProbe {
                 .build_done_watcher
                 .send(1)
                 .map_err(|_| ErrorCode::TokioError("build_done_watcher channel is closed"))?;
-            self.join_probe_state
-                .final_probe_workers
-                .store(self.join_probe_state.processor_count, Ordering::Relaxed);
         }
         self.outer_scan_finished = false;
+        self.join_probe_state.restore_barrier.wait().await;
         Ok(())
     }
 }
@@ -516,7 +514,7 @@ impl Processor for TransformHashJoinProbe {
                     }
                 }
                 self.join_probe_state.restore_barrier.wait().await;
-                self.reset()?;
+                self.reset().await?;
             }
             HashJoinProbeStep::FinalScan | HashJoinProbeStep::FastReturn => unreachable!(),
         };

tests/sqllogictests/suites/query/spill.test

@@ -4,7 +4,7 @@ statement ok
 set disable_join_reorder = 1;
 
 statement ok
-set join_spilling_threshold = 0;
+set join_spilling_threshold = 1;
 
 statement ok
 create table t3(a int);
@@ -65,7 +65,7 @@ statement ok
 create table t3 as select number as a from numbers(1000000);
 
 statement ok
-set join_spilling_threshold = 0;
+set join_spilling_threshold = 100;
 
 query I
 select count() from t3 inner join numbers(1000000) on t3.a = number;
@@ -74,7 +74,7 @@ select count() from t3 inner join numbers(1000000) on t3.a = number;
 
 onlyif mysql
 statement ok
-set join_spilling_threshold = 0;
+set join_spilling_threshold = 1024 * 1024 * 1;
 
 onlyif mysql
 query I