From 9bf2a1b2079819394fbb7fd5a98fa5bd9c39826d Mon Sep 17 00:00:00 2001 From: "Eduard R." Date: Thu, 6 Aug 2026 00:34:34 +0200 Subject: [PATCH 1/2] =?UTF-8?q?pd:=20bound=20every=20wait=20on=20the=20PD?= =?UTF-8?q?=20path=20=E2=80=94=20a=20frozen=20peer=20must=20not=20hang=20t?= =?UTF-8?q?he=20client?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The TSO path had no deadline, keepalive, or health check anywhere: a request parked in pending_requests waits on responses.next() unbounded. When the stream terminates that is recoverable — the pending oneshots drop, callers error, the retry layer reconnects — but a stream that goes silent WITHOUT terminating (the peer's VM paused by its hypervisor, the process frozen under resource starvation, a middlebox silently discarding the connection) strands every pending and future timestamp request forever. The PD dial had the same defect: a frozen peer's kernel still completes the TCP handshake, then the HTTP/2 preface never arrives and connect() hangs unbounded — Connection:: connect even took a timeout parameter and ignored it (_timeout), and its get_members call carried no bound either. That combination is the mechanism behind the #516 CI hang. The runner VM froze (all three TiKVs logged "monotonic time jumped back" by ~65s; PD measured a 2.6s scheduling gap against its 1.67s election keep-alive, resigned, and closed its TSO allocator). PD recovered ~15s later, but the test's in-flight current_timestamp() hung against the zombie stream until nextest's 600s kill. Reproduced locally with docker pause on PD (SIGSTOP: TCP open, no FIN): the in-flight request waited the full 45s pause, resolving only because unpause terminated the stream; reconnects hung the same way at the dial. Follow client-go's shape — its TSO dispatcher arms a deadline per batch with defaultPDTimeout = 3s and cancels the stream on expiry (pd/client/clients/tso/ dispatcher.go), bounds stream creation via checkStreamTimeout (stream.go), and bounds dials — using the client's configured request timeout (Config::timeout, whose documented contract covers PD requests). Every bound is armed only while something is owed: - The TSO receive loop waits in whole timeout windows while batches are outstanding: a batch pending across one full silent window declares the stream stalled and the worker exits. A batch dispatched mid-window is granted one full window, so teardown happens only after between one and two windows of genuine silence. Idle streams (nothing owed) are never torn down. - Stream creation is guarded by a watchdog with the same window rule, armed only while a batch is pending: PD (a gRPC-go server) sends response headers with its first response, so tso(..) legitimately waits while the client is idle — HTTP/2 lets request batches flow (and become pending) before response headers arrive, so a stalled creation is still detected. - The worker's exit is SIGNALED to callers: get_timestamp selects on a watch channel the worker closes on exit — while ENQUEUEING as well as while awaiting the response, since a dead worker's request channel can remain open with a full buffer (held by the frozen connection's task, which may not run again until the connection thaws) and send would otherwise block forever. - Dials are bounded, HTTP/2 handshake included: SecurityManager grows connect_with_timeout (the existing connect signature is preserved, now bounded by the default request timeout), Connection::connect uses its hitherto-ignored timeout, and its get_members call is bounded like every other request in that path. Stream termination is also made observable instead of silent: a clean end with nothing owed stays info-level; stalls, creation timeouts, termination with batches outstanding, and stream status errors (previously swallowed) surface as warnings with the outstanding count — resolving the old TODO about indistinguishable terminations. Considered and omitted: a timer-based backstop in get_timestamp (client-go has no analogue — ambient contexts cover it there — and the worker-exit signal is event-driven and precise), and gRPC keepalive (connection-level health, client-go applies it to TiKV connections; a follow-up). Verification: 10 unit tests — the receive loop and the creation watchdog under a paused clock (stall detected within [1, 2] windows; a mid-window batch granted its full window; idle streams and idle creation never torn down; termination clean only when nothing is owed; a stream status error surfaced; the happy path), get_timestamp failing fast when the worker dies, and enqueueing observing worker death with a full request buffer. The stall test bites: with the stall branch disabled it runs unbounded. 82 lib tests and make check green. Manually, with docker pause on PD for 45s: before, one silent 44.5s hang ended only by the unpause; after, a 5s-idle client keeps its worker silently, the stall is detected in ~4s, reconnects against the frozen peer fail in 2s each, callers receive bounded, descriptive errors while the freeze lasts, and the client recovers within ~3s of the peer thawing. Ref: https://github.com/tikv/client-rust/issues/516 Signed-off-by: Eduard R. --- Cargo.toml | 2 +- src/common/security.rs | 29 ++- src/config.rs | 2 +- src/pd/cluster.rs | 24 ++- src/pd/timestamp.rs | 424 +++++++++++++++++++++++++++++++++++++++-- src/store/client.rs | 2 +- 6 files changed, 449 insertions(+), 34 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index ca1029f3..6290cd6a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -59,7 +59,7 @@ rstest = "0.18.2" serde_json = "1" serial_test = "0.5.0" tempfile = "=3.14.0" -tokio = { version = "1", features = ["sync", "rt-multi-thread", "macros"] } +tokio = { version = "1", features = ["sync", "rt-multi-thread", "macros", "test-util"] } [[test]] name = "failpoint_tests" diff --git a/src/common/security.rs b/src/common/security.rs index 89d60060..ee6874f7 100644 --- a/src/common/security.rs +++ b/src/common/security.rs @@ -76,10 +76,27 @@ impl SecurityManager { } /// Connect to gRPC server using TLS connection. If TLS is not configured, use normal connection. - pub async fn connect( + /// + /// The dial is bounded by the default request timeout — see + /// [`Self::connect_with_timeout`] to choose the bound. + pub async fn connect(&self, addr: &str, factory: Factory) -> Result + where + Factory: FnOnce(Channel) -> Client, + { + self.connect_with_timeout(addr, crate::config::DEFAULT_REQUEST_TIMEOUT, factory) + .await + } + + /// Like [`Self::connect`], with `connect_timeout` bounding the whole dial, + /// HTTP/2 handshake included. A frozen peer can complete the TCP handshake + /// from its kernel (which keeps running when the process does not) and then + /// never answer the HTTP/2 preface, hanging an unbounded connect until the + /// peer thaws (#516). + pub async fn connect_with_timeout( &self, // env: Arc, addr: &str, + connect_timeout: Duration, factory: Factory, ) -> Result where @@ -91,7 +108,15 @@ impl SecurityManager { } else { self.default_channel(addr).await? }; - let ch = channel.connect().await?; + let ch = tokio::time::timeout(connect_timeout, channel.connect()) + .await + .map_err(|_| { + internal_err!( + "connecting to {} timed out after {:?}", + addr, + connect_timeout + ) + })??; Ok(factory(ch)) } diff --git a/src/config.rs b/src/config.rs index 92a67e65..35ab2e89 100644 --- a/src/config.rs +++ b/src/config.rs @@ -28,7 +28,7 @@ pub struct Config { pub keyspace: Option, } -const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(2); +pub(crate) const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(2); const DEFAULT_GRPC_MAX_DECODING_MESSAGE_SIZE: usize = 4 * 1024 * 1024; // 4MB impl Default for Config { diff --git a/src/pd/cluster.rs b/src/pd/cluster.rs index 5ac0b0f3..adc34970 100644 --- a/src/pd/cluster.rs +++ b/src/pd/cluster.rs @@ -128,7 +128,7 @@ impl Connection { let members = self.validate_endpoints(endpoints, timeout).await?; let (client, keyspace_client, members) = self.try_connect_leader(&members, timeout).await?; let id = members.header.as_ref().unwrap().cluster_id; - let tso = TimestampOracle::new(id, &client)?; + let tso = TimestampOracle::new(id, &client, timeout)?; let cluster = Cluster { id, client, @@ -145,7 +145,7 @@ impl Connection { let start = Instant::now(); let (client, keyspace_client, members) = self.try_connect_leader(&cluster.members, timeout).await?; - let tso = TimestampOracle::new(cluster.id, &client)?; + let tso = TimestampOracle::new(cluster.id, &client, timeout)?; *cluster = Cluster { id: cluster.id, client, @@ -213,7 +213,7 @@ impl Connection { async fn connect( &self, addr: &str, - _timeout: Duration, + timeout: Duration, ) -> Result<( pdpb::pd_client::PdClient, keyspacepb::keyspace_client::KeyspaceClient, @@ -221,19 +221,25 @@ impl Connection { )> { let mut client = self .security_mgr - .connect(addr, pdpb::pd_client::PdClient::::new) + .connect_with_timeout(addr, timeout, pdpb::pd_client::PdClient::::new) .await?; let keyspace_client = self .security_mgr - .connect( + .connect_with_timeout( addr, + timeout, keyspacepb::keyspace_client::KeyspaceClient::::new, ) .await?; - let resp: pdpb::GetMembersResponse = client - .get_members(pdpb::GetMembersRequest::default()) - .await? - .into_inner(); + // Bounded like every other request in this path: this is the first RPC + // on a fresh connection, exactly where a frozen peer would hang it. + let resp: pdpb::GetMembersResponse = tokio::time::timeout( + timeout, + client.get_members(pdpb::GetMembersRequest::default()), + ) + .await + .map_err(|_| internal_err!("get_members timed out after {:?}", timeout))?? + .into_inner(); if let Some(err) = resp .header .as_ref() diff --git a/src/pd/timestamp.rs b/src/pd/timestamp.rs index 30cfd8ee..dd591bc7 100644 --- a/src/pd/timestamp.rs +++ b/src/pd/timestamp.rs @@ -14,6 +14,7 @@ use std::collections::VecDeque; use std::pin::Pin; use std::sync::Arc; +use std::time::Duration; use futures::pin_mut; use futures::prelude::*; @@ -22,10 +23,13 @@ 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; use tokio::sync::oneshot; +use tokio::sync::watch; use tokio::sync::Mutex; +use tokio::time::timeout; use tonic::transport::Channel; use crate::internal_err; @@ -50,27 +54,76 @@ pub(crate) struct TimestampOracle { /// In the working thread, the `TimestampRequest`, which is actually a one channel sender, /// is used to send back the timestamp result. request_tx: mpsc::Sender, + /// Closed (its sender dropped) when the TSO worker exits, however it exits. + /// `get_timestamp` watches this so its callers fail fast when the worker is + /// gone. Relying on the worker's teardown to drop the request channel is + /// not enough: the channel is owned by the gRPC request body inside the + /// connection's task, and on a frozen connection that task may not run + /// again — releasing nothing — until the connection thaws. + worker_closed: watch::Receiver<()>, } impl TimestampOracle { - pub(crate) fn new(cluster_id: u64, pd_client: &PdClient) -> Result { + /// `timeout` bounds every wait against the PD server: stream creation, and + /// each silent window of the response stream while request batches are + /// outstanding (see [`allocate_from_stream`]). It is the client's + /// configured request timeout (`Config::timeout`), whose contract covers + /// requests to PD nodes. A healthy PD answers a TSO request in + /// milliseconds, but a stream can go silent without terminating — the + /// peer's VM paused by its hypervisor, the process frozen under resource + /// starvation, a middlebox silently discarding the connection — and + /// nothing else in this path bounds the wait, so a request could otherwise + /// hang forever (#516). client-go bounds its TSO batches the same way: the + /// dispatcher arms a deadline watcher per batch with `defaultPDTimeout` = + /// 3s and cancels the stream on expiry (pd/client/clients/tso/dispatcher.go), + /// and bounds stream creation via `checkStreamTimeout` + /// (pd/client/clients/tso/stream.go). + pub(crate) fn new( + cluster_id: u64, + pd_client: &PdClient, + timeout: Duration, + ) -> Result { let pd_client = pd_client.clone(); let (request_tx, request_rx) = mpsc::channel(MAX_BATCH_SIZE); + let (close_on_exit, worker_closed) = watch::channel(()); // Start a background thread to handle TSO requests and responses - tokio::spawn(run_tso(cluster_id, pd_client, request_rx)); + tokio::spawn(async move { + let _close_on_exit = close_on_exit; + run_tso(cluster_id, pd_client, request_rx, timeout).await + }); - Ok(TimestampOracle { request_tx }) + Ok(TimestampOracle { + request_tx, + worker_closed, + }) } - pub(crate) async fn get_timestamp(self) -> Result { + pub(crate) async fn get_timestamp(mut self) -> Result { debug!("getting current timestamp"); let (request, response) = oneshot::channel(); - self.request_tx - .send(request) - .await - .map_err(|_| internal_err!("TimestampRequest channel is closed"))?; - Ok(response.await?) + // The worker's death must be observed while ENQUEUEING too: a dead + // worker's request channel can remain open (held by the frozen + // connection's task) and, once its buffer fills, `send` would block + // forever without ever failing. + tokio::select! { + biased; + sent = self.request_tx.send(request) => { + sent.map_err(|_| internal_err!("TimestampRequest channel is closed"))? + } + _ = self.worker_closed.changed() => { + return Err(internal_err!("the TSO worker terminated")); + } + } + tokio::select! { + // Deliver a timestamp that is already there even if the worker died + // in the same instant. + biased; + res = response => Ok(res?), + _ = self.worker_closed.changed() => { + Err(internal_err!("the TSO worker terminated")) + } + } } } @@ -78,6 +131,7 @@ async fn run_tso( cluster_id: u64, mut pd_client: PdClient, request_rx: mpsc::Receiver, + deadline: Duration, ) -> Result<()> { // The `TimestampRequest`s which are waiting for the responses from the PD server let pending_requests = Arc::new(Mutex::new(VecDeque::with_capacity(MAX_PENDING_COUNT))); @@ -94,21 +148,117 @@ async fn run_tso( self_waker: sending_future_waker.clone(), }; - // let send_requests = rpc_sender.send_all(&mut request_stream); - let mut responses = pd_client.tso(request_stream).await?.into_inner(); + let result = async { + // Stream creation needs the same bound as the receive loop, but ARMED + // ONLY WHILE A BATCH IS PENDING: PD (a gRPC-go server) sends response + // headers with its first response, so `tso(..)` legitimately waits as + // long as the client is idle — while against a frozen peer it would + // hang with a request owed (client-go: `checkStreamTimeout`). HTTP/2 + // lets the request stream flow before response headers arrive, so + // batches do get dispatched (and become pending) during this await. + let responses = tokio::select! { + created = pd_client.tso(request_stream) => created?.into_inner(), + stalled = stall_watchdog(&pending_requests, deadline) => return Err(stalled), + }; - while let Some(Ok(resp)) = responses.next().await { - { - let mut pending_requests = pending_requests.lock().await; - allocate_timestamps(&resp, &mut pending_requests)?; + allocate_from_stream( + responses, + &pending_requests, + &sending_future_waker, + deadline, + ) + .await + } + .await; + + match &result { + // The stream ended with nothing owed — e.g. the client is shutting down. + Ok(()) => info!("TSO stream terminated"), + // Failed to come up, stalled, or terminated with batches outstanding. + // The worker's exit closes the oracle's `worker_closed` watch, failing + // every waiting and future request fast, and the retry layer reconnects. + Err(e) => warn!("TSO stream failed: {}", e), + } + result +} + +/// Resolve (never `Ok`) once a full `deadline` window has passed with a batch +/// pending the whole time and no progress possible — the stall signal for +/// stream creation, where no response can arrive by definition. The same +/// [1, 2]-window guarantee as [`allocate_from_stream`]: a batch dispatched +/// mid-window is granted the following full window. +async fn stall_watchdog( + pending_requests: &Mutex>, + deadline: Duration, +) -> crate::Error { + loop { + let pending_at_window_start = !pending_requests.lock().await.is_empty(); + tokio::time::sleep(deadline).await; + if pending_at_window_start && !pending_requests.lock().await.is_empty() { + return internal_err!( + "TSO stream creation stalled: no stream within {:?} with requests pending", + deadline + ); } + } +} - // Wake up the sending future blocked by too many pending requests or locked. - sending_future_waker.wake(); +/// Drive the TSO response stream, allocating timestamps to the pending request +/// batches, until the stream ends or goes silent past `deadline` with batches +/// outstanding. +/// +/// The deadline is evaluated in whole windows: a batch dispatched near the end +/// of a silent window is granted one more full window before the stream is +/// declared stalled, so teardown happens only after between one and two +/// deadlines of genuine silence. (client-go stamps a deadline per batch at +/// dispatch time; windows approximate that without tracking send times.) +async fn allocate_from_stream( + mut responses: impl Stream> + Unpin, + pending_requests: &Mutex>, + sending_future_waker: &AtomicWaker, + deadline: Duration, +) -> Result<()> { + loop { + let pending_at_window_start = !pending_requests.lock().await.is_empty(); + match timeout(deadline, responses.next()).await { + Ok(Some(Ok(resp))) => { + { + let mut pending_requests = pending_requests.lock().await; + allocate_timestamps(&resp, &mut pending_requests)?; + } + // Wake up the sending future blocked by too many pending requests or locked. + sending_future_waker.wake(); + } + Ok(Some(Err(status))) => return Err(status.into()), + Ok(None) => { + let outstanding = pending_requests.lock().await.len(); + if outstanding == 0 { + return Ok(()); + } + return Err(internal_err!( + "TSO stream terminated with {} batches outstanding", + outstanding + )); + } + Err(_elapsed) => { + let outstanding = pending_requests.lock().await.len(); + if outstanding == 0 { + // A quiet stream with nothing owed is healthy idleness. + continue; + } + if pending_at_window_start { + // Owed a response for the entire silent window: stalled. + return Err(internal_err!( + "TSO stream stalled: no response within {:?} with {} batches outstanding", + deadline, + outstanding + )); + } + // The batch was dispatched mid-window; grant it one full window + // (the next iteration starts with it pending). + } + } } - // TODO: distinguish between unexpected stream termination and expected end of test - info!("TSO stream terminated"); - Ok(()) } struct RequestGroup { @@ -216,3 +366,237 @@ fn allocate_timestamps( }; Ok(()) } + +#[cfg(test)] +mod tests { + + use futures::stream; + + use super::*; + + const DEADLINE: Duration = Duration::from_secs(3); + + fn group_of_one() -> (RequestGroup, oneshot::Receiver) { + let (tx, rx) = oneshot::channel(); + let group = RequestGroup { + tso_request: TsoRequest { + count: 1, + ..Default::default() + }, + requests: vec![tx], + }; + (group, rx) + } + + /// The #516 hang: a stream that never yields while a batch is owed. The + /// deadline must tear the loop down — and the paused clock proves it takes + /// between one and two deadline windows, not forever. + #[tokio::test(start_paused = true)] + async fn a_stalled_stream_with_batches_outstanding_fails_within_two_deadlines() { + let pending = Mutex::new(VecDeque::new()); + let (group, _rx) = group_of_one(); + pending.lock().await.push_back(group); + let waker = AtomicWaker::new(); + + let t0 = tokio::time::Instant::now(); + let result = allocate_from_stream(stream::pending(), &pending, &waker, DEADLINE).await; + let elapsed = t0.elapsed(); + + let err = result.unwrap_err(); + assert!( + err.to_string().contains("stalled"), + "want the stall error, got {err:?}" + ); + assert!( + elapsed >= DEADLINE && elapsed <= DEADLINE * 2, + "want teardown within [1, 2] deadlines, took {elapsed:?}" + ); + } + + /// A batch dispatched mid-window must not be torn down at that window's + /// end — it gets one full window of its own. + #[tokio::test(start_paused = true)] + async fn a_batch_dispatched_mid_window_gets_a_full_window() { + let pending = Arc::new(Mutex::new(VecDeque::new())); + let waker = Arc::new(AtomicWaker::new()); + + let task = { + let pending = pending.clone(); + let waker = waker.clone(); + tokio::spawn(async move { + let t0 = tokio::time::Instant::now(); + let result = + allocate_from_stream(stream::pending(), &pending, &waker, DEADLINE).await; + (result, t0.elapsed()) + }) + }; + + // Dispatch the batch halfway through the first (idle) window. + tokio::time::sleep(DEADLINE / 2).await; + let (group, _rx) = group_of_one(); + pending.lock().await.push_back(group); + + let (result, elapsed) = task.await.unwrap(); + assert!(result.is_err()); + assert!( + elapsed >= DEADLINE * 2, + "the mid-window batch must get a full window before teardown, got {elapsed:?}" + ); + } + + /// Healthy idleness — no batches owed — must never be torn down, no matter + /// how long the stream stays quiet. + #[tokio::test(start_paused = true)] + async fn an_idle_stream_is_left_alone() { + let pending = Mutex::new(VecDeque::new()); + let waker = AtomicWaker::new(); + + let outcome = timeout( + DEADLINE * 20, + allocate_from_stream(stream::pending(), &pending, &waker, DEADLINE), + ) + .await; + assert!( + outcome.is_err(), + "an idle stream must keep waiting, but the loop returned {outcome:?}" + ); + } + + /// Termination with batches owed is a failure (the callers' requests are + /// dropped and fail fast); termination when idle is a clean end. + #[tokio::test] + async fn termination_is_clean_only_when_nothing_is_owed() { + let pending = Mutex::new(VecDeque::new()); + let waker = AtomicWaker::new(); + let empty: Vec> = vec![]; + let result = allocate_from_stream(stream::iter(empty), &pending, &waker, DEADLINE).await; + assert!(result.is_ok()); + + let (group, _rx) = group_of_one(); + pending.lock().await.push_back(group); + let empty: Vec> = vec![]; + let err = allocate_from_stream(stream::iter(empty), &pending, &waker, DEADLINE) + .await + .unwrap_err(); + assert!( + err.to_string().contains("terminated"), + "want the terminated-with-outstanding error, got {err:?}" + ); + } + + /// A status error carried by the stream surfaces as-is instead of being + /// silently swallowed (which the old loop did). + #[tokio::test] + async fn a_stream_error_propagates() { + let pending = Mutex::new(VecDeque::new()); + let waker = AtomicWaker::new(); + let items = vec![Err(tonic::Status::unavailable("pd gone"))]; + let err = allocate_from_stream(stream::iter(items), &pending, &waker, DEADLINE) + .await + .unwrap_err(); + assert!( + matches!(err, crate::Error::GrpcAPI(_)), + "want the tonic status surfaced, got {err:?}" + ); + } + + /// A client that is merely idle after construction must keep its worker: + /// PD sends response headers only with its first response, so stream + /// creation legitimately waits while nothing is owed. The watchdog must + /// never fire with the pending queue empty. + #[tokio::test(start_paused = true)] + async fn an_idle_worker_is_not_torn_down_during_stream_creation() { + let pending = Mutex::new(VecDeque::new()); + let outcome = timeout(DEADLINE * 20, stall_watchdog(&pending, DEADLINE)).await; + assert!( + outcome.is_err(), + "the watchdog must keep waiting while nothing is owed, got {outcome:?}" + ); + } + + /// With a batch owed and no stream forthcoming, the watchdog fires within + /// the same [1, 2]-window guarantee as the receive loop. + #[tokio::test(start_paused = true)] + async fn the_creation_watchdog_fires_with_a_batch_pending() { + let pending = Mutex::new(VecDeque::new()); + let (group, _rx) = group_of_one(); + pending.lock().await.push_back(group); + + let t0 = tokio::time::Instant::now(); + let err = stall_watchdog(&pending, DEADLINE).await; + let elapsed = t0.elapsed(); + assert!(err.to_string().contains("creation stalled")); + assert!( + elapsed >= DEADLINE && elapsed <= DEADLINE * 2, + "want the watchdog within [1, 2] windows, took {elapsed:?}" + ); + } + + /// Enqueueing must observe worker death too: a dead worker's request + /// channel can stay open (held by a frozen connection's task) with a full + /// buffer, and `send` would otherwise block forever. + #[tokio::test] + async fn enqueueing_observes_worker_death() { + let (request_tx, request_rx) = mpsc::channel(1); + let (close_on_exit, worker_closed) = watch::channel(()); + // The buffer is full and the receiver alive-but-unpolled: the zombie. + let (dummy, _dummy_rx) = oneshot::channel(); + request_tx.try_send(dummy).unwrap(); + // The worker is dead. + drop(close_on_exit); + + let oracle = TimestampOracle { + request_tx, + worker_closed, + }; + let outcome = timeout(Duration::from_secs(30), oracle.get_timestamp()) + .await + .expect("a full queue on a dead worker must fail, not block"); + assert!(outcome.is_err()); + drop(request_rx); + } + + /// The oracle's callers must fail fast when the worker is gone — the + /// worker's teardown cannot rely on dropping the request channel (a frozen + /// connection's task may hold it), so `get_timestamp` watches the worker's + /// exit directly. A lazy channel to a closed port makes the worker die at + /// stream creation. + #[tokio::test] + async fn get_timestamp_fails_fast_when_the_worker_dies() { + let channel = tonic::transport::Endpoint::from_static("http://127.0.0.1:1").connect_lazy(); + let pd_client = PdClient::new(channel); + let oracle = TimestampOracle::new(1, &pd_client, DEADLINE).unwrap(); + + let outcome = timeout(Duration::from_secs(30), oracle.get_timestamp()) + .await + .expect("a request against a dead worker must fail, not hang"); + assert!(outcome.is_err()); + } + + /// The happy path: a response allocates a timestamp to the pending batch, + /// and a subsequent clean end reports Ok. + #[tokio::test] + async fn responses_allocate_timestamps_to_pending_batches() { + let pending = Mutex::new(VecDeque::new()); + let (group, rx) = group_of_one(); + pending.lock().await.push_back(group); + let waker = AtomicWaker::new(); + + let items = vec![Ok(TsoResponse { + count: 1, + timestamp: Some(Timestamp { + physical: 7, + logical: 42, + ..Default::default() + }), + ..Default::default() + })]; + let result = allocate_from_stream(stream::iter(items), &pending, &waker, DEADLINE).await; + assert!(result.is_ok(), "clean end after all batches answered"); + + let ts = rx + .await + .expect("the pending request must receive its timestamp"); + assert_eq!((ts.physical, ts.logical), (7, 42)); + } +} diff --git a/src/store/client.rs b/src/store/client.rs index 1c873285..fea9080a 100644 --- a/src/store/client.rs +++ b/src/store/client.rs @@ -35,7 +35,7 @@ impl KvConnect for TikvConnect { async fn connect(&self, address: &str) -> Result { self.security_mgr - .connect(address, move |channel| { + .connect_with_timeout(address, self.timeout, move |channel| { TikvClient::new(channel) .max_decoding_message_size(self.grpc_max_decoding_message_size) .accept_compressed(CompressionEncoding::Gzip) From f60393daa79cdaa45ff23c04419efc8541c64748 Mon Sep 17 00:00:00 2001 From: "Eduard R." Date: Sat, 17 Oct 2026 13:22:43 +0200 Subject: [PATCH 2/2] pd: count queued TSO requests as owed, and keep the request stream's wakeup Two review findings on the TSO deadlines: 1. Lost wakeup. The watchdog and the receive loop inspected the pending queue through a tokio Mutex. If `TsoRequestStream::poll_next` met that lock held, it dropped its lock future (and with it the mutex waiter) and returned Pending; releasing the inspection's guard then woke nobody, and a queued request was never dispatched. The queue is now a std Mutex, never held across an await, so `poll_next` waits out the lock instead of parking. 2. Queued requests ignored. Only batches already taken by the transport counted as owed. The transport polls the request stream only when it can send, so a request enqueued while it cannot (connection not ready, no send window) left the deadline disarmed and the caller waiting forever on a live worker. What is owed is now pending batches plus requests still in the channel, read through a WeakSender so the worker does not hold the channel open; client-go likewise arms its batch deadline before the RPC is sent. Regression tests: the stream polled while another thread holds the lock; the creation watchdog and the receive loop each with a request queued but not dispatched (paused clock); and end to end through a tonic channel whose connection never completes. Each fails or hangs with its fix disabled. Ref: https://github.com/tikv/client-rust/issues/516 Signed-off-by: Eduard R. --- Cargo.toml | 1 + src/pd/timestamp.rs | 299 +++++++++++++++++++++++++++++++++----------- 2 files changed, 230 insertions(+), 70 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 6290cd6a..e1925718 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -52,6 +52,7 @@ tonic = { version = "0.12", features = ["tls", "gzip"] } clap = "2" env_logger = "0.10" fail = { version = "0.4", features = ["failpoints"] } +hyper-util = { version = "0.1", features = ["tokio"] } proptest = "1" proptest-derive = "0.5.1" reqwest = { version = "0.11", features = ["json", "native-tls-vendored"] } diff --git a/src/pd/timestamp.rs b/src/pd/timestamp.rs index dd591bc7..3507c819 100644 --- a/src/pd/timestamp.rs +++ b/src/pd/timestamp.rs @@ -14,9 +14,9 @@ use std::collections::VecDeque; use std::pin::Pin; use std::sync::Arc; +use std::sync::Mutex; use std::time::Duration; -use futures::pin_mut; use futures::prelude::*; use futures::task::AtomicWaker; use futures::task::Context; @@ -28,7 +28,6 @@ use pin_project::pin_project; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio::sync::watch; -use tokio::sync::Mutex; use tokio::time::timeout; use tonic::transport::Channel; @@ -88,9 +87,10 @@ impl TimestampOracle { let (close_on_exit, worker_closed) = watch::channel(()); // Start a background thread to handle TSO requests and responses + let queued = request_tx.downgrade(); tokio::spawn(async move { let _close_on_exit = close_on_exit; - run_tso(cluster_id, pd_client, request_rx, timeout).await + run_tso(cluster_id, pd_client, request_rx, queued, timeout).await }); Ok(TimestampOracle { @@ -131,9 +131,12 @@ async fn run_tso( cluster_id: u64, mut pd_client: PdClient, request_rx: mpsc::Receiver, + queued: mpsc::WeakSender, deadline: Duration, ) -> Result<()> { - // The `TimestampRequest`s which are waiting for the responses from the PD server + // The `TimestampRequest`s which are waiting for the responses from the PD server. + // A std mutex, never held across an await: `TsoRequestStream::poll_next` + // must not return `Pending` on contention, since nothing would wake it. let pending_requests = Arc::new(Mutex::new(VecDeque::with_capacity(MAX_PENDING_COUNT))); // When there are too many pending requests, the `send_request` future will refuse to fetch @@ -147,27 +150,23 @@ async fn run_tso( pending_requests: pending_requests.clone(), self_waker: sending_future_waker.clone(), }; + let owed = Owed { + pending_requests: &pending_requests, + queued, + }; let result = async { // Stream creation needs the same bound as the receive loop, but ARMED - // ONLY WHILE A BATCH IS PENDING: PD (a gRPC-go server) sends response + // ONLY WHILE A REQUEST IS OWED: PD (a gRPC-go server) sends response // headers with its first response, so `tso(..)` legitimately waits as // long as the client is idle — while against a frozen peer it would - // hang with a request owed (client-go: `checkStreamTimeout`). HTTP/2 - // lets the request stream flow before response headers arrive, so - // batches do get dispatched (and become pending) during this await. + // hang with a request owed (client-go: `checkStreamTimeout`). let responses = tokio::select! { created = pd_client.tso(request_stream) => created?.into_inner(), - stalled = stall_watchdog(&pending_requests, deadline) => return Err(stalled), + stalled = stall_watchdog(&owed, deadline) => return Err(stalled), }; - allocate_from_stream( - responses, - &pending_requests, - &sending_future_waker, - deadline, - ) - .await + allocate_from_stream(responses, &owed, &sending_future_waker, deadline).await } .await; @@ -182,80 +181,105 @@ async fn run_tso( result } -/// Resolve (never `Ok`) once a full `deadline` window has passed with a batch -/// pending the whole time and no progress possible — the stall signal for +/// What the worker owes its callers: batches dispatched to PD and awaiting a +/// response, plus requests enqueued but not yet taken by the transport. The +/// latter count too — the transport polls the request stream only when it can +/// send (HTTP/2 flow control, connection readiness), so against a peer that +/// grants no send window a request is owed without ever becoming a batch. +/// client-go likewise arms its batch deadline before the RPC is sent. +struct Owed<'a> { + pending_requests: &'a Mutex>, + /// Weak, so the worker does not keep the request channel open by itself. + queued: mpsc::WeakSender, +} + +impl Owed<'_> { + fn batches(&self) -> usize { + self.pending_requests.lock().unwrap().len() + } + + /// Requests in the channel, including senders holding a permit. + fn queued(&self) -> usize { + self.queued + .upgrade() + .map_or(0, |tx| tx.max_capacity() - tx.capacity()) + } + + fn any(&self) -> bool { + self.batches() > 0 || self.queued() > 0 + } +} + +/// Resolve (never `Ok`) once a full `deadline` window has passed with a +/// request owed the whole time and no progress possible — the stall signal for /// stream creation, where no response can arrive by definition. The same -/// [1, 2]-window guarantee as [`allocate_from_stream`]: a batch dispatched +/// [1, 2]-window guarantee as [`allocate_from_stream`]: a request enqueued /// mid-window is granted the following full window. -async fn stall_watchdog( - pending_requests: &Mutex>, - deadline: Duration, -) -> crate::Error { +async fn stall_watchdog(owed: &Owed<'_>, deadline: Duration) -> crate::Error { loop { - let pending_at_window_start = !pending_requests.lock().await.is_empty(); + let owed_at_window_start = owed.any(); tokio::time::sleep(deadline).await; - if pending_at_window_start && !pending_requests.lock().await.is_empty() { + if owed_at_window_start && owed.any() { return internal_err!( - "TSO stream creation stalled: no stream within {:?} with requests pending", - deadline + "TSO stream creation stalled: no stream within {:?} with {} batches outstanding and {} requests queued", + deadline, + owed.batches(), + owed.queued() ); } } } /// Drive the TSO response stream, allocating timestamps to the pending request -/// batches, until the stream ends or goes silent past `deadline` with batches -/// outstanding. +/// batches, until the stream ends or goes silent past `deadline` with requests +/// owed (see [`Owed`]). /// -/// The deadline is evaluated in whole windows: a batch dispatched near the end +/// The deadline is evaluated in whole windows: a request owed from near the end /// of a silent window is granted one more full window before the stream is /// declared stalled, so teardown happens only after between one and two /// deadlines of genuine silence. (client-go stamps a deadline per batch at /// dispatch time; windows approximate that without tracking send times.) async fn allocate_from_stream( mut responses: impl Stream> + Unpin, - pending_requests: &Mutex>, + owed: &Owed<'_>, sending_future_waker: &AtomicWaker, deadline: Duration, ) -> Result<()> { loop { - let pending_at_window_start = !pending_requests.lock().await.is_empty(); + let owed_at_window_start = owed.any(); match timeout(deadline, responses.next()).await { Ok(Some(Ok(resp))) => { - { - let mut pending_requests = pending_requests.lock().await; - allocate_timestamps(&resp, &mut pending_requests)?; - } + allocate_timestamps(&resp, &mut owed.pending_requests.lock().unwrap())?; // Wake up the sending future blocked by too many pending requests or locked. sending_future_waker.wake(); } Ok(Some(Err(status))) => return Err(status.into()), Ok(None) => { - let outstanding = pending_requests.lock().await.len(); - if outstanding == 0 { + if !owed.any() { return Ok(()); } return Err(internal_err!( - "TSO stream terminated with {} batches outstanding", - outstanding + "TSO stream terminated with {} batches outstanding and {} requests queued", + owed.batches(), + owed.queued() )); } Err(_elapsed) => { - let outstanding = pending_requests.lock().await.len(); - if outstanding == 0 { + if !owed.any() { // A quiet stream with nothing owed is healthy idleness. continue; } - if pending_at_window_start { + if owed_at_window_start { // Owed a response for the entire silent window: stalled. return Err(internal_err!( - "TSO stream stalled: no response within {:?} with {} batches outstanding", + "TSO stream stalled: no response within {:?} with {} batches outstanding and {} requests queued", deadline, - outstanding + owed.batches(), + owed.queued() )); } - // The batch was dispatched mid-window; grant it one full window - // (the next iteration starts with it pending). + // The request became owed mid-window; grant it one full window + // (the next iteration starts with it owed). } } } @@ -281,15 +305,7 @@ impl Stream for TsoRequestStream { fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll> { 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) - { - pending_requests - } else { - this.self_waker.register(cx.waker()); - return Poll::Pending; - }; + let mut pending_requests = this.pending_requests.lock().unwrap(); let mut requests = Vec::new(); while requests.len() < MAX_BATCH_SIZE && pending_requests.len() < MAX_PENDING_COUNT { @@ -376,6 +392,15 @@ mod tests { const DEADLINE: Duration = Duration::from_secs(3); + /// Nothing queued: the request channel is already gone. + fn owed(pending: &Mutex>) -> Owed<'_> { + let (tx, _) = mpsc::channel(1); + Owed { + pending_requests: pending, + queued: tx.downgrade(), + } + } + fn group_of_one() -> (RequestGroup, oneshot::Receiver) { let (tx, rx) = oneshot::channel(); let group = RequestGroup { @@ -395,11 +420,12 @@ mod tests { async fn a_stalled_stream_with_batches_outstanding_fails_within_two_deadlines() { let pending = Mutex::new(VecDeque::new()); let (group, _rx) = group_of_one(); - pending.lock().await.push_back(group); + pending.lock().unwrap().push_back(group); let waker = AtomicWaker::new(); let t0 = tokio::time::Instant::now(); - let result = allocate_from_stream(stream::pending(), &pending, &waker, DEADLINE).await; + let result = + allocate_from_stream(stream::pending(), &owed(&pending), &waker, DEADLINE).await; let elapsed = t0.elapsed(); let err = result.unwrap_err(); @@ -426,7 +452,8 @@ mod tests { tokio::spawn(async move { let t0 = tokio::time::Instant::now(); let result = - allocate_from_stream(stream::pending(), &pending, &waker, DEADLINE).await; + allocate_from_stream(stream::pending(), &owed(&pending), &waker, DEADLINE) + .await; (result, t0.elapsed()) }) }; @@ -434,7 +461,7 @@ mod tests { // Dispatch the batch halfway through the first (idle) window. tokio::time::sleep(DEADLINE / 2).await; let (group, _rx) = group_of_one(); - pending.lock().await.push_back(group); + pending.lock().unwrap().push_back(group); let (result, elapsed) = task.await.unwrap(); assert!(result.is_err()); @@ -453,7 +480,7 @@ mod tests { let outcome = timeout( DEADLINE * 20, - allocate_from_stream(stream::pending(), &pending, &waker, DEADLINE), + allocate_from_stream(stream::pending(), &owed(&pending), &waker, DEADLINE), ) .await; assert!( @@ -469,13 +496,14 @@ mod tests { let pending = Mutex::new(VecDeque::new()); let waker = AtomicWaker::new(); let empty: Vec> = vec![]; - let result = allocate_from_stream(stream::iter(empty), &pending, &waker, DEADLINE).await; + let result = + allocate_from_stream(stream::iter(empty), &owed(&pending), &waker, DEADLINE).await; assert!(result.is_ok()); let (group, _rx) = group_of_one(); - pending.lock().await.push_back(group); + pending.lock().unwrap().push_back(group); let empty: Vec> = vec![]; - let err = allocate_from_stream(stream::iter(empty), &pending, &waker, DEADLINE) + let err = allocate_from_stream(stream::iter(empty), &owed(&pending), &waker, DEADLINE) .await .unwrap_err(); assert!( @@ -491,7 +519,7 @@ mod tests { let pending = Mutex::new(VecDeque::new()); let waker = AtomicWaker::new(); let items = vec![Err(tonic::Status::unavailable("pd gone"))]; - let err = allocate_from_stream(stream::iter(items), &pending, &waker, DEADLINE) + let err = allocate_from_stream(stream::iter(items), &owed(&pending), &waker, DEADLINE) .await .unwrap_err(); assert!( @@ -507,7 +535,7 @@ mod tests { #[tokio::test(start_paused = true)] async fn an_idle_worker_is_not_torn_down_during_stream_creation() { let pending = Mutex::new(VecDeque::new()); - let outcome = timeout(DEADLINE * 20, stall_watchdog(&pending, DEADLINE)).await; + let outcome = timeout(DEADLINE * 20, stall_watchdog(&owed(&pending), DEADLINE)).await; assert!( outcome.is_err(), "the watchdog must keep waiting while nothing is owed, got {outcome:?}" @@ -520,10 +548,10 @@ mod tests { async fn the_creation_watchdog_fires_with_a_batch_pending() { let pending = Mutex::new(VecDeque::new()); let (group, _rx) = group_of_one(); - pending.lock().await.push_back(group); + pending.lock().unwrap().push_back(group); let t0 = tokio::time::Instant::now(); - let err = stall_watchdog(&pending, DEADLINE).await; + let err = stall_watchdog(&owed(&pending), DEADLINE).await; let elapsed = t0.elapsed(); assert!(err.to_string().contains("creation stalled")); assert!( @@ -532,6 +560,136 @@ mod tests { ); } + /// A request the transport has not yet taken from the channel is owed + /// too: against a peer that grants no send window the request stream is + /// never polled, so no batch becomes pending, yet a caller is waiting. + #[tokio::test(start_paused = true)] + async fn the_creation_watchdog_fires_with_a_request_queued() { + let pending = Mutex::new(VecDeque::new()); + let (tx, _rx) = mpsc::channel(MAX_BATCH_SIZE); + let (request, _response) = oneshot::channel(); + tx.try_send(request).unwrap(); + let owed = Owed { + pending_requests: &pending, + queued: tx.downgrade(), + }; + + let t0 = tokio::time::Instant::now(); + let err = stall_watchdog(&owed, DEADLINE).await; + let elapsed = t0.elapsed(); + assert!(err.to_string().contains("1 requests queued"), "{err}"); + assert!( + elapsed >= DEADLINE && elapsed <= DEADLINE * 2, + "want the watchdog within [1, 2] windows, took {elapsed:?}" + ); + } + + /// The same for an established stream: a queued request the transport + /// cannot send is a stall, not idleness. + #[tokio::test(start_paused = true)] + async fn a_stalled_stream_with_a_request_queued_fails_within_two_deadlines() { + let pending = Mutex::new(VecDeque::new()); + let (tx, _rx) = mpsc::channel(MAX_BATCH_SIZE); + let (request, _response) = oneshot::channel(); + tx.try_send(request).unwrap(); + let owed = Owed { + pending_requests: &pending, + queued: tx.downgrade(), + }; + let waker = AtomicWaker::new(); + + let t0 = tokio::time::Instant::now(); + let err = allocate_from_stream(stream::pending(), &owed, &waker, DEADLINE) + .await + .unwrap_err(); + let elapsed = t0.elapsed(); + assert!(err.to_string().contains("stalled"), "{err}"); + assert!( + elapsed >= DEADLINE && elapsed <= DEADLINE * 2, + "want teardown within [1, 2] deadlines, took {elapsed:?}" + ); + } + + /// Inspecting the pending queue (the watchdog, the receive loop) must not + /// cost the request stream its wakeup: if `poll_next` met a held lock and + /// returned `Pending`, nothing would wake it once the lock was released, + /// and a queued request would never be dispatched. `poll_next` waits out + /// the lock instead. + #[test] + fn the_request_stream_is_not_parked_by_lock_contention() { + let (tx, request_rx) = mpsc::channel(MAX_BATCH_SIZE); + let (request, _response) = oneshot::channel(); + tx.try_send(request).unwrap(); + let pending_requests = Arc::new(Mutex::new(VecDeque::new())); + let mut stream = TsoRequestStream { + cluster_id: 1, + request_rx, + pending_requests: pending_requests.clone(), + self_waker: Arc::new(AtomicWaker::new()), + }; + + let held = Arc::new(std::sync::Barrier::new(2)); + let inspector = { + let pending_requests = pending_requests.clone(); + let held = held.clone(); + std::thread::spawn(move || { + let _guard = pending_requests.lock().unwrap(); + held.wait(); + std::thread::sleep(Duration::from_millis(50)); + }) + }; + held.wait(); + + let mut cx = Context::from_waker(futures::task::noop_waker_ref()); + let polled = Pin::new(&mut stream).poll_next(&mut cx); + inspector.join().unwrap(); + assert!( + matches!(polled, Poll::Ready(Some(ref req)) if req.count == 1), + "the request must be dispatched despite the contention, got {polled:?}" + ); + assert_eq!(pending_requests.lock().unwrap().len(), 1); + } + + /// End to end through a real tonic channel whose transport never becomes + /// ready — its connection attempt never completes — so the request body + /// is never polled and the request stays in the channel while the worker + /// is alive. The caller must still fail within the [1, 2]-window bound. + #[tokio::test(start_paused = true)] + async fn a_transport_that_never_becomes_ready_fails_within_two_deadlines() { + struct NeverConnects; + impl tonic::codegen::Service for NeverConnects { + type Response = hyper_util::rt::TokioIo; + type Error = std::io::Error; + type Future = future::Pending>; + + fn poll_ready( + &mut self, + _: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, _: tonic::transport::Uri) -> Self::Future { + future::pending() + } + } + + let channel = tonic::transport::Endpoint::from_static("http://127.0.0.1:1") + .connect_with_connector_lazy(NeverConnects); + let oracle = TimestampOracle::new(1, &PdClient::new(channel), DEADLINE).unwrap(); + + let t0 = tokio::time::Instant::now(); + let outcome = timeout(DEADLINE * 20, oracle.get_timestamp()) + .await + .expect("a request the transport never takes must fail, not hang"); + let elapsed = t0.elapsed(); + assert!(outcome.is_err()); + assert!( + elapsed >= DEADLINE && elapsed <= DEADLINE * 2, + "want failure within [1, 2] windows, took {elapsed:?}" + ); + } + /// Enqueueing must observe worker death too: a dead worker's request /// channel can stay open (held by a frozen connection's task) with a full /// buffer, and `send` would otherwise block forever. @@ -579,7 +737,7 @@ mod tests { async fn responses_allocate_timestamps_to_pending_batches() { let pending = Mutex::new(VecDeque::new()); let (group, rx) = group_of_one(); - pending.lock().await.push_back(group); + pending.lock().unwrap().push_back(group); let waker = AtomicWaker::new(); let items = vec![Ok(TsoResponse { @@ -591,7 +749,8 @@ mod tests { }), ..Default::default() })]; - let result = allocate_from_stream(stream::iter(items), &pending, &waker, DEADLINE).await; + let result = + allocate_from_stream(stream::iter(items), &owed(&pending), &waker, DEADLINE).await; assert!(result.is_ok(), "clean end after all batches answered"); let ts = rx