Harden TSO waker edge cases and refresh benchmarks

Author: Michael Ingley

doc/tso_waker_criterion.md

@@ -1,6 +1,6 @@
 # TSO Waker Criterion Benchmark
 
-Date: 2026-02-06
+Date: 2026-02-09
 Repo: `tikv/client-rust`
 Branch: `mingley/tso-waker-criterion`
 Host: macOS 26.2 (Darwin 25.2.0), Apple M4 Pro, arm64
@@ -46,15 +46,15 @@ benefit of skipping redundant wake calls under this simulation pattern.
 
 From Criterion output (`time` line):
 
-- `tso_waker_policy/response/old`: `[3.2519 ns 3.2712 ns 3.2926 ns]`
-- `tso_waker_policy/response/new`: `[763.41 ps 766.39 ps 769.43 ps]`
+- `tso_waker_policy/response/old`: `[3.4869 ns 3.4909 ns 3.4954 ns]`
+- `tso_waker_policy/response/new`: `[797.61 ps 798.59 ps 799.71 ps]`
 
-- `tso_waker_policy/register/old`: `[2.3768 ns 2.3819 ns 2.3874 ns]`
-- `tso_waker_policy/register/new`: `[286.76 ps 287.51 ps 288.27 ps]`
+- `tso_waker_policy/register/old`: `[2.3964 ns 2.4009 ns 2.4058 ns]`
+- `tso_waker_policy/register/new`: `[290.92 ps 291.35 ps 291.85 ps]`
 
 Median-based speedups:
-- response path: `3.2712 ns / 0.76639 ns = 4.27x`
-- registration path: `2.3819 ns / 0.28751 ns = 8.28x`
+- response path: `3.4909 ns / 0.79859 ns = 4.37x`
+- registration path: `2.4009 ns / 0.29135 ns = 8.24x`
 
 ## Interpretation
 

src/pd/timestamp.rs

@@ -13,6 +13,8 @@
 
 use std::collections::VecDeque;
 use std::pin::Pin;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
 use std::sync::Arc;
 
 use futures::pin_mut;
@@ -21,7 +23,6 @@ use futures::task::AtomicWaker;
 use futures::task::Context;
 use futures::task::Poll;
 use log::debug;
-use log::info;
 use log::warn;
 use pin_project::pin_project;
 use tokio::sync::mpsc;
@@ -93,12 +94,16 @@ async fn run_tso(
     // more requests from the bounded channel. This waker is used to wake up the sending future
     // if the queue containing pending requests is no longer full.
     let sending_future_waker = Arc::new(AtomicWaker::new());
+    // This flag indicates the sender stream could not acquire `pending_requests` lock in poll
+    // and needs an explicit wake from the response path.
+    let sender_waiting_on_lock = Arc::new(AtomicBool::new(false));
 
     let request_stream = TsoRequestStream {
         cluster_id,
         request_rx,
         pending_requests: pending_requests.clone(),
         self_waker: sending_future_waker.clone(),
+        sender_waiting_on_lock: sender_waiting_on_lock.clone(),
     };
 
     // let send_requests = rpc_sender.send_all(&mut request_stream);
@@ -112,15 +117,24 @@ async fn run_tso(
             allocate_timestamps(&resp, &mut pending_requests)?;
             was_full && pending_requests.len() < MAX_PENDING_COUNT
         };
+        let sender_blocked_by_lock = sender_waiting_on_lock.swap(false, Ordering::AcqRel);
 
-        // Only wake sender when a previously full queue gains capacity.
-        if should_wake_sender {
+        // Wake sender when:
+        // 1. a previously full queue gains capacity, or
+        // 2. sender was blocked on `pending_requests` mutex contention.
+        if should_wake_sender || sender_blocked_by_lock {
             sending_future_waker.wake();
         }
     }
-    // TODO: distinguish between unexpected stream termination and expected end of test
-    info!("TSO stream terminated");
-    Ok(())
+    let pending_count = pending_requests.lock().await.len();
+    if pending_count == 0 {
+        Ok(())
+    } else {
+        Err(internal_err!(
+            "TSO stream terminated with {} pending requests",
+            pending_count
+        ))
+    }
 }
 
 struct RequestGroup {
@@ -135,6 +149,7 @@ struct TsoRequestStream {
     request_rx: mpsc::Receiver<oneshot::Sender<Timestamp>>,
     pending_requests: Arc<Mutex<VecDeque<RequestGroup>>>,
     self_waker: Arc<AtomicWaker>,
+    sender_waiting_on_lock: Arc<AtomicBool>,
 }
 
 impl Stream for TsoRequestStream {
@@ -147,8 +162,18 @@ impl Stream for TsoRequestStream {
         pin_mut!(pending_requests);
         let mut pending_requests = if let Poll::Ready(pending_requests) = pending_requests.poll(cx)
         {
+            this.sender_waiting_on_lock.store(false, Ordering::Release);
             pending_requests
         } else {
+            // The lock future is dropped at the end of this poll, so record the lock-wait state
+            // and rely on the response path to issue a wake after it releases the lock.
+            this.sender_waiting_on_lock.store(true, Ordering::Release);
+            this.self_waker.register(cx.waker());
+            // If response path consumed the wait flag before registration, its wake might be lost.
+            // Trigger one local wake to guarantee another poll.
+            if !this.sender_waiting_on_lock.load(Ordering::Acquire) {
+                cx.waker().wake_by_ref();
+            }
             return Poll::Pending;
         };
         let mut requests = Vec::new();
@@ -230,3 +255,214 @@ fn allocate_timestamps(
     };
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::atomic::AtomicUsize;
+    use std::sync::Arc;
+
+    use futures::executor::block_on;
+    use futures::task::noop_waker_ref;
+    use futures::task::waker;
+    use futures::task::ArcWake;
+
+    use super::*;
+
+    struct WakeCounter {
+        wakes: AtomicUsize,
+    }
+
+    impl ArcWake for WakeCounter {
+        fn wake_by_ref(arc_self: &Arc<Self>) {
+            arc_self.wakes.fetch_add(1, Ordering::SeqCst);
+        }
+    }
+
+    fn test_tso_request(count: u32) -> TsoRequest {
+        TsoRequest {
+            header: Some(RequestHeader {
+                cluster_id: 1,
+                sender_id: 0,
+            }),
+            count,
+            dc_location: String::new(),
+        }
+    }
+
+    fn test_tso_response(count: u32, logical: i64) -> TsoResponse {
+        TsoResponse {
+            header: None,
+            count,
+            timestamp: Some(Timestamp {
+                physical: 123,
+                logical,
+                suffix_bits: 0,
+            }),
+        }
+    }
+
+    type TestStreamContext = (
+        TsoRequestStream,
+        mpsc::Sender<TimestampRequest>,
+        Arc<Mutex<VecDeque<RequestGroup>>>,
+        Arc<AtomicWaker>,
+        Arc<AtomicBool>,
+    );
+
+    fn new_test_stream() -> TestStreamContext {
+        let (request_tx, request_rx) = mpsc::channel(MAX_BATCH_SIZE);
+        let pending_requests = Arc::new(Mutex::new(VecDeque::new()));
+        let self_waker = Arc::new(AtomicWaker::new());
+        let sender_waiting_on_lock = Arc::new(AtomicBool::new(false));
+        let stream = TsoRequestStream {
+            cluster_id: 1,
+            request_rx,
+            pending_requests: pending_requests.clone(),
+            self_waker: self_waker.clone(),
+            sender_waiting_on_lock: sender_waiting_on_lock.clone(),
+        };
+        (
+            stream,
+            request_tx,
+            pending_requests,
+            self_waker,
+            sender_waiting_on_lock,
+        )
+    }
+
+    #[test]
+    fn allocate_timestamps_successfully_assigns_monotonic_timestamps() {
+        let (tx1, rx1) = oneshot::channel();
+        let (tx2, rx2) = oneshot::channel();
+        let (tx3, rx3) = oneshot::channel();
+        let mut pending_requests = VecDeque::new();
+        pending_requests.push_back(RequestGroup {
+            tso_request: test_tso_request(3),
+            requests: vec![tx1, tx2, tx3],
+        });
+
+        allocate_timestamps(&test_tso_response(3, 100), &mut pending_requests).unwrap();
+        assert!(pending_requests.is_empty());
+
+        let ts1 = block_on(rx1).unwrap();
+        let ts2 = block_on(rx2).unwrap();
+        let ts3 = block_on(rx3).unwrap();
+        assert_eq!(ts1.logical, 98);
+        assert_eq!(ts2.logical, 99);
+        assert_eq!(ts3.logical, 100);
+    }
+
+    #[test]
+    fn allocate_timestamps_errors_without_timestamp() {
+        let (tx, _rx) = oneshot::channel();
+        let mut pending_requests = VecDeque::new();
+        pending_requests.push_back(RequestGroup {
+            tso_request: test_tso_request(1),
+            requests: vec![tx],
+        });
+        let resp = TsoResponse {
+            header: None,
+            count: 1,
+            timestamp: None,
+        };
+
+        let err = allocate_timestamps(&resp, &mut pending_requests).unwrap_err();
+        assert!(format!("{err:?}").contains("No timestamp in TsoResponse"));
+    }
+
+    #[test]
+    fn allocate_timestamps_errors_when_count_mismatches() {
+        let (tx, _rx) = oneshot::channel();
+        let mut pending_requests = VecDeque::new();
+        pending_requests.push_back(RequestGroup {
+            tso_request: test_tso_request(2),
+            requests: vec![tx],
+        });
+
+        let err =
+            allocate_timestamps(&test_tso_response(1, 10), &mut pending_requests).unwrap_err();
+        assert!(format!("{err:?}").contains("different number of timestamps"));
+    }
+
+    #[test]
+    fn allocate_timestamps_errors_on_extra_response() {
+        let mut pending_requests = VecDeque::new();
+        let err =
+            allocate_timestamps(&test_tso_response(1, 10), &mut pending_requests).unwrap_err();
+        assert!(format!("{err:?}").contains("more TsoResponse than expected"));
+    }
+
+    #[test]
+    fn poll_next_emits_request_and_enqueues_request_group() {
+        let (stream, request_tx, pending_requests, _self_waker, sender_waiting_on_lock) =
+            new_test_stream();
+        let (tx, _rx) = oneshot::channel();
+        request_tx.try_send(tx).unwrap();
+
+        let mut stream = Box::pin(stream);
+        let mut cx = Context::from_waker(noop_waker_ref());
+        let polled = stream.as_mut().poll_next(&mut cx);
+        let req = match polled {
+            Poll::Ready(Some(req)) => req,
+            other => panic!("expected Poll::Ready(Some(_)), got {:?}", other),
+        };
+
+        assert_eq!(req.count, 1);
+        assert!(!sender_waiting_on_lock.load(Ordering::SeqCst));
+        let queued = block_on(async { pending_requests.lock().await.len() });
+        assert_eq!(queued, 1);
+    }
+
+    #[test]
+    fn poll_next_registers_self_waker_when_pending_queue_is_full() {
+        let (stream, _request_tx, pending_requests, self_waker, _sender_waiting_on_lock) =
+            new_test_stream();
+        block_on(async {
+            let mut guard = pending_requests.lock().await;
+            for _ in 0..MAX_PENDING_COUNT {
+                guard.push_back(RequestGroup {
+                    tso_request: test_tso_request(0),
+                    requests: Vec::new(),
+                });
+            }
+        });
+
+        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);
+
+        let polled = stream.as_mut().poll_next(&mut cx);
+        assert!(matches!(polled, Poll::Pending));
+        assert_eq!(wake_counter.wakes.load(Ordering::SeqCst), 0);
+
+        self_waker.wake();
+        assert_eq!(wake_counter.wakes.load(Ordering::SeqCst), 1);
+    }
+
+    #[test]
+    fn poll_next_marks_waiting_flag_when_lock_is_contended() {
+        let (stream, _request_tx, pending_requests, self_waker, sender_waiting_on_lock) =
+            new_test_stream();
+        let lock_guard = block_on(pending_requests.lock());
+
+        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);
+
+        let polled = stream.as_mut().poll_next(&mut cx);
+        assert!(matches!(polled, Poll::Pending));
+        assert!(sender_waiting_on_lock.load(Ordering::SeqCst));
+
+        drop(lock_guard);
+        if sender_waiting_on_lock.swap(false, Ordering::AcqRel) {
+            self_waker.wake();
+        }
+        assert!(wake_counter.wakes.load(Ordering::SeqCst) >= 1);
+    }
+}