feat: Add ReplyScope to SquashPipeline

src/facade/dragonfly_connection.cc

@@ -1796,17 +1796,24 @@ void Connection::SquashPipeline() {
     return;
   }
 
-  // Send the deferred replies in linked-list order, then release the commands. SendReply() may
-  // preempt, allowing the producer to append new commands, so re-read the next pointer after it
-  // and advance the list incrementally to keep it consistent at any preemption point.
+  // Send all replies under a ReplyScope before releasing the commands.
+  // This allows the reply builder to flush without copies
+  {
+    SinkReplyBuilder::ReplyScope scope(reply_builder_.get());
+    auto* cmd = parsed_head_;
+    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;
+      if (cmd->IsSuspendedReply())
+        pause.emplace(reply_builder_.get());
+      cmd->SendReply();
+      conn_stats.pipelined_wait_latency += CycleClock::ToUsec(start - cmd->parsed_cycle);
+    }
+  }
+
   for (unsigned i = 0; i < squashed && parsed_head_; ++i) {
     auto* current = parsed_head_;
-
-    DCHECK(current->CanReply());
-    current->SendReply();
-
-    conn_stats.pipelined_wait_latency += CycleClock::ToUsec(start - current->parsed_cycle);
-
     auto* next = current->next;
     ReleasePipelinedCommand(current);
     AdvanceParsedHead(next);

src/facade/parsed_command.cc

@@ -115,8 +115,8 @@ bool ParsedCommand::CanReply() const {
 }
 
 void ParsedCommand::SendReply() {
-  auto payload_handler = [this](payload::Payload& pl) {
-    CapturingReplyBuilder::Apply(std::move(pl), rb_);
+  auto payload_handler = [this](const payload::Payload& pl) {
+    CapturingReplyBuilder::Apply(pl, rb_);
   };
   auto task_handler = [](SuspendedCommand& task) {
     DCHECK(task.coro);

src/facade/parsed_command.h

@@ -154,6 +154,11 @@ class ParsedCommand : public cmn::BackedArguments {
     return std::get<SuspendedCommand>(reply_).blocker;
   }
 
+  // Whether this command will resume a coroutine on SendReply().
+  bool IsSuspendedReply() const {
+    return std::holds_alternative<SuspendedCommand>(reply_);
+  }
+
   // Requires CanReply(). Either sends stored reply or wakes up coroutine to send reply
   void SendReply();
 

src/facade/reply_builder.cc

@@ -15,6 +15,7 @@
 #include "absl/types/span.h"
 #include "base/cycle_clock.h"
 #include "base/logging.h"
+#include "common/borrowed_string.h"
 #include "facade/error.h"
 #include "util/fibers/proactor_base.h"
 
@@ -344,7 +345,7 @@ void RedisReplyBuilderBase::SendBulkString(std::string_view str) {
   WritePieces(kCRLF);
 }
 
-void RedisReplyBuilderBase::SendBulkStringBorrowed(cmn::BorrowedString bs) {
+void RedisReplyBuilderBase::SendBulkStringBorrowed(const cmn::BorrowedString& bs) {
   ReplyScope scope(this);
   tl_facade_stats->reply_stats.borrowed_string_sent_cnt++;
 
@@ -359,11 +360,10 @@ void RedisReplyBuilderBase::SendBulkStringBorrowed(cmn::BorrowedString bs) {
     WriteRef(bs.view());
   }
   WritePieces(kCRLF);
+}
 
-  // 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();
+void RedisReplyBuilderBase::SendBulkStringBorrowed(cmn::BorrowedString&& bs) {
+  SendBulkStringBorrowed(static_cast<const cmn::BorrowedString&>(bs));
 }
 
 void RedisReplyBuilderBase::SendLong(long val) {

src/facade/reply_builder.h

@@ -14,12 +14,19 @@
 #include "facade/facade_types.h"
 #include "io/io.h"
 
+namespace cmn {
+class BorrowedString;
+}  // namespace cmn
+
 namespace facade {
 
 enum class RespVersion { kResp2, kResp3 };
 
 // Base class for all reply builders. Offer a simple high level interface for controlling output
 // modes and sending basic response types.
+// By default, pointer validity for reference types (string_view, const BorrowedString&) is only
+// assumed for the Send...() function call. Use ReplyScope if the lifetime is longer to avoid copies
+// and enable zero-copy vectorized IO.
 class SinkReplyBuilder {
   struct GuardBase {
     bool prev;
@@ -59,6 +66,17 @@ class SinkReplyBuilder {
     ~ReplyScope();
   };
 
+  // Temporarily pauses an active ReplyScope. While paused, Send calls copy data as usual
+  // (no zero-copy refs). Restores the previous scoped_ state on destruction.
+  struct ScopePause : GuardBase {
+    explicit ScopePause(SinkReplyBuilder* rb) : GuardBase{std::exchange(rb->scoped_, false), rb} {
+    }
+
+    ~ScopePause() {
+      rb->scoped_ = prev;
+    }
+  };
+
   // Aggregator reduces the number of raw send calls by copying data in an intermediate buffer.
   // Prefer ReplyScope if possible to additionally reduce the number of copies.
   struct ReplyAggregator : GuardBase {
@@ -197,15 +215,11 @@ class RedisReplyBuilderBase : public SinkReplyBuilder {
   void SendSimpleString(std::string_view str) override;
   virtual void SendBulkString(std::string_view str);  // RESP: Blob String
 
-  // Like SendBulkString, but takes ownership of a move-only cmn::BorrowedString
-  // whose bytes are borrowed from a stable in-memory source (e.g. a CompactObj
-  // LargeString under the read-only invariant). The default impl handles both
-  // raw (iovec ref) and packed (chunked decode into scratch) encodings and
-  // Flushes before returning, so ~BorrowedString releases the pin only after
-  // the source bytes are in the kernel. CapturingReplyBuilder overrides this
-  // to move `bs` into the captured payload, extending the pin's lifetime
-  // through replay.
-  virtual void SendBulkStringBorrowed(cmn::BorrowedString bs);
+  void SendBulkStringBorrowed(const cmn::BorrowedString& bs);
+
+  // The interface exposes only the rvalue function to allow squashing to "steal" the value,
+  // the real builder implmenetation forwards it to the constref version
+  virtual void SendBulkStringBorrowed(cmn::BorrowedString&& bs);
 
   void SendLong(long val) override;
   virtual void SendDouble(double val);  // RESP: Number

src/facade/reply_capture.cc

@@ -56,7 +56,7 @@ void CapturingReplyBuilder::SendBulkString(std::string_view str) {
 
 // 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) {
   SKIP_LESS(ReplyMode::FULL);
   Capture(std::move(bs));
 }
@@ -130,8 +130,8 @@ struct CaptureVisitor {
     static_cast<RedisReplyBuilder*>(rb)->SendBulkString(bs);
   }
 
-  void operator()(cmn::BorrowedString bs) {
-    static_cast<RedisReplyBuilder*>(rb)->SendBulkStringBorrowed(std::move(bs));
+  void operator()(const cmn::BorrowedString& bs) {
+    static_cast<RedisReplyBuilder*>(rb)->SendBulkStringBorrowed(bs);
   }
 
   void operator()(payload::Null) {
@@ -166,8 +166,12 @@ void CapturingReplyBuilder::Apply(Payload&& pl, SinkReplyBuilder* rb) {
     return;
   }
 
+  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);
 }
 
 void CapturingReplyBuilder::SetReplyMode(ReplyMode mode) {

src/facade/reply_capture.h

@@ -31,7 +31,8 @@ class CapturingReplyBuilder : public RedisReplyBuilder {
   void SendDouble(double val) override;
   void SendSimpleString(std::string_view str) override;
   void SendBulkString(std::string_view str) override;
-  void SendBulkStringBorrowed(cmn::BorrowedString bs) override;
+
+  void SendBulkStringBorrowed(cmn::BorrowedString&& bs) override;
 
   void StartCollection(unsigned len, CollectionType type) override;
   void SendNullArray() override;
@@ -56,6 +57,10 @@ class CapturingReplyBuilder : public RedisReplyBuilder {
   // Send payload to builder.
   static void Apply(Payload&& pl, SinkReplyBuilder* builder);
 
+  // Send payload to builder without consuming it. String refs into the payload remain valid,
+  // making this safe to use under a ReplyScope when the payload outlives the scope.
+  static void Apply(const Payload& pl, SinkReplyBuilder* builder);
+
   // If an error is stored inside payload, get a reference to it.
   static std::optional<ErrorRef> TryExtractError(const Payload& pl);
 

tests/dragonfly/connection_test.py

@@ -732,6 +732,37 @@ async def measure(aw):
     assert await measure(async_client.ft("i1").search("*")) <= 2
 
 
+@dfly_args({"proactor_threads": "1", "pipeline_squash": 1})
+async def test_squashed_reply_count(async_client: aioredis.Redis):
+    """Akin to test_reply_count but with pipelines"""
+
+    await async_client.ping()
+
+    # Create some keys to read
+    await async_client.set("long-str", "a" * 256)
+    await async_client.lpush("long-list", *(f"{i}" for i in range(100)))
+
+    pre = int((await async_client.info("stats"))["total_writes_processed"])
+    before = int((await async_client.info("stats"))["total_writes_processed"])
+    info_diff = before - pre
+
+    # Send a large pipeline of simple commands that will be squashed together.
+    p = async_client.pipeline(transaction=False)
+    for _ in range(10):
+        p.set("some-keys", "some-values")
+    for _ in range(50):
+        p.get("long-str")
+    for _ in range(50):
+        p.lrange("long-list", 0, -1)
+    await p.execute()
+
+    after = int((await async_client.info("stats"))["total_writes_processed"])
+    writes = after - before
+
+    # 90% of the commands don't cause writes. The replies should have been squashed
+    assert writes - info_diff < 11
+
+
 @dfly_args({"enable_resp_io_loop_v2": "true"})
 async def test_v2_conditional_flush_no_stall(df_server: DflyInstance):
     """Verify that V2 conditional flushing never stalls client replies.