feat: Add ReplyScope to SquashPipeline#7630
Conversation
dea237e to
17b19a0
Compare
| // Drain all iovecs (including any WriteRef into bs.view()) before this | ||
| // function returns. ~BorrowedString releases the pin afterward — by then | ||
| // the source bytes are already in the kernel. | ||
| Flush(); |
There was a problem hiding this comment.
This was totally misleading and tried to make a normal use-case look like the exception.
First, all Send functions assume that the lifetime of the passed reference (or reference object like string_view) does not extend beyond the function. This is why there is a ReplyScope at the top of the function - and when it drops it calls FinishScope that decides on it's own whether to flush or to copy based on it's allowed copy size (all borrowed refs exceeed it, so it would just flush)
There was a problem hiding this comment.
I haven't had a deep though when reviewing the PRs, but now we temporarily end up with two functions (const lvalue, rvalue). I think there would have been more elegant solutions
There was a problem hiding this comment.
I do not understand. ReplyScope is not aware of BorrowedString semantics.
Suppose BorrowedString is passed, reply scope says references can be kept, but then BorrowedString goes out of scope and pin is raised.
when ReplyScope finishes, it will reference possible invalid string.
Maybe I broke ReplyScope contract but why removing Flush is valid?
FYI, I assumed that flushing a BorrowedString does not significantly impact performance compared to sending strings themselves.
There was a problem hiding this comment.
Because you are in full control of the lifetime. I've added a comment at the top that the reply builder accepts only reference types and assumes each of them goes out if scope immediately - it never says "references can be kept" if you never tell it
There was a problem hiding this comment.
If you want to break that assumption, use ReplyScopes
There was a problem hiding this comment.
Maybe I broke ReplyScope contract but why removing Flush is valid?
Because there is a ReplyScope inside the function and when ut goes out of scope it will either flush or copy the data
There was a problem hiding this comment.
But if there is another replyscope above this function? The scope will persist, right?
There was a problem hiding this comment.
I think I can only exend the scoping, I can not disable it.
There was a problem hiding this comment.
Maybe replacing ReplyScope in this function with ScopePause - would work.
dranikpg
changed the title
PR Summary by QodoAdd ReplyScope batching to SquashPipeline for zero-copy replies Description
Diagram
High-Level Assessment
Files changed (7)
|
🤖 Augment PR SummarySummary: This PR optimizes pipelined reply emission by using a single Changes:
🤖 Was this summary useful? React with 👍 or 👎 |
Code Review by Qodo
1. Const replay still consumes
|
![qodo-free-for-open-source-projects[bot]](https://avatars.githubusercontent.com/in/1068000?s=60&v=4){kind=small}
| Apply(static_cast<const Payload&>(pl), rb); | ||
| } | ||
|
|
||
| void CapturingReplyBuilder::Apply(const Payload& pl, SinkReplyBuilder* rb) { | ||
| CaptureVisitor cv{rb}; | ||
| visit(cv, std::move(pl)); | ||
| visit(cv, pl); | ||
| } |
There was a problem hiding this comment.
1. Const replay still consumes 🐞 Bug ≡ Correctness
Reply replay under an outer ReplyScope can queue iovecs that outlive the underlying bytes because collection replay in CapturingReplyBuilder::Apply(const Payload&) still moves elements out of CollectionPayload::arr, and RedisReplyBuilderBase::SendBulkStringBorrowed(BorrowedString&&) no longer takes ownership or extends the pin lifetime (it forwards as const&). As a result, moved-from/temporary std::string buffers or BorrowedString pins may be destroyed/unpinned before Connection::SquashPipeline’s outer scope finally flushes, leaving iovecs pointing at freed/unpinned memory.
Agent Prompt
## Issue description
Fix reply building/replay so that any bytes referenced by queued iovecs remain valid until the deferred flush completes when an outer `ReplyScope` batches multiple replies. Currently, `CapturingReplyBuilder::Apply(const Payload&)` unintentionally consumes collection elements by moving out of `CollectionPayload::arr`, and `RedisReplyBuilderBase::SendBulkStringBorrowed(cmn::BorrowedString&&)` does not actually take ownership/extend pin lifetime, allowing `std::string` buffers or `BorrowedString` pins to be destroyed/unpinned before the outer scope flush sends the queued iovecs.
## Issue Context
`Connection::SquashPipeline` places multiple `cmd->SendReply()` calls under a single outer `ReplyScope`, so inner `Send...()` calls enqueue `WriteRef(...)` without flushing until that outer scope ends. In this environment, replaying a captured collection payload currently does `visit(*this, std::move(pl))` for each element, which can free underlying storage/pins early, and the `BorrowedString&&` overload now forwards as `const&` (no move, no parking), contradicting call-site expectations (see `string_family.cc` comments) and making it easy for temporaries/rvalues to unpin before the deferred flush.
## Fix Focus Areas
- src/facade/reply_capture.cc[145-175]
- src/facade/reply_builder.cc[348-367]
- src/facade/reply_builder.h[215-220]
- src/server/string_family.cc[708-715]
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
![augmentcode[bot]](https://avatars.githubusercontent.com/in/1027498?s=60&v=4){kind=small}
| // function returns. ~BorrowedString releases the pin afterward — by then | ||
| // the source bytes are already in the kernel. | ||
| Flush(); | ||
| void RedisReplyBuilderBase::SendBulkStringBorrowed(cmn::BorrowedString&& bs) { |
There was a problem hiding this comment.
RedisReplyBuilderBase::SendBulkStringBorrowed(cmn::BorrowedString&&) forwards to the const cmn::BorrowedString& overload, which can enqueue a WriteRef(bs.view()) that outlives bs when an outer SinkReplyBuilder::ReplyScope is active (nested ReplyScope won’t call FinishScope() when prev==true). This can lead to use-after-free / early pin release if callers pass temporaries or otherwise short-lived BorrowedStrings under a ReplyScope.
Severity: high
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
| Apply(static_cast<const Payload&>(pl), rb); | ||
| } | ||
|
|
||
| void CapturingReplyBuilder::Apply(const Payload& pl, SinkReplyBuilder* rb) { |
There was a problem hiding this comment.
CapturingReplyBuilder::Apply(Payload&&, SinkReplyBuilder*) now forwards to Apply(const Payload&) for non-capturing builders, so it no longer consumes/moves pl even though the signature is rvalue. This makes it easier for callers to accidentally pass a temporary payload under an active ReplyScope, leaving WriteRef iovecs pointing into a payload that is destroyed right after Apply returns.
Severity: medium
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
| for (unsigned i = 0; i < squashed && cmd; i++, cmd = cmd->next) { | ||
| DCHECK(cmd->CanReply()); | ||
| // TODO: Coroutine replies don't preserve lifetimes after running, so pause the scope | ||
| std::optional<SinkReplyBuilder::ScopePause> pause; |
There was a problem hiding this comment.
Instead of ScopePause you can add ReplyScope::Pause() / UnPause().
| @@ -56,6 +57,10 @@ class CapturingReplyBuilder : public RedisReplyBuilder { | |||
| // Send payload to builder. | |||
| static void Apply(Payload&& pl, SinkReplyBuilder* builder); | |||
There was a problem hiding this comment.
There are too many semantics now. Do we still need "r-value" interface?
| // through replay. | ||
| virtual void SendBulkStringBorrowed(cmn::BorrowedString bs); | ||
| void SendBulkStringBorrowed(const cmn::BorrowedString& bs); | ||
| virtual void SendBulkStringBorrowed(cmn::BorrowedString&& bs); |
There was a problem hiding this comment.
@dranikpg why do we need in this case virtual void SendBulkStringBorrowed(cmn::BorrowedString&& bs); with temporary value semantics?
do you agree that this function can not run under any scope?
| // Capture the borrow into the payload, extending the pin's lifetime until | ||
| // replay moves it into the real sink where it is parked across the writev. | ||
| void CapturingReplyBuilder::SendBulkStringBorrowed(cmn::BorrowedString bs) { | ||
| void CapturingReplyBuilder::SendBulkStringBorrowed(cmn::BorrowedString&& bs) { |
There was a problem hiding this comment.
string_family.cc is the only user of this function. I think this signature should be removed
and the code below should be adjusted to use const ref
void Send(StringResult&& res) const {
if (holds_alternative<std::string>(res))
return Send(get<std::string>(res));
if (holds_alternative<cmn::BorrowedString>(res)) {
// Move the BorrowedString into SendBulkStringBorrowed; the reply builder takes ownership
// of the borrow (and associated pin) and parks the pin until all the writes complete.
rb->SendBulkStringBorrowed(std::move(get<cmn::BorrowedString>(res)));
return;
}
auto fut = get<TieredStorage::TResult<std::string>>(std::move(res));
io::Result<std::string> iores = fut.Get();
if (iores.has_value())
Send(*iores);
else
Send(iores.error().message());
}
a0716f7 to
efd3bf8
Compare
Signed-off-by: Vladislav Oleshko <vlad@dragonflydb.io>
efd3bf8 to
db79cb2
Compare
Signed-off-by: Roman Gershman <romange@gmail.com>
2 participants
Send replies from
SquashPipelineunder a ReplyScope to avoid copies and use the zero-copy vectorized IO features of the sink reply builder. Deallocate commands only after all replies were sentPerformance comparision against the main branch for GET benchmark. With 10+ connections and long pipelines (50) the difference in throughput becomes 15% and more