From d2e3e984f6ea54c63b4b131813f7bc64a76836bf Mon Sep 17 00:00:00 2001
From: Anton Borisov <anton.borisov@fresha.com>
Date: Tue, 21 Jul 2026 02:32:13 +0100
Subject: [PATCH] [rust] Fix lost wakeup in LogFetchBuffer::await_not_empty

---
 .../src/client/table/log_fetch_buffer.rs      | 44 ++++++++++++++-----
 .../crates/fluss/tests/integration/utils.rs   |  1 -
 2 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/fluss-rust/crates/fluss/src/client/table/log_fetch_buffer.rs b/fluss-rust/crates/fluss/src/client/table/log_fetch_buffer.rs
index 794fe9df92..51354f5409 100644
--- a/fluss-rust/crates/fluss/src/client/table/log_fetch_buffer.rs
+++ b/fluss-rust/crates/fluss/src/client/table/log_fetch_buffer.rs
@@ -139,35 +139,34 @@ impl LogFetchBuffer {
     pub async fn await_not_empty(&self, timeout: Duration) -> Result<bool> {
         let deadline = Instant::now() + timeout;
 
+        // Arm before checking: notify_waiters() stores no permit and only wakes already-armed
+        // waiters, so arming late would drop a notification that races the emptiness check.
+        let notified = self.not_empty_notify.notified();
+        tokio::pin!(notified);
+
         loop {
-            // Check if buffer is not empty
+            notified.as_mut().enable();
+
             if !self.is_empty() {
                 return Ok(true);
             }
 
-            // Check if woken up
             if self.woken_up.swap(false, Ordering::Acquire) {
                 return Err(Error::WakeupError {
                     message: "The await operation was interrupted by wakeup.".to_string(),
                 });
             }
 
-            // Check if timeout
             let now = Instant::now();
             if now >= deadline {
                 return Ok(false);
             }
 
-            // Wait for notification with remaining time
             let remaining = deadline - now;
-            let notified = self.not_empty_notify.notified();
             tokio::select! {
-                _ = tokio::time::sleep(remaining) => {
-                    return Ok(false); // Timeout
-                }
-                _ = notified => {
-                    // Got notification, check again
-                    continue;
+                _ = tokio::time::sleep(remaining) => return Ok(false),
+                _ = notified.as_mut() => {
+                    notified.set(self.not_empty_notify.notified()); // re-arm
                 }
             }
         }
@@ -1064,6 +1063,29 @@ mod tests {
         assert!(completed.take_error().is_some());
     }
 
+    #[tokio::test]
+    async fn await_not_empty_returns_true_when_data_arrives_during_wait() {
+        let buffer = Arc::new(LogFetchBuffer::new(test_resolver().unwrap()));
+        let table_bucket = TableBucket::new(1, 0);
+
+        let producer = Arc::clone(&buffer);
+        let producer_bucket = table_bucket.clone();
+        tokio::spawn(async move {
+            tokio::time::sleep(Duration::from_millis(20)).await;
+            producer.pend(Box::new(ErrorPendingFetch {
+                table_bucket: producer_bucket.clone(),
+            }));
+            producer.try_complete(&producer_bucket);
+        });
+
+        // Timeout far exceeds the producer delay: must return as soon as data lands, and must
+        // not strand until the deadline even if the notification races with the internal check.
+        let start = Instant::now();
+        let result = buffer.await_not_empty(Duration::from_secs(5)).await;
+        assert!(matches!(result, Ok(true)));
+        assert!(start.elapsed() < Duration::from_secs(1));
+    }
+
     #[test]
     fn default_completed_fetch_reads_records() -> Result<()> {
         let row_type = RowType::new(vec![
diff --git a/fluss-rust/crates/fluss/tests/integration/utils.rs b/fluss-rust/crates/fluss/tests/integration/utils.rs
index 58d7862901..2ad0e6c368 100644
--- a/fluss-rust/crates/fluss/tests/integration/utils.rs
+++ b/fluss-rust/crates/fluss/tests/integration/utils.rs
@@ -230,7 +230,6 @@ pub fn make_int_array(values: &[Option<i32>]) -> FlussArray {
 // ─── Poll helpers ────────────────────────────────────────────────────────────
 
 pub const DEFAULT_POLL_TIMEOUT: Duration = Duration::from_secs(10);
-pub const DEFAULT_POLL_INTERVAL: Duration = Duration::from_millis(500);
 
 /// Polls repeatedly, accumulating items until `expected_count` is reached or `timeout` elapses.
 pub async fn poll_until_count<T>(