perf - Reduce TSO waker churn and quantify impact with Criterion

[mingley]

Summary

This PR improves correctness and efficiency of TSO request-stream wake coordination in src/pd/timestamp.rs.

Quality objective:

• prevent sender stalls under lock contention,
• reduce unnecessary wake/waker churn,
• preserve timestamp allocation correctness and backpressure behavior.

Rationale

Problem statement

The prior sender/response coordination could generate redundant wake operations and had a lock-contention interleaving that could miss a needed wake signal.

Design changes

• Response wake policy is conditional, not unconditional:
  • wake when pending queue transitions from full -> not full,
  • or when sender explicitly signaled lock-wait.
• Sender lock-contention path now uses try_lock plus a register-and-retry handshake:
  • register waker,
  • set sender_waiting_on_lock,
  • retry lock acquisition once immediately to close the lost-wake window.
• Self-waker registration remains scoped to full-queue blocking paths.
• Error propagation and observability are preserved:
  • TSO stream errors are surfaced,
  • non-empty batch emission calls observe_tso_batch(...).

Risk analysis

• Primary risk: race regressions in async wake sequencing.
• Mitigation: targeted unit tests for contention and wake edge cases, plus full clippy/test gates.

File Scope

• src/pd/timestamp.rs
• Cargo.toml (branch-level context change already present in PR)

Testing Done

Executed locally:

1. cargo fmt -> pass
2. cargo clippy --all-targets --all-features -- -D warnings -> pass
3. cargo test -> pass

Focused concurrency coverage in src/pd/timestamp.rs includes:

• poll_next_marks_waiting_flag_when_lock_is_contended_and_response_wakes
• register_sender_wait_sets_waiting_flag_and_registers_waker_on_retry_failure
• register_sender_wait_retries_once_and_clears_waiting_flag_when_lock_reacquires
• poll_next_clears_waiting_flag_on_lock_acquire
• poll_next_registers_self_waker_when_pending_queue_is_full
• poll_next_does_not_register_self_waker_when_queue_not_full
• timestamp allocation invariants and error-path tests

Compatibility

No public API surface change is introduced by this PR.

[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 you06 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

[ti-chi-bot]

Welcome @mingley!

It looks like this is your first PR to tikv/client-rust 🎉.

I'm the bot to help you request reviewers, add labels and more, See available commands.

We want to make sure your contribution gets all the attention it needs!



Thank you, and welcome to tikv/client-rust. 😃

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Adds a Criterion benchmark and docs for TSO waker policies, and modifies TSO timestamp handling to track batch sizes, introduce an AtomicBool for lock-wait detection, adjust wake/register semantics based on pending-queue fullness transitions, and add unit tests for the new wake behavior.

Changes

Cohort / File(s)	Summary
Build config Cargo.toml	Added criterion = "0.5" to [dev-dependencies] and a [[bench]] entry tso_waker_policy (harness = false).
Benchmarks benches/tso_waker_policy.rs	New Criterion micro-benchmark file comparing old vs new TSO waker wake/register policies across synthetic pending-queue scenarios; noop waker and grouped measurements.
Documentation doc/tso_waker_criterion.md	New documentation describing the benchmark goal, method (cargo bench / Criterion), target file, config, usage, and reported latency comparisons.
TSO runtime & poll logic src/pd/timestamp.rs	Spawn TSO loop as async task; add AtomicBool sender_waiting_on_lock; propagate it into TsoRequestStream and run_tso; compute should_wake_sender from full→non-full transitions or cleared lock-wait; wake sender only on capacity gain or when contention cleared; register self-waker only when still blocked by full queue; call observe_tso_batch(...) for non-empty batches; add unit tests for wake semantics and contention handling.

Sequence Diagram(s)

mermaid
sequenceDiagram
participant Client
participant TSO_Task as "TSO Task"
participant PendingQ as "Pending Queue"
participant Sender as "Sender Waker"

Client->>TSO_Task: submit TSO request
TSO_Task->>PendingQ: attempt append / allocate (read prev_full)
PendingQ-->>TSO_Task: prev_full_state
TSO_Task->>PendingQ: update queue (compute curr_full_state)
alt prev_full -> now not full OR sender_waiting_on_lock set
    TSO_Task->>Sender: wake sender
    Sender-->>Client: resume sender processing
else still full
    TSO_Task->>TSO_Task: register self-waker (blocked)
end
TSO_Task->>TSO_Task: observe_tso_batch(size) when batch formed

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Poem

I rabbit-hop through benches bright,
I time each wake in morning light,
Queues shrank, I peeked, then gave a thump,
A tiny wake, a noiseless jump—
I count the ticks and kiss the heap. 🐇

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 31.82% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: optimizing TSO waker behavior and adding benchmarking with Criterion to measure the performance impact.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

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 and usage tips.}

[coderabbitai]

Actionable comments posted: 0

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

src/pd/timestamp.rs (1)

78-82: ⚠️ Potential issue | 🟡 Minor

Background task errors are silently discarded due to missing JoinHandle handling.

The Result<()> return type and explicit Ok(()) at line 116 align cleanly with the ? usage inside the function. However, the JoinHandle returned by tokio::spawn(run_tso(...)) at line 62 is never stored or awaited. Since errors can occur at both the pd_client.tso() call (line 99) and within allocate_timestamps() (line 105), failures in the background task will go unnoticed and the connection closure will only be discovered when callers receive a channel-closed error instead of the root cause.

Consider storing the JoinHandle and handling its potential error, or spawn a task that logs/propagates failures.

🧹 Nitpick comments (3)

benches/tso_waker_policy.rs (2)

20-36: The "old" and "new" response benchmarks have asymmetric work, which is expected but worth noting.

In response_policy_old, wake() is called unconditionally on every iteration, while in response_policy_new, it's called only on the full→non-full transition (~once per 1024 iterations). The reported speedup primarily measures the cost of not calling wake(), rather than the overhead of the conditional check itself. This is fine for validating the optimization's effect, but the doc and PR description should be clear that the speedup reflects the amortized skip rate under this specific simulation pattern.

Also applies to: 38-57

---

9-11: Consider documenting what FULL_EVERY and FULL_WINDOW represent.

These simulation parameters control how often the queue becomes full in the benchmark, directly affecting the measured speedup ratio. A brief comment explaining their role would help future readers understand and tune the benchmark.

doc/tso_waker_criterion.md (1)

1-58: Consider noting that results should be re-run when the benchmark or production code changes.

Hardcoded benchmark results in committed documentation risk becoming stale as the code evolves. Consider adding a note that these numbers are a point-in-time snapshot and should be re-collected after significant changes to timestamp.rs or the benchmark itself.

[mingley]

Addressed the CodeRabbit feedback in commit cd6aa8f:

1. Background TSO task errors are now surfaced via logs

• Updated src/pd/timestamp.rs so the spawned task wraps run_tso(...) and logs any Err(...) with context (warn!) instead of silently discarding it.

2. Benchmark assumptions are now explicit

• Added comments in benches/tso_waker_policy.rs for FULL_EVERY / FULL_WINDOW.
• Added inline notes clarifying old vs new wake behavior (unconditional vs transition-only wake).

3. Benchmark doc now includes interpretation/staleness guidance

• Updated doc/tso_waker_criterion.md to state these numbers are point-in-time and should be re-run after changes to src/pd/timestamp.rs or the benchmark.
• Added a note that the response benchmark speedup reflects amortized skipped wake calls under this simulation pattern.

[mingley]

@coderabbitai review

[coderabbitai]

✅ Actions performed

Review triggered.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[mingley]

Addressed local review finding #1 in bba9a1c:

• src/pd/timestamp.rs: changed stream loop from while let Some(Ok(resp)) to while let Some(resp) + let resp = resp?; so gRPC stream errors are propagated out of run_tso and logged by the spawned-task wrapper instead of being silently treated as normal termination.