chore: use BackedArguments for implementing a pipeline message

Remove duplicate code for building backed arguments, similarly to #6171

Signed-off-by: Roman Gershman <[email protected]>

Author: romange

src/common/backed_args.h

@@ -18,7 +18,15 @@ class BackedArguments {
   BackedArguments() {
   }
 
-  template <typename I> BackedArguments(I begin, I end);
+  // Construct the arguments from iterator range.
+  // TODO: In general we could get away without the len argument,
+  // but that would require fixing base::it::CompoundIterator tu support substraction.
+  // Similarly, I wish that CompoundIterator support implement -> operator.
+  template <typename I> BackedArguments(I begin, I end, size_t len) {
+    Assign(begin, end, len);
+  }
+
+  template <typename I> void Assign(I begin, I end, size_t len);
 
   size_t HeapMemory() const {
     size_t s1 = lens_.capacity() <= kLenCap ? 0 : lens_.capacity() * sizeof(uint32_t);
@@ -57,22 +65,29 @@ class BackedArguments {
 
 static_assert(sizeof(BackedArguments) == 128);
 
-template <typename I> BackedArguments::BackedArguments(I begin, I end) {
-  lens_.reserve(end - begin);
+template <typename I> void BackedArguments::Assign(I begin, I end, size_t len) {
+  lens_.resize(len);
   size_t total_size = 0;
+  unsigned idx = 0;
   for (auto it = begin; it != end; ++it) {
-    lens_.push_back(it->size());
-    total_size += it->size() + 1;  // +1 for '\0'
+    size_t sz = (*it).size();
+    lens_[idx++] = sz;
+    total_size += sz + 1;  // +1 for '\0'
   }
   storage_.resize(total_size);
 
+  // Reclaim memory if we have too much allocated.
+  if (storage_.capacity() > kStorageCap && total_size < storage_.capacity() / 2)
+    storage_.shrink_to_fit();
+
   char* next = storage_.data();
   for (auto it = begin; it != end; ++it) {
-    if (!it->empty()) {
-      memcpy(next, it->data(), it->size());
+    size_t sz = (*it).size();
+    if (sz > 0) {
+      memcpy(next, (*it).data(), sz);
     }
-    next[it->size()] = '\0';
-    next += it->size() + 1;
+    next[sz] = '\0';
+    next += sz + 1;
   }
 }
 

src/facade/dragonfly_connection.cc

@@ -152,6 +152,10 @@ void UpdateIoBufCapacity(const io::IoBuf& io_buf, ConnectionStats* stats,
   }
 }
 
+size_t UsedMemoryInternal(const Connection::PipelineMessage& msg) {
+  return sizeof(msg) + msg.HeapMemory();
+}
+
 struct TrafficLogger {
   // protects agains closing the file while writing or data races when opening the file.
   // Also, makes sure that LogTraffic are executed atomically.
@@ -377,19 +381,6 @@ class PipelineCacheSizeTracker {
 
 thread_local PipelineCacheSizeTracker tl_pipe_cache_sz_tracker;
 
-void Connection::PipelineMessage::SetArgs(const RespVec& args) {
-  auto* next = storage.data();
-  for (size_t i = 0; i < args.size(); ++i) {
-    RespExpr::Buffer buf = args[i].GetBuf();
-    size_t s = buf.size();
-    if (s)
-      memcpy(next, buf.data(), s);
-    next[s] = '\0';
-    this->args[i] = MutableSlice(next, s);
-    next += (s + 1);
-  }
-}
-
 Connection::MCPipelineMessage::MCPipelineMessage(MemcacheParser::Command cmd_in,
                                                  std::string_view value_in)
     : cmd{std::move(cmd_in)}, value{value_in}, backing_size{0} {
@@ -425,23 +416,13 @@ Connection::MCPipelineMessage::MCPipelineMessage(MemcacheParser::Command cmd_in,
   }
 }
 
-void Connection::PipelineMessage::Reset(size_t nargs, size_t capacity) {
-  storage.resize(capacity);
-  args.resize(nargs);
-}
-
-size_t Connection::PipelineMessage::StorageCapacity() const {
-  return storage.capacity() + args.capacity();
-}
-
 size_t Connection::MessageHandle::UsedMemory() const {
   struct MessageSize {
     size_t operator()(const PubMessagePtr& msg) {
       return sizeof(PubMessage) + (msg->channel.size() + msg->message.size());
     }
     size_t operator()(const PipelineMessagePtr& msg) {
-      return sizeof(PipelineMessage) + msg->args.capacity() * sizeof(MutableSlice) +
-             msg->storage.capacity();
+      return UsedMemoryInternal(*msg);
     }
     size_t operator()(const MonitorMessage& msg) {
       return msg.capacity();
@@ -555,11 +536,10 @@ void Connection::AsyncOperations::operator()(const PubMessage& pub_msg) {
 }
 
 void Connection::AsyncOperations::operator()(Connection::PipelineMessage& msg) {
-  DVLOG(2) << "Dispatching pipeline: " << ToSV(msg.args.front());
+  DVLOG(2) << "Dispatching pipeline: " << msg.Front();
 
   ++self->local_stats_.cmds;
-  self->service_->DispatchCommand(ParsedArgs{msg.args}, self->reply_builder_.get(),
-                                  self->cc_.get());
+  self->service_->DispatchCommand(ParsedArgs{msg}, self->reply_builder_.get(), self->cc_.get());
 
   self->last_interaction_ = time(nullptr);
   self->skip_next_squashing_ = false;
@@ -644,7 +624,7 @@ Connection::Connection(Protocol protocol, util::HttpListenerBase* http_listener,
   static atomic_uint32_t next_id{1};
 
   constexpr size_t kReqSz = sizeof(Connection::PipelineMessage);
-  static_assert(kReqSz <= 256 && kReqSz >= 200);
+  static_assert(kReqSz <= 256);
 
   switch (protocol) {
     case Protocol::REDIS:
@@ -1537,16 +1517,17 @@ void Connection::SquashPipeline() {
   DCHECK_EQ(reply_builder_->GetProtocol(), Protocol::REDIS);  // Only Redis is supported.
 
   vector<ParsedArgs> squash_cmds;
-  squash_cmds.reserve(dispatch_q_.size());
+
+  squash_cmds.reserve(std::min<uint32_t>(dispatch_q_.size(), pipeline_squash_limit_cached));
 
   uint64_t start = CycleClock::Now();
 
   for (const auto& msg : dispatch_q_) {
     CHECK(holds_alternative<PipelineMessagePtr>(msg.handle))
         << msg.handle.index() << " on " << DebugInfo();
 
-    auto& pmsg = get<PipelineMessagePtr>(msg.handle);
-    squash_cmds.emplace_back(ParsedArgs(pmsg->args));
+    const auto& pmsg = get<PipelineMessagePtr>(msg.handle);
+    squash_cmds.emplace_back(*pmsg);
     if (squash_cmds.size() >= pipeline_squash_limit_cached) {
       // We reached the limit of commands to squash, so we dispatch them.
       break;
@@ -1755,25 +1736,15 @@ void Connection::AsyncFiber() {
 }
 
 Connection::PipelineMessagePtr Connection::FromArgs(const RespVec& args) {
-  DCHECK(!args.empty());
-  size_t backed_sz = 0;
-  for (const auto& arg : args) {
-    CHECK_EQ(RespExpr::STRING, arg.type);
-    backed_sz += arg.GetBuf().size() + 1;  // for '\0'
-  }
-  DCHECK(backed_sz);
-
-  static_assert(alignof(PipelineMessage) == 8);
-
   PipelineMessagePtr ptr;
-  if (ptr = GetFromPipelinePool(); ptr) {
-    ptr->Reset(args.size(), backed_sz);
-  } else {
+  if (ptr = GetFromPipelinePool(); !ptr) {
     // We must construct in place here, since there is a slice that uses memory locations
-    ptr = make_unique<PipelineMessage>(args.size(), backed_sz);
+    ptr = make_unique<PipelineMessage>();
   }
 
-  ptr->SetArgs(args);
+  auto map = [](const RespExpr& expr) { return expr.GetView(); };
+  auto range = base::it::Transform(map, base::it::Range(args.begin(), args.end()));
+  ptr->Assign(range.begin(), range.end(), args.size());
   return ptr;
 }
 
@@ -1782,7 +1753,7 @@ void Connection::ShrinkPipelinePool() {
     return;
 
   if (tl_pipe_cache_sz_tracker.CheckAndUpdateWatermark(pipeline_req_pool_.size())) {
-    stats_->pipeline_cmd_cache_bytes -= pipeline_req_pool_.back()->StorageCapacity();
+    stats_->pipeline_cmd_cache_bytes -= UsedMemoryInternal(*pipeline_req_pool_.back());
     pipeline_req_pool_.pop_back();
   }
 }
@@ -1792,7 +1763,7 @@ Connection::PipelineMessagePtr Connection::GetFromPipelinePool() {
     return nullptr;
 
   auto ptr = std::move(pipeline_req_pool_.back());
-  stats_->pipeline_cmd_cache_bytes -= ptr->StorageCapacity();
+  stats_->pipeline_cmd_cache_bytes -= UsedMemoryInternal(*ptr);
   pipeline_req_pool_.pop_back();
   return ptr;
 }
@@ -1968,7 +1939,7 @@ void Connection::RecycleMessage(MessageHandle msg) {
     pending_pipeline_cmd_cnt_--;
     pending_pipeline_bytes_ -= used_mem;
     if (stats_->pipeline_cmd_cache_bytes < qbp.pipeline_cache_limit) {
-      stats_->pipeline_cmd_cache_bytes += (*pipe)->StorageCapacity();
+      stats_->pipeline_cmd_cache_bytes += UsedMemoryInternal(*(*pipe));
       pipeline_req_pool_.push_back(std::move(*pipe));
     }
   }

src/facade/dragonfly_connection.h

@@ -13,6 +13,7 @@
 #include <utility>
 #include <variant>
 
+#include "common/backed_args.h"
 #include "facade/acl_commands_def.h"
 #include "facade/facade_types.h"
 #include "facade/memcache_parser.h"
@@ -80,23 +81,7 @@ class Connection : public util::Connection {
   };
 
   // Pipeline message, accumulated Redis command to be executed.
-  struct PipelineMessage {
-    PipelineMessage(size_t nargs, size_t capacity) : args(nargs), storage(capacity) {
-    }
-
-    void Reset(size_t nargs, size_t capacity);
-
-    void SetArgs(const RespVec& args);
-
-    size_t StorageCapacity() const;
-
-    // mi_stl_allocator uses mi heap internally.
-    // The capacity is chosen so that we allocate a fully utilized (256 bytes) block.
-    using StorageType = absl::InlinedVector<char, kReqStorageSize>;
-
-    absl::InlinedVector<std::string_view, 6> args;
-    StorageType storage;
-  };
+  using PipelineMessage = cmn::BackedArguments;
 
   // Pipeline message, accumulated Memcached command to be executed.
   struct MCPipelineMessage {

src/server/conn_context.cc

@@ -36,7 +36,7 @@ static void SendSubscriptionChangedResponse(string_view action, std::optional<st
 
 StoredCmd::StoredCmd(const CommandId* cid, facade::ArgSlice args, facade::ReplyMode mode)
     : cid_{cid}, args_{args}, reply_mode_{mode} {
-  backed_ = std::make_unique<cmn::BackedArguments>(args.begin(), args.end());
+  backed_ = std::make_unique<cmn::BackedArguments>(args.begin(), args.end(), args.size());
   args_ = facade::ParsedArgs{*backed_};
 }