Arm the per-row request timeout only when the decode suspends

[Saurabh Singh (saurabh500)]

Description

Rows decoded from an already-buffered TDS packet usually complete on their first poll. The row path still wrapped every decode in tokio::time::timeout, paying for a clock read, Sleep construction, and a larger wrapper future even though no wait occurred.

This PR polls the decode future once and creates the timeout only when that poll returns Pending. Tokio also polls the inner future before its delay, so the ready path keeps the same behavior while avoiding the unused timeout setup.

The lazy timeout is used by the four row-data operations in NetworkTransport:

• receive_row_into
• receive_row_header
• resume_row_into
• read_active_plp_bytes

Connection, cancellation-drain, and token-boundary timeouts stay eager because they are not paid per row or column. The PLP streaming site is changed for consistency with the adjacent row-resume path; it was not covered by the benchmark.

Timeout behavior

The remaining request budget can be Some(Duration::ZERO). The implementation therefore keys on whether the decode suspends, not whether the budget is zero:

first poll	result
Ready	succeeds, matching Tokio
Pending	arms the timeout and returns Elapsed when the exhausted budget wins

Tests compare these cases directly with tokio::time::timeout, including a future that suspends once and then becomes ready.

The implementation is an inline macro because the equivalent async fn helper measured about 4% slower by growing the per-row state machine.

Performance

PR #264 is now in main and removes the dominant boxed packet-reader overhead from row decoding. On that optimized path, a paired ABBA benchmark measured this change at about 254 ns/row, or 15% of isolated decode time for a 48-column row; the eager arm was slower in 202 of 205 pairs. See the measurement on #271.

This isolates row-decode CPU cost. End-to-end impact will be smaller when network, TLS, or server time dominates.

The future-size guard still passes and now also covers read_active_plp_chunk.

Cancellation handling is unchanged: each operation still uses CancelHandle::run_until_cancelled.

Related Issues

Fixes #271

Checklist

• cargo bfmt passes
• cargo bclippy passes
• cargo btest passes
• New/changed functionality has tests
• Public API changes are documented (none)

[Copilot]

Pull request overview

Optimizes row decoding by deferring request-timeout creation until decoding suspends.

Changes:

• Adds a lazy timeout macro with behavioral tests.
• Applies it to three row/column decoding paths.
• Registers the timeout module.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
request_timeout.rs	Adds lazy timeout logic and tests.
network_transport.rs	Uses lazy timeouts in hot decode paths.
transport.rs	Registers the new internal module.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[github-actions]

📊 Code Coverage Report

🔥 Diff Coverage 100%

🎯 Overall Coverage 91.5%

📦 Project: mssql-tds + mssql-odbc + mssql-py-core ℹ️ Note: diff coverage is reported, not enforced.

---

Diff Coverage

Diff: main...HEAD, staged and unstaged changes

• mssql-tds/src/connection/tds_client.rs (100%)
• mssql-tds/src/connection/transport/request_timeout.rs (100%)

Summary

• Total: 83 lines
• Missing: 0 lines
• Coverage: 100%

---

🔗 Quick Links

View Azure DevOps Build · Coverage Report

[David Engel (David-Engel)]

Summary

Replaces eager tokio::time::timeout with a probe-then-arm macro on the four per-row/per-column read paths. The core reasoning is sound and the implementation looks correct to me:

• the probe's poll_fn always returns Ready, so it never suspends and no wakeup can be lost;
• the waker the inner future registers on the probe is the same task waker Timeout uses on its immediate next poll, so there is no gap;
• Option<Duration> is Copy, so $budget is evaluated exactly once;
• macro-local fut / budget are hygienic against the caller's expressions;
• cancellation ordering (run_until_cancelled nested inside the budget) is preserved at all four sites.

Keying the arming decision on suspension rather than on Duration::ZERO is the right call, and the rationale is documented well.

No blocking findings. Suggestions are inline; two items that fall on unchanged lines are below.

Verification I ran

Fresh worktree at db45b6b8:

• cargo clippy -p mssql-tds --all-targets — clean
• cargo fmt --all -- --check — clean
• all 6 new request_timeout tests pass
• instrumented row_fetch_futures_stay_small on this branch and on origin/main to compare future sizes (see the inline comment on the doc block)
• confirmed receive_token really is not in the per-row loops (tds_client.rs get_next_row_into / next_row_cursor use receive_row_into / receive_row_header), so leaving it eager is correctly scoped

Follow-up beyond this PR (relevant to #271)

These land on unchanged lines, so noting them here rather than inline.

tds_client.rs:3631 (get_next_row_into) and tds_client.rs:3696 (next_row_cursor) still take an unconditional Instant::now() per row before every receive_row_* call, and update_remaining_timeout takes a second clock read when a budget is set. When no command timeout is configured, that first read is pure overhead on exactly the hot path this PR is optimizing. Skipping it when remaining_request_timeout.is_none() looks like a cheap next slice. Out of scope here — just flagging it.

Nit on the checklist

cargo btest is the one unchecked box in the description. The ADO validation run is green, so I assume it just needs ticking.

[Saurabh Singh (saurabh500)]

cargo btest is the one unchecked box in the description.

Ticked. The ADO validation run on 1aaf17f5 is green.

tds_client.rs:3631 (get_next_row_into) and tds_client.rs:3696 (next_row_cursor) still take an unconditional Instant::now() per row

Confirmed, and it is worse than "one extra clock read" when no timeout is configured: update_remaining_timeout only touches start inside remaining_request_timeout.map(..), so with None that read is never observed at all. Filed as #297 under #247, including the other call sites that share the shape (:3532, :3799, :3880, :3908, :2371), rather than widening this PR.

Note

Auto-replied by the GitHub Copilot app.

[David Engel (David-Engel)]

Re-review at 1aaf17f5 — looks good to me

All four actionable items are addressed and both declines are well argued. Locally on this head: cargo fmt --all -- --check clean, cargo clippy -p mssql-tds --all-targets clean, all 7 request_timeout tests pass. All PR checks green.

I would approve, but I have been in this thread as a reviewer of your measurements rather than as an owner, so treat this as a non-blocking sign-off.

On the new guard test

I verified ready_path_never_constructs_a_sleep independently rather than taking the claim on trust. Replacing the macro body with plain eager tokio::time::timeout and re-running the module: that test FAILS (panics at the Sleep construction) while the other six still pass. It isolates exactly the property the PR exists for, which is what was missing before. Nice.

Retraction

My "multi-kilobyte" finding was wrong — details in that thread. async_trait boxes the trait method, so the sizes I quoted never bounded the moved future. Re-measured at the call site: receive_row_into_internal is 1,408 B and the CancelHandle::run_until_cancelled future handed to the macro is 4,328 B, matching your numbers. Your doc wording stands.

Follow-up that came out of that measurement (not for this PR)

Decomposing the 3.07x:

future	size	vs. inner
receive_row_into_internal	1,408 B	1.00x
tokio-util CancellationToken::run_until_cancelled(inner)	2,904 B	2.06x
CancelHandle::run_until_cancelled(None, inner) (core.rs:43)	4,328 B	3.07x

The outer wrapper adds 1,424 B on top of tokio-util's — very close to one whole extra copy of inner. That matches its shape: f is live in two distinct suspend states (the Some arm's combinator and None => f.await), so the state machine reserves room for both. Collapsing it to a single await point, for example resolving to one CancellationToken up front and using a never-cancelled token when cancel_handle is None, should land near 2,904 B for the cost of an Arc bump on the no-handle path.

This is worth more than a memcpy: async_trait expands to Box::pin(async move { ... }) per call, so that 4,328 B is a per-row heap allocation, not just a per-row move. Same hot path as #271 and a larger term than the 254 ns this PR reclaims. Happy to open it as its own issue if you would rather not carry it here.

[Saurabh Singh (saurabh500)]

Happy to open it as its own issue if you would rather not carry it here.

Filed as #298 so it does not ride along here. Your diagnosis is right and the size prediction is essentially exact: a pin-once prototype (pin f before the match, pass f.as_mut() into CancellationToken::run_until_cancelled, keep the direct None => f.await arm) measured 4,280 B -> 2,912 B (-32.0%) on post-#264 main, against your predicted ~2,904 B.

The reason I did not carry it as a follow-on PR is that the shrink did not translate into time. Seven fresh processes, 31 paired ABBA pairs each, 20,000 buffered rows per pass: no handle median +12.2 ns/row (0.82% slower), live handle median +6.8 ns/row (0.43% slower), per-process range -42.7 to +75.7 ns/row with the sign flipping across processes, while the old-vs-old A/A control sat at -0.8 ns/row. The harness is tight; the wrapper swing is code layout.

One correction on the framing: shrinking the future shrinks the Box::pin allocation but does not remove it, since async_trait still allocates once per call at any size. Removing that allocation is #265. All of the above is in the issue so the next attempt starts from the measurement rather than the prediction.

Thanks for re-measuring the size decomposition and for independently verifying ready_path_never_constructs_a_sleep against an eager control rather than taking the claim on trust.

Note

Auto-replied by the GitHub Copilot app.

[David Engel (David-Engel)]

Correction to my own follow-up review above — please disregard the cancellation-wrapper section.

It is a duplicate. #265 already carries it as "Adjacent finding: ~70% of the row future is cancellation plumbing", down to the same mechanism at core.rs:50-56 (f awaited in both match arms, separate suspend points, so rustc stores both rather than overlapping) and the same arithmetic corroboration. My 1,408 / 2,904 / 4,328 B are that table re-measured on a dyn RowWriter + Send instantiation; #265 has 1,328 / 2,744 / 4,088 on DefaultRowWriter. I should have checked the issue tree before writing it up as new.

And my conclusion was backwards. I said it was a larger term than what this PR reclaims. #265's paired wrapper ladder already priced it:

Mode	median ms	vs bare
bare	32.811	—
run_until_cancelled(None, ..)	34.709	+5.8%
+ tokio::time::timeout(..)	46.751	+42.5%

So the cancellation plumbing is ~70% of the future by size but only ~5% by time, while the per-row timeout is +112 B and ~36% of the cost — the wrapper this PR removes is the expensive one. #265 says it plainly: "Size and time are anti-correlated here." I argued from allocation size against a measurement in this repo that already said not to do that. That is the same mistake as the async_trait one earlier in this review: I measured a proxy and reported it as the result.

For the record, the size does still matter, but for a reason that belongs to #265 rather than the one I gave: at ~1.4 KB the row future would fit the 4096 B row_fetch_futures_stay_small guard unboxed, which would remove the per-row allocation outright instead of merely preserving it.

No new issue needed — #265 owns this, and #297 now covers the Instant::now() item. Nothing here changes the review: #275 still looks good to me.

[David Engel (David-Engel)]

We crossed by about thirty seconds — #298 was filed while I was posting, so ignore the "no new issue needed" line in my previous comment. #298 is the right home and my comment predates it.

Your ABBA run settles it better than my retraction did. The prototype landing at 2,912 B against a predicted ~2,904 B and still measuring +12.2 / +6.8 ns/row — with the sign flipping across processes against a −0.8 ns/row A/A control — is the cleanest statement of the point: an exact model of the wrong quantity. Worth keeping in #298 so the next person does not re-derive the byte table and assume it implies time.

Taking your framing correction as well. Shrinking the future does not remove the Box::pin; async_trait allocates once per call at any size, and removal is #265. My "would fit the guard unboxed" sentence was quoting #265 about the whole row future, but placed where it was it reads as though the wrapper shrink gets you there. It does not.

Nothing changes on the PR — #275 still looks good.