pd: reopen the TSO stream after it fails

[dina-kar]

Motivation

After a PD stall, restart or leadership blip, begin fails with

InternalError { message: "[src/pd/timestamp.rs:72]: TimestampRequest channel is closed" }

even though PD is already serving again.

The cause is in pd/timestamp.rs. When the TSO stream ends or errors, run_tso returns and drops the request receiver, and the TimestampOracle is dead from then on. The only way back is RetryClient::reconnect, which builds a new Cluster and oracle. But reconnect skips reconnecting when the previous reconnect was less than RECONNECT_INTERVAL_SEC (1 s) ago. When the fresh stream also dies right away (PD answers ErrGenerateTimestamp ... requested pd is not leader of cluster for a moment after it resumes), each later attempt fails in microseconds on the closed channel. The reconnect is skipped every time, and all ten attempts of retry_core! are used up within one interval.

Reproduced on a tiup playground (v8.5.8): a loop runs begin / put / commit every 500 ms while PD gets SIGSTOP for 5 s. After SIGCONT, two consecutive transactions failed with TimestampRequest channel is closed (one after waiting 4.7 s, one in 111 µs) before the client recovered. An older release in the same setup never recovered without rebuilding the client.

What is changed and how it works

The oracle now survives the loss of its stream:

• The request receiver is shared by the oracle's successive streams (Arc<Mutex<Receiver>>). The request channel stays open for as long as any TimestampOracle handle exists.
• When a stream fails, StreamState::fail marks it dead and fails the requests waiting on it with TSO stream failed: <cause>. A dead stream takes no more requests from the channel: the check runs under the pending-queue lock, so a request is either failed or left in the channel for the next stream, and never lost.
• The background task reopens the stream for the next request, with backoff (50 ms doubling to 1 s, reset after a stream that answered). The request that woke it is sent first. An idle oracle does not keep dialling an unavailable PD.
• The task ends when every oracle handle is dropped, as before. The channel closes, the request side returns None, and the task tells this apart from a stream failure.
• The oneshot now carries Result<Timestamp>, so callers see the stream's error instead of a bare RecvError.
• Opening the stream goes through a small private TsoConnector trait (production: PdClient::tso), so the oracle can be unit-tested with in-memory streams.

RetryClient is unchanged. A leader change still goes through retry_core! and reconnect, which replace the oracle. The difference is that the retries of a failed get_timestamp now wait for the oracle's reopened stream instead of failing at once on a closed channel.

Tests

• Unit (pd::timestamp::tests, in-memory PD):
  • requests_after_a_stream_failure_are_served_by_a_new_stream
  • a_refused_stream_is_reopened_with_backoff
  • concurrent_requests_survive_repeated_stream_failures: 32 callers, streams that fail after 3, 0 and 5 requests; all timestamps unique.
  • the_background_task_ends_when_the_oracle_is_dropped
• Manual, on a tiup playground v8.5.8 (API v2): the begin/commit loop above with PD SIGSTOPped for 5 s and then 12 s. No transaction failed after either stall. The client logged TSO stream failed: ... not leader ...; reopening it in 50ms/100ms and carried on.
• cargo test --lib, cargo clippy --all-targets --features integration-tests -- -D clippy::all, cargo fmt --check: clean. The failpoint (9), integration (27) and sync (28) suites pass on a tiup playground v8.5.8.

Notes for reviewers

• The error text of a request that was waiting on a failed stream changes from TimestampRequest channel is closed (or Channel(RecvError)) to TSO stream failed: .... TimestampRequest channel is closed can still happen, but only when every oracle handle has been dropped.
• No new dependencies (async-trait and futures were already used).

Check list

• Unit tests
• Manual test (playground, PD SIGSTOP)
• Integration test: needs a way to stall PD from the test; happy to add one gated behind an env var if wanted.
• DCO: the commit is signed off.

Release note: the transaction client recovers from a PD stall without failing begin with "TimestampRequest channel is closed".

Summary by CodeRabbit

• Reliability
  • Timestamp requests now recover from stream failures by reconnecting with backoff. Requests waiting on a failed stream receive an error, while new requests can trigger a fresh connection.
  • The background connection task exits when all oracle handles are dropped.

[ti-chi-bot]

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by:
Once this PR has been reviewed and has the lgtm label, please assign andremouche for approval. For more information see the Code Review Process.
Please ensure that each of them provides their approval before proceeding.

The full list of commands accepted by this bot can be found here.

Details

Needs approval from an approver in each of these files:

• OWNERS

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

[coderabbitai]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Advanced

Run ID: 9dc7fb38-c51b-46cf-a460-bdcf51aa642c

📥 Commits

Reviewing files that changed from the base of the PR and between ab4be1c and 7acba21.

📒 Files selected for processing (1)

• src/pd/timestamp.rs

Included review availability: This review used your included allowance. Your plan provides up to 2 included reviews per hour; 1 remain after this review.

---

📝 Walkthrough

Walkthrough

TimestampOracle now uses a persistent background task to manage TSO streams. The task reports stream failures to pending requests, reconnects with capped exponential backoff, and opens a replacement stream when a request is available. Added tests cover reconnects, concurrent requests, refused stream opens, and shutdown.

Changes

Timestamp Oracle reconnect handling

Layer / File(s)	Summary
Request result contract and stream connector src/pd/timestamp.rs	Timestamp requests now carry results that can include stream errors. TimestampOracle::new uses with_connector to set up the background task and TSO connector.
Stream serving, reconnects, and validation src/pd/timestamp.rs	run_tso manages stream failures, pending requests, capped backoff, and replacement streams. Tests cover reconnects, refused opens, concurrent requests, and shutdown.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Bug fix

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant run_tso
  participant TsoConnector
  participant TsoRequestStream
  participant PD_TSO
  Caller->>run_tso: Submit timestamp request
  run_tso->>TsoConnector: Open stream when a request is available
  TsoConnector->>PD_TSO: Create TSO stream
  run_tso->>TsoRequestStream: Serve stream
  TsoRequestStream->>PD_TSO: Send request
  PD_TSO->>run_tso: Return response
  run_tso->>Caller: Send Ok(timestamp)
  PD_TSO->>run_tso: Report stream failure
  run_tso->>Caller: Fail pending requests

Loading

Merge Risk: ⚪ Minimal · up to 7acba

No actionable merge-blocking risk was established for the reconnect change. It is mergeable after normal checks.

Security Architecture Review

Security architecture risk: 🔵 Low · up to 7acba

The new recovery path reuses the existing PD connection and does not appear to expose a new endpoint or weaken its configured connection controls. The principal remaining availability limitations also existed before this change, although recovery under every failure mode has not been verified.

Retained concerns
No architecture-level concerns identified.

Security review details

Security Blast Radius

• inferred — Failure or delay in PD timestamp delivery can affect client operations that require a timestamp. The inspected path does not establish a new independently reachable endpoint or broader privilege.

Trust Boundaries and Controls

• observed — The production connector sends requests over the existing PD client; PD channel construction applies the configured security settings. Reopening a stream does not itself create another credential path.

Resilience and Maintainability Implications

• inferred — The new retry loop handles returned stream errors, but cannot recover while the existing response wait remains pending indefinitely. This limitation predates the PR; its manifestation for a particular PD failure mode is not established here.

Hardening Proposals

• proposed — Consider a bounded response wait and shutdown detection independent of pending-queue capacity, so silent PD behavior cannot indefinitely retain in-flight requests or the oracle task.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 20 functions across 1 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: reopening the TSO stream after a failure.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

Warning

Some tools did not complete. Review the errors below.

🔧 Clippy (1.98.1)

Clippy execution failed

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}