pd: fix tso sender lost-wake race

Author: Michael Ingley

src/pd/timestamp.rs

@@ -17,7 +17,6 @@ use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
 use std::sync::Arc;
 
-use futures::pin_mut;
 use futures::prelude::*;
 use futures::task::AtomicWaker;
 use futures::task::Context;
@@ -158,25 +157,22 @@ impl Stream for TsoRequestStream {
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
         let mut this = self.project();
 
-        let pending_requests = this.pending_requests.lock();
-        pin_mut!(pending_requests);
-        let mut pending_requests = if let Poll::Ready(pending_requests) = pending_requests.poll(cx)
-        {
+        let mut pending_requests = if let Ok(pending_requests) = this.pending_requests.try_lock() {
             this.sender_waiting_on_lock.store(false, Ordering::SeqCst);
             pending_requests
         } else {
-            // Lock is held by the response path. Register waker first so any
-            // subsequent wake() targets the correct waker, then advertise that
-            // we are waiting.
-            this.self_waker.register(cx.waker());
-            this.sender_waiting_on_lock.store(true, Ordering::SeqCst);
-            // If the response path cleared the flag between our register and
-            // store, its wake may have targeted a stale waker. Self-wake to
-            // guarantee we get re-polled.
-            if !this.sender_waiting_on_lock.load(Ordering::SeqCst) {
-                cx.waker().wake_by_ref();
+            let pending_requests = register_sender_wait_for_pending_lock(
+                cx,
+                this.self_waker.as_ref(),
+                this.sender_waiting_on_lock.as_ref(),
+                || this.pending_requests.try_lock().ok(),
+            );
+
+            if let Some(pending_requests) = pending_requests {
+                pending_requests
+            } else {
+                return Poll::Pending;
             }
-            return Poll::Pending;
         };
         let mut requests = Vec::new();
 
@@ -219,6 +215,29 @@ impl Stream for TsoRequestStream {
     }
 }
 
+fn register_sender_wait_for_pending_lock<F, G>(
+    cx: &mut Context<'_>,
+    self_waker: &AtomicWaker,
+    sender_waiting_on_lock: &AtomicBool,
+    mut try_lock_after_register: F,
+) -> Option<G>
+where
+    F: FnMut() -> Option<G>,
+{
+    // Register first so a wake from the response path targets the current task.
+    self_waker.register(cx.waker());
+    sender_waiting_on_lock.store(true, Ordering::SeqCst);
+
+    // Retry once after advertising waiting to close the race where the response path
+    // already checked/cleared the flag before this store and therefore will not wake.
+    if let Some(guard) = try_lock_after_register() {
+        sender_waiting_on_lock.store(false, Ordering::SeqCst);
+        Some(guard)
+    } else {
+        None
+    }
+}
+
 fn allocate_timestamps(
     resp: &TsoResponse,
     pending_requests: &mut VecDeque<RequestGroup>,
@@ -445,7 +464,7 @@ mod tests {
     }
 
     #[test]
-    fn poll_next_marks_waiting_flag_when_lock_is_contended() {
+    fn poll_next_marks_waiting_flag_when_lock_is_contended_and_response_wakes() {
         let (stream, _request_tx, pending_requests, self_waker, sender_waiting_on_lock) =
             new_test_stream();
         let lock_guard = block_on(pending_requests.lock());
@@ -469,68 +488,55 @@ mod tests {
         assert!(wake_counter.wakes.load(Ordering::SeqCst) >= 1);
     }
 
-    /// Simulate the race where the response path clears the flag *before*
-    /// poll_next sets it. The self-wake guard must fire.
     #[test]
-    fn poll_next_self_wakes_when_flag_cleared_before_store() {
-        let (stream, _request_tx, pending_requests, _self_waker, sender_waiting_on_lock) =
-            new_test_stream();
-        // Hold lock so poll returns Pending.
-        let lock_guard = block_on(pending_requests.lock());
-
+    fn register_sender_wait_sets_waiting_flag_and_registers_waker_on_retry_failure() {
+        let self_waker = AtomicWaker::new();
+        let sender_waiting_on_lock = AtomicBool::new(false);
         let wake_counter = Arc::new(WakeCounter {
             wakes: AtomicUsize::new(0),
         });
         let test_waker = waker(wake_counter.clone());
         let mut cx = Context::from_waker(&test_waker);
-        let mut stream = Box::pin(stream);
-
-        // Pre-clear the flag (simulates response path racing ahead).
-        sender_waiting_on_lock.store(false, Ordering::SeqCst);
-
-        let polled = stream.as_mut().poll_next(&mut cx);
-        assert!(matches!(polled, Poll::Pending));
 
-        // The flag should have been set to true by poll_next.
-        // Because the flag was not externally cleared *after* the store,
-        // no self-wake is needed — the flag stays true for the response path
-        // to observe normally.
+        let reacquired = register_sender_wait_for_pending_lock(
+            &mut cx,
+            &self_waker,
+            &sender_waiting_on_lock,
+            || None::<()>,
+        );
+        assert!(reacquired.is_none());
         assert!(sender_waiting_on_lock.load(Ordering::SeqCst));
 
-        drop(lock_guard);
+        self_waker.wake();
+        assert_eq!(wake_counter.wakes.load(Ordering::SeqCst), 1);
     }
 
-    /// Verify that after the response path clears the flag *and* we simulate
-    /// that clearing happening between register and store, the sender detects
-    /// it and self-wakes.
     #[test]
-    fn poll_next_detects_flag_cleared_after_store_and_self_wakes() {
-        let (stream, _request_tx, pending_requests, _self_waker, sender_waiting_on_lock) =
-            new_test_stream();
-        let lock_guard = block_on(pending_requests.lock());
-
+    fn register_sender_wait_retries_once_and_clears_waiting_flag_when_lock_reacquires() {
+        let self_waker = AtomicWaker::new();
+        let sender_waiting_on_lock = AtomicBool::new(false);
+        let mut retry_count = 0;
         let wake_counter = Arc::new(WakeCounter {
             wakes: AtomicUsize::new(0),
         });
         let test_waker = waker(wake_counter.clone());
         let mut cx = Context::from_waker(&test_waker);
-        let mut stream = Box::pin(stream);
 
-        // poll_next will: register waker, store(true), then load to re-check.
-        // We can't interleave mid-poll, but we can verify the steady-state:
-        // after poll returns Pending, simulate response clearing the flag
-        // and confirm wake propagation via self_waker.
-        let polled = stream.as_mut().poll_next(&mut cx);
-        assert!(matches!(polled, Poll::Pending));
-        assert!(sender_waiting_on_lock.load(Ordering::SeqCst));
-
-        // Simulate response path: clear the flag (as if between store and load).
-        sender_waiting_on_lock.store(false, Ordering::SeqCst);
-        // The registered waker is current, so waking self_waker delivers correctly.
-        _self_waker.wake();
-        assert_eq!(wake_counter.wakes.load(Ordering::SeqCst), 1);
-
-        drop(lock_guard);
+        // Simulates the lost-wake interleaving: initial lock contention, then lock
+        // becomes available before any response-side wake.
+        let reacquired = register_sender_wait_for_pending_lock(
+            &mut cx,
+            &self_waker,
+            &sender_waiting_on_lock,
+            || {
+                retry_count += 1;
+                Some(())
+            },
+        );
+        assert_eq!(retry_count, 1);
+        assert!(reacquired.is_some());
+        assert!(!sender_waiting_on_lock.load(Ordering::SeqCst));
+        assert_eq!(wake_counter.wakes.load(Ordering::SeqCst), 0);
     }
 
     /// After acquiring the lock, the waiting flag must be cleared.