Add try_encode_bounded: budget-checked single-pass encode entry points#320
Conversation
Writers behind a transport or storage size cap previously had to call
try_encoded_len then encode — two full tree walks — to check a caller
budget before encoding. The SizeCache already holds everything write_to
needs after compute_size, but no public entry point exploited it.
This commit adds:
- Message::try_encode_bounded(max_bytes, buf) -> Result<u32, EncodeError>
- Message::try_encode_bounded_with_cache(max_bytes, cache, buf) -> Result<u32, EncodeError>
- ViewEncode::try_encode_bounded / try_encode_bounded_with_cache (same shape)
- SizeCachePool::try_encode_bounded / try_encode_view_bounded (pooled equivalents)
- EncodeError::ExceedsBudget { len, max_bytes } — distinct from MessageTooLarge;
MessageTooLarge (> 2 GiB) takes precedence when both limits are exceeded.
The implementation runs one compute_size pass, checks both the protobuf
limit (via checked_encode_size) and max_bytes before write_to, so on Err
nothing is written to the buffer and pooled spill buffers are returned.
Returns the encoded body length on Ok for use in metrics or framing.
Closes anthropics#283.
|
All contributors have signed the CLA ✍️ ✅ |
"Nothing written on Err" matters most for a sink that already holds bytes, which is how a caller framing several messages into one buffer uses it: there a partial write corrupts the preceding message rather than leaving an obviously empty buffer. Every test started from an empty Vec, so that case was the one going unchecked. Moving the write ahead of the budget check fails this test, so it does hold the line. Also fix the changelog: six methods return the encoded length, not four, and the *_with_cache forms were introduced twice. Formatting on the view test is what cargo fmt produces; CI's check would have rejected it as it stood. No-Verification-Needed: test and changelog only; the claim is mutation-tested
iainmcgin
left a comment
There was a problem hiding this comment.
[claude code] Approving. This is a clean implementation of #283 and the risky part is right.
The thing I most wanted to check is the invariant the optimization rests on. Collapsing two size passes into one means write_to consumes the SizeCache that compute_size populated, in the same order — and a desync there corrupts wire output silently, with wrong lengths and no panic, which a round-trip test on a small message would not necessarily catch. It holds: try_encode_bounded_with_cache is the same clear → compute_size → write_to shape as the existing try_encode_with_cache, with only a local len read and an early return between the passes. Nothing touches the cache.
Also verified: checked_encode_size(...)? runs first, so MessageTooLarge correctly wins over ExceedsBudget; a max_bytes above the protobuf limit is simply never the binding constraint (no clamp needed); max_bytes == 0 on an empty message succeeds as you documented; the pool variants release(cache) unconditionally, so the spill buffer returns on both paths, and your ..._returned tests prove it by encoding again afterwards; and there is not a single as usize/as u32 cast in the implementation, so nothing truncates on the 32-bit target CI builds.
Pushed one commit (597091c):
An append-semantics test. Your doc promises "nothing reaches buf on Err", but every test starts from an empty Vec, so the case where it matters most was unpinned: a caller framing several messages into one buffer, where a partial write corrupts the preceding message rather than leaving an obviously empty buffer. The new test pre-fills the sink, triggers ExceedsBudget, and requires the buffer to be byte-identical afterwards. It is load-bearing — moving write_to ahead of the budget check fails it (along with two of yours).
Changelog: it said "All four return the encoded body length" but six do, and the *_with_cache forms were introduced twice. Reworded.
Formatting: the pool.try_encode_view_bounded(...) call in the view test needed cargo fmt; CI's fmt check would have rejected it as it stood. Nothing to do with your logic — worth running cargo fmt --all before pushing next time and it'd never surface.
Two things I deliberately left alone, noted so they read as decisions rather than oversights: the bounded path doesn't call debug_assert_two_pass, which is consistent with the existing try_encode/try_encode_with_cache (a generic EncodeSink has no cheap written-byte count); and there's no try_encode_bounded_to_vec/_to_bytes convenience, a mild asymmetry with encode_to_vec that the sink form covers. The latter is a reasonable follow-up if callers ask.
Nice work — the PR body's claims all checked out against the code, which made this a verification rather than an investigation.
iainmcgin
left a comment
There was a problem hiding this comment.
[claude code] Re-approving: the push of 597091c (append-semantics test + changelog wording + fmt) dismissed the prior approval under the stale-review rule. All checks green on the updated head; the review above stands.
Summary
Writers behind a transport or storage size cap (RPC frame limit, message-queue
cap, storage row limit) previously had to call
try_encoded_lenthenencodeto guard a caller budget before encoding — two full tree walks over the message.
The
SizeCachealready holds everythingwrite_toneeds aftercompute_size,but no public entry point exploited this.
Message::try_encode_bounded(max_bytes, buf)— single size pass, budgetcheck before
write_to; nothing written onErr; returns encoded bodylength on
OkMessage::try_encode_bounded_with_cache— same, caller-supplied cacheViewEncode::try_encode_bounded/try_encode_bounded_with_cache— viewequivalents
SizeCachePool::try_encode_bounded/try_encode_view_bounded— pooledequivalents (spill buffer returned to pool on both
OkandErr)EncodeError::ExceedsBudget { len, max_bytes }— new variant carryingthe exact encoded size and the exceeded budget; distinct from
MessageTooLarge(> 2 GiB), which takes precedence when both limitsare exceeded
Test plan
try_encode_bounded_within_budget_encodes_and_returns_len— within budget, output matches plainencodetry_encode_bounded_at_exact_limit_succeeds— budget == encoded_len succeedstry_encode_bounded_over_budget_errors_and_writes_nothing— one byte under, Err, empty buffertry_encode_bounded_zero_budget_with_empty_message_succeeds— zero-length message + zero budgettry_encode_bounded_over_protobuf_limit_errors_as_too_large— over 2 GiB → MessageTooLargeViewEncodeside (view::tests::view_try_encode_bounded_*,pool_try_encode_view_bounded_*)Closes #283.