From 5ab074709206c539657d91ac8766065addd3b084 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Wed, 21 Jun 2023 19:04:34 +0900 Subject: [PATCH] test(storage): wait task before checking scheduled task num (#1811) --- src/storage/src/scheduler.rs | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/src/storage/src/scheduler.rs b/src/storage/src/scheduler.rs index 2c8d1c162a..4f6d1689cd 100644 --- a/src/storage/src/scheduler.rs +++ b/src/storage/src/scheduler.rs @@ -312,7 +312,7 @@ where #[cfg(test)] mod tests { - use std::sync::atomic::AtomicI32; + use std::sync::atomic::{AtomicBool, AtomicI32}; use std::time::Duration; use store_api::storage::RegionId; @@ -564,7 +564,9 @@ mod tests { let task_scheduled = Arc::new(AtomicI32::new(0)); let task_scheduled_cloned = task_scheduled.clone(); - common_runtime::spawn_write(async move { + let scheduling = Arc::new(AtomicBool::new(true)); + let scheduling_clone = scheduling.clone(); + let handle = common_runtime::spawn_write(async move { for i in 0..10000 { if let Ok(res) = scheduler_cloned.schedule(MockRequest { region_id: i as RegionId, @@ -573,12 +575,19 @@ mod tests { task_scheduled_cloned.fetch_add(1, Ordering::Relaxed); } } + + if !scheduling_clone.load(Ordering::Relaxed) { + break; + } } }); - tokio::time::sleep(Duration::from_millis(1)).await; scheduler.stop(true).await.unwrap(); + scheduling.store(false, Ordering::Relaxed); + let finished = finished.load(Ordering::Relaxed); + handle.await.unwrap(); + assert_eq!(finished, task_scheduled.load(Ordering::Relaxed)); } }