feat: Inline buffer for captured replies

src/common/backed_args.h

@@ -14,7 +14,7 @@ namespace cmn {
 
 class BackedArguments {
   constexpr static size_t kLenCap = 5;
-  constexpr static size_t kStorageCap = 88;
+  constexpr static size_t kStorageCap = 128;
 
   constexpr static size_t kShrinkFloor = 64 << 10;  // 64 KiB
 
@@ -130,15 +130,19 @@ class BackedArguments {
     storage_.resize(last_offs);
   }
 
+  // Inlined buffer where the command arguments are stored.
+  // Can be used as a local buffer to store captured replies
+  std::span<char> GetInlineBuffer() {
+    return {storage_.begin(), std::max(kStorageCap, storage_.size())};
+  }
+
  protected:
   absl::InlinedVector<uint32_t, kLenCap> offsets_;
 
-  // The capacity is chosen so that we allocate a fully utilized (128 bytes) block.
+  // See CommandContext for storage cap size limits
   absl::InlinedVector<char, kStorageCap> storage_;
 };
 
-static_assert(sizeof(BackedArguments) == 128);
-
 template <typename I> void BackedArguments::Assign(I begin, I end, size_t len) {
   offsets_.resize(len);
   size_t total_size = 0;

src/facade/dragonfly_connection.cc

@@ -750,9 +750,6 @@ Connection::Connection(Protocol protocol, util::HttpListenerBase* http_listener,
       ssl_ctx_(ctx),
       service_(service),
       flags_(0) {
-  constexpr size_t kReqSz = sizeof(ParsedCommand);
-  static_assert(kReqSz <= 256);
-
   // TODO: to move parser initialization to where we initialize the reply builder.
   switch (protocol) {
     case Protocol::REDIS:
@@ -2980,7 +2977,7 @@ void Connection::EnqueueParsedCommand(ParsedCommand* cmd) {
   }
   parsed_tail_ = cmd;
 
-  size_t used_mem = cmd->EnqueuedBytes();
+  size_t used_mem = cmd->UsedMemory();
   parsed_cmd_q_len_++;
   dispatch_waiting_count_++;  // the newly appended tail command is not yet dispatched
   parsed_cmd_q_bytes_ += used_mem;
@@ -3013,7 +3010,7 @@ void Connection::ReleasePipelinedCommand(ParsedCommand* cmd) {
 }
 
 void Connection::ReleaseParsedCommand(ParsedCommand* cmd) {
-  size_t used_mem = cmd->EnqueuedBytes();
+  size_t used_mem = cmd->UsedMemory();
   auto& conn_stats = tl_facade_stats->conn_stats;
 
   DCHECK_GT(parsed_cmd_q_len_, 0u);
@@ -3042,6 +3039,16 @@ void Connection::ReleaseParsedCommand(ParsedCommand* cmd) {
   }
 }
 
+void Connection::AdjustParsedCmdBytes(ssize_t delta) {
+  if (parsed_cmd_q_bytes_ == 0)
+    return;  // command dispatched synchronously, not in pipeline queue
+  auto& conn_stats = tl_facade_stats->conn_stats;
+  DCHECK_GE(static_cast<ssize_t>(parsed_cmd_q_bytes_) + delta, 0);
+  DCHECK_GE(static_cast<ssize_t>(conn_stats.pipeline_queue_bytes) + delta, 0);
+  parsed_cmd_q_bytes_ += delta;
+  conn_stats.pipeline_queue_bytes += delta;
+}
+
 void Connection::DestroyParsedQueue() {
   ConnectionMemoryTracker memory_tracker(this);
   while (parsed_head_ != nullptr) {

src/facade/dragonfly_connection.h

@@ -333,6 +333,11 @@ class Connection : public util::Connection {
 
   void MarkForClose();
 
+  // Adjusts the pipeline queue byte counter by `delta`.
+  // Used when command dispatch mutates a ParsedCommand's backing storage
+  // (e.g. SwapArgs during MULTI/EXEC collection).
+  void AdjustParsedCmdBytes(ssize_t delta);
+
  protected:
   void OnShutdown() override;
   void OnPreMigrateThread() override;

src/facade/parsed_command.h

@@ -106,13 +106,8 @@ class ParsedCommand : public cmn::BackedArguments {
     return sz;
   }
 
-  size_t EnqueuedBytes() const {
-    return enqueued_bytes_;
-  }
-
   void FinalizeParsing() {
     parsed_cycle = base::CycleClock::Now();
-    enqueued_bytes_ = UsedMemory();
   }
 
   // Marks this command as having reply stored in its payload instead of being sent directly.
@@ -207,13 +202,7 @@ class ParsedCommand : public cmn::BackedArguments {
   // otherwise, moved asynchronously into reply_payload_
   bool is_deferred_reply_ = false;
 
-  size_t enqueued_bytes_ = 0;
-
   std::variant<payload::Payload, SuspendedCommand> reply_;
 };
 
-#ifdef __linux__
-static_assert(sizeof(ParsedCommand) == 240);
-#endif
-
 }  // namespace facade

src/facade/reply_capture.cc

@@ -5,6 +5,7 @@
 
 #include "absl/types/span.h"
 #include "base/logging.h"
+#include "facade/reply_payload.h"
 #include "reply_capture.h"
 
 #define SKIP_LESS(needed)     \
@@ -51,7 +52,12 @@ void CapturingReplyBuilder::SendSimpleString(std::string_view str) {
 
 void CapturingReplyBuilder::SendBulkString(std::string_view str) {
   SKIP_LESS(ReplyMode::FULL);
-  Capture(BulkString{string{str}});
+  if (str.size() < 12 || str.size() > inline_buffer_.size())
+    return Capture(BulkString{std::string{str}});
+
+  memcpy(inline_buffer_.data(), str.data(), str.size());
+  Capture(BulkStringRef{std::string_view{inline_buffer_.data(), str.size()}});
+  inline_buffer_ = inline_buffer_.subspan(str.size());
 }
 
 void CapturingReplyBuilder::SendVerbatimString(std::string_view str, VerbatimFormat format) {
@@ -144,6 +150,10 @@ struct CaptureVisitor {
     static_cast<RedisReplyBuilder*>(rb)->SendBulkStringBorrowed(bs);
   }
 
+  void operator()(const payload::BulkStringRef& bs) {
+    static_cast<RedisReplyBuilder*>(rb)->SendBulkString(bs);
+  }
+
   void operator()(payload::Null) {
     static_cast<RedisReplyBuilder*>(rb)->SendNull();
   }
@@ -164,7 +174,7 @@ struct CaptureVisitor {
     }
     builder->StartCollection(cp->len, cp->type);
     for (auto& pl : cp->arr)
-      visit(*this, std::move(pl));
+      visit(*this, pl);
   }
 
   SinkReplyBuilder* rb;

src/facade/reply_capture.h

@@ -65,6 +65,11 @@ class CapturingReplyBuilder : public RedisReplyBuilder {
   // If an error is stored inside payload, get a reference to it.
   static std::optional<ErrorRef> TryExtractError(const Payload& pl);
 
+  // Provide inline buffer to avoid allocations. Payload references it - track the lifetime!
+  void ProvideInlineBuffer(std::span<char> buf) {
+    inline_buffer_ = buf;
+  }
+
  private:
   // Send payload directly, bypassing external interface. For efficient passing between two
   // captures.
@@ -78,6 +83,9 @@ class CapturingReplyBuilder : public RedisReplyBuilder {
 
   ReplyMode reply_mode_;
 
+  // A buffer that can be used to avoid allocations. The lifetime is guaranteed by the user
+  std::span<char> inline_buffer_;
+
   // List of nested active collections that are being built.
   std::stack<std::pair<std::unique_ptr<payload::CollectionPayload>, int>> stack_;
 

src/facade/reply_payload.h

@@ -22,17 +22,18 @@ using Error = std::unique_ptr<std::pair<std::string, std::string>>;
 using Null = std::nullptr_t;  // SendNull or SendNullArray
 
 struct CollectionPayload;
-struct SimpleString : public std::string {};  // SendSimpleString
-struct BulkString : public std::string {};    // SendBulkString
+struct SimpleString : public std::string {};        // SendSimpleString
+struct BulkString : public std::string {};          // SendBulkString
+struct BulkStringRef : public std::string_view {};  // SendBulkString with guaranteed lifetime
 
 struct VerbatimString {
   std::string str;
   uint8_t format;  // RedisReplyBuilder::VerbatimFormat avoided include
 };
 
 using Payload = std::variant<std::monostate, Null, Error, long, double, SimpleString, BulkString,
-                             std::unique_ptr<VerbatimString> /* big struct */, cmn::BorrowedString,
-                             std::unique_ptr<CollectionPayload>>;
+                             BulkStringRef, std::unique_ptr<VerbatimString> /* big struct */,
+                             cmn::BorrowedString, std::unique_ptr<CollectionPayload>>;
 
 #if defined(__linux__) && !defined(_LIBCPP_VERSION)
 static_assert(sizeof(Payload) == 40);

src/server/conn_context.h

@@ -462,4 +462,8 @@ class CommandContext : public facade::ParsedCommand {
   const CommandId* cid_ = nullptr;
 };
 
+// 320 is a mi_good_size boundary.
+// The previous boundary of 256 would require making backed_args buffers much smaller
+static_assert(sizeof(CommandContext) == 320);
+
 }  // namespace dfly

src/server/http_api.cc

@@ -9,6 +9,7 @@
 #include "core/flatbuffers.h"
 #include "facade/conn_context.h"
 #include "facade/reply_capture.h"
+#include "facade/reply_payload.h"
 #include "server/conn_context.h"
 #include "server/main_service.h"
 #include "util/http/http_common.h"
@@ -120,6 +121,10 @@ struct CaptureVisitor {
     absl::StrAppend(&str, JsonEscape(bs));
   }
 
+  void operator()(const payload::BulkStringRef& bs) {
+    absl::StrAppend(&str, JsonEscape(bs));
+  }
+
   void operator()(const unique_ptr<payload::VerbatimString>& vs) {
     absl::StrAppend(&str, JsonEscape(vs->str));
   }

src/server/main_service.cc

@@ -88,7 +88,7 @@ ABSL_FLAG(bool, multi_exec_squash, true,
 
 ABSL_FLAG(bool, lua_resp2_legacy_float, false,
           "Return rounded down integers instead of floats for lua scripts with RESP2");
-ABSL_FLAG(uint32_t, multi_eval_squash_buffer, 4096, "Max buffer for squashed commands per script");
+ABSL_FLAG(uint32_t, multi_eval_squash_buffer, 8096, "Max buffer for squashed commands per script");
 
 ABSL_DECLARE_FLAG(bool, primary_port_http_enabled);
 ABSL_FLAG(size_t, listpack_max_field_len, 64,
@@ -883,8 +883,15 @@ void StoreInMultiBlock(ConnectionContext* dfly_cntx, const CommandId* cid,
   auto& exec_info = dfly_cntx->conn_state.exec_info;
   const size_t old_size = exec_info.GetStoredCmdBytes();
 
-  // Moves arguments from parsed_cmd to body.
+  // SwapArgs inside StoredCmd moves heap storage out of parsed_cmd,
+  // so we must adjust the pipeline queue accounting for the delta.
+  size_t before = parsed_cmd->UsedMemory();
   exec_info.body.emplace_back(cid, parsed_cmd, tail_index);
+  size_t after = parsed_cmd->UsedMemory();
+  if (before != after) {
+    dfly_cntx->conn()->AdjustParsedCmdBytes(ssize_t(after) - ssize_t(before));
+  }
+
   exec_info.stored_cmd_bytes += exec_info.body.back().UsedMemory();
   exec_info.is_write |= cid->IsJournaled();
   ServerState::tlocal()->stats.stored_cmd_bytes += exec_info.GetStoredCmdBytes() - old_size;

src/server/multi_command_squasher.cc

@@ -219,6 +219,13 @@ OpStatus MultiCommandSquasher::SquashedHopCb(EngineShard* es, RespVersion resp_v
     auto* ctx = &local_cntx;
     crb.SetReplyMode(dispatched.reply_mode);
 
+    // Allow captured replies to be stored in argument storage.
+    // Some commands might include arguments in replies, so we have a limited set
+    if (dispatched.cmd_cntx && dispatched.cid->SupportsAsync())
+      crb.ProvideInlineBuffer(dispatched.cmd_cntx->GetInlineBuffer());
+    else
+      crb.ProvideInlineBuffer({});  // reset buffer
+
     // With tiered storage enabled, it makes sense to dispatch async commands concurrently
     // to allow concurrent disk operations. Tiered futures are only blocked on during replies
     bool do_async = es->tiered_storage() && !IsAtomic() && opts_.pipeline_mode &&