From b68a8aceb9a3502d654c64bbc5836901e818cd6d Mon Sep 17 00:00:00 2001
From: WenyXu <wenymedia@gmail.com>
Date: Wed, 27 Dec 2023 15:20:36 +0000
Subject: [PATCH] chore: apply suggestions from CR

---
 src/datanode/src/heartbeat/task_tracker.rs | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/datanode/src/heartbeat/task_tracker.rs b/src/datanode/src/heartbeat/task_tracker.rs
index c78efa1a1450..626754722204 100644
--- a/src/datanode/src/heartbeat/task_tracker.rs
+++ b/src/datanode/src/heartbeat/task_tracker.rs
@@ -249,12 +249,13 @@ mod tests {
     async fn test_async_task_tracker_wait_timeout() {
         let tracker = TaskTracker::<TestResult>::new();
         let region_id = RegionId::new(1024, 1);
+        let (tx, rx) = oneshot::channel::<()>();
 
         let result = tracker
             .try_register(
                 region_id,
                 Box::pin(async move {
-                    tokio::time::sleep(Duration::from_millis(200)).await;
+                    let _ = rx.await;
                     Ok(TestResult { value: 1024 })
                 }),
             )
@@ -264,7 +265,8 @@ mod tests {
         let result = tracker.wait(&mut watcher, Duration::from_millis(100)).await;
         assert!(result.is_timeout());
 
-        tokio::time::sleep(Duration::from_millis(150)).await;
+        // Triggers first future return.
+        tx.send(()).unwrap();
         let result = tracker
             .wait(&mut watcher, Duration::from_millis(100))
             .await