diff --git a/Cargo.toml b/Cargo.toml index ca1029f3..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"] } @@ -59,7 +60,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 d8d013ec..44193b54 100644 --- a/src/common/security.rs +++ b/src/common/security.rs @@ -76,35 +76,36 @@ impl SecurityManager { } /// Connect to gRPC server using TLS connection. If TLS is not configured, use normal connection. - pub async fn connect( - &self, - // env: Arc, - addr: &str, - factory: Factory, - ) -> Result - where - Factory: FnOnce(Channel) -> Client, - { - self.connect_inner(addr, None, factory).await - } - - /// Connect with an explicit TCP connection timeout. - pub(crate) async fn connect_with_timeout( - &self, - addr: &str, - timeout: Duration, - factory: Factory, - ) -> Result + /// + /// 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_inner(addr, Some(timeout), factory).await + self.connect_with_timeout(addr, crate::config::DEFAULT_REQUEST_TIMEOUT, factory) + .await } - async fn connect_inner( + /// 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). + /// + /// The same bound is handed to the endpoint as tonic's `connect_timeout`, + /// which caps the TCP connection phase — tonic installs it on hyper's + /// `HttpConnector` — and keeps applying to reconnects the channel makes on + /// its own. It stops there, though: TLS and the HTTP/2 handshake both run + /// after the connector returns, so only the outer deadline covers them, + /// and the handshake is where a frozen peer actually strands the dial. + /// Whichever bound trips first decides the error a caller sees — tonic's + /// transport error from the connector, or the descriptive one below. + pub async fn connect_with_timeout( &self, + // env: Arc, addr: &str, - timeout: Option, + connect_timeout: Duration, factory: Factory, ) -> Result where @@ -116,11 +117,16 @@ impl SecurityManager { } else { self.default_channel(addr).await? }; - let channel = match timeout { - Some(timeout) => channel.connect_timeout(timeout), - None => channel, - }; - let ch = channel.connect().await?; + let channel = channel.connect_timeout(connect_timeout); + 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..3507c819 100644 --- a/src/pd/timestamp.rs +++ b/src/pd/timestamp.rs @@ -14,18 +14,21 @@ 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; 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::Mutex; +use tokio::sync::watch; +use tokio::time::timeout; use tonic::transport::Channel; use crate::internal_err; @@ -50,27 +53,77 @@ 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)); + let queued = request_tx.downgrade(); + tokio::spawn(async move { + let _close_on_exit = close_on_exit; + run_tso(cluster_id, pd_client, request_rx, queued, 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,8 +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 @@ -93,22 +150,139 @@ 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 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`). + let responses = tokio::select! { + created = pd_client.tso(request_stream) => created?.into_inner(), + stalled = stall_watchdog(&owed, deadline) => return Err(stalled), + }; + + allocate_from_stream(responses, &owed, &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 +} - // let send_requests = rpc_sender.send_all(&mut request_stream); - let mut responses = pd_client.tso(request_stream).await?.into_inner(); +/// 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 + } +} - while let Some(Ok(resp)) = responses.next().await { - { - let mut pending_requests = pending_requests.lock().await; - allocate_timestamps(&resp, &mut pending_requests)?; +/// 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 request enqueued +/// mid-window is granted the following full window. +async fn stall_watchdog(owed: &Owed<'_>, deadline: Duration) -> crate::Error { + loop { + let owed_at_window_start = owed.any(); + tokio::time::sleep(deadline).await; + if owed_at_window_start && owed.any() { + return internal_err!( + "TSO stream creation stalled: no stream within {:?} with {} batches outstanding and {} requests queued", + deadline, + owed.batches(), + owed.queued() + ); } + } +} - // 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 requests +/// owed (see [`Owed`]). +/// +/// 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, + owed: &Owed<'_>, + sending_future_waker: &AtomicWaker, + deadline: Duration, +) -> Result<()> { + loop { + let owed_at_window_start = owed.any(); + match timeout(deadline, responses.next()).await { + Ok(Some(Ok(resp))) => { + 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) => { + if !owed.any() { + return Ok(()); + } + return Err(internal_err!( + "TSO stream terminated with {} batches outstanding and {} requests queued", + owed.batches(), + owed.queued() + )); + } + Err(_elapsed) => { + if !owed.any() { + // A quiet stream with nothing owed is healthy idleness. + continue; + } + 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 and {} requests queued", + deadline, + owed.batches(), + owed.queued() + )); + } + // The request became owed mid-window; grant it one full window + // (the next iteration starts with it owed). + } + } } - // TODO: distinguish between unexpected stream termination and expected end of test - info!("TSO stream terminated"); - Ok(()) } struct RequestGroup { @@ -131,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 { @@ -216,3 +382,380 @@ fn allocate_timestamps( }; Ok(()) } + +#[cfg(test)] +mod tests { + + use futures::stream; + + use super::*; + + 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 { + 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().unwrap().push_back(group); + let waker = AtomicWaker::new(); + + let t0 = tokio::time::Instant::now(); + let result = + allocate_from_stream(stream::pending(), &owed(&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(), &owed(&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().unwrap().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(), &owed(&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), &owed(&pending), &waker, DEADLINE).await; + assert!(result.is_ok()); + + let (group, _rx) = group_of_one(); + pending.lock().unwrap().push_back(group); + let empty: Vec> = vec![]; + let err = allocate_from_stream(stream::iter(empty), &owed(&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), &owed(&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(&owed(&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().unwrap().push_back(group); + + let t0 = tokio::time::Instant::now(); + let err = stall_watchdog(&owed(&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:?}" + ); + } + + /// 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. + #[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().unwrap().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), &owed(&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)); + } +}