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 to support subtraction.
+  // Similarly, I wish that CompoundIterator supported the -> 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:
@@ -1536,36 +1516,32 @@ void Connection::SquashPipeline() {
   DCHECK_EQ(dispatch_q_.size(), pending_pipeline_cmd_cnt_);
   DCHECK_EQ(reply_builder_->GetProtocol(), Protocol::REDIS);  // Only Redis is supported.
 
-  vector<ParsedArgs> squash_cmds;
-  squash_cmds.reserve(dispatch_q_.size());
+  unsigned pipeline_count = 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();
+  // We use indexes as iterators are invalidated when pushing into the queue.
+  auto get_next_fn = [i = 0, this]() mutable -> ParsedArgs {
+    const auto& elem = dispatch_q_[i++];
+    CHECK(holds_alternative<PipelineMessagePtr>(elem.handle));
+    const auto& pmsg = get<PipelineMessagePtr>(elem.handle);
 
-    auto& pmsg = get<PipelineMessagePtr>(msg.handle);
-    squash_cmds.emplace_back(ParsedArgs(pmsg->args));
-    if (squash_cmds.size() >= pipeline_squash_limit_cached) {
-      // We reached the limit of commands to squash, so we dispatch them.
-      break;
-    }
-  }
+    return *pmsg;
+  };
 
   // async_dispatch is a guard to prevent concurrent writes into reply_builder_, hence
   // it must guard the Flush() as well.
   cc_->async_dispatch = true;
 
   DispatchManyResult result =
-      service_->DispatchManyCommands(absl::MakeSpan(squash_cmds), reply_builder_.get(), cc_.get());
+      service_->DispatchManyCommands(get_next_fn, pipeline_count, reply_builder_.get(), cc_.get());
 
   uint32_t dispatched = result.processed;
   uint64_t before_flush = CycleClock::Now();
   //
   // TODO: to investigate if always flushing will improve P99 latency because otherwise we
   // wait for the next batch to finish before fully flushing the current response.
-  if (pending_pipeline_cmd_cnt_ == squash_cmds.size() ||
+  if (pending_pipeline_cmd_cnt_ == pipeline_count ||
       always_flush_pipeline_cached) {  // Flush if no new commands appeared
     reply_builder_->Flush();
     reply_builder_->SetBatchMode(false);  // in case the next dispatch is sync
@@ -1597,7 +1573,7 @@ void Connection::SquashPipeline() {
   dispatch_q_.erase(it, it + dispatched);
 
   // If interrupted due to pause, fall back to regular dispatch
-  skip_next_squashing_ = dispatched != squash_cmds.size();
+  skip_next_squashing_ = dispatched != pipeline_count;
 }
 
 void Connection::ClearPipelinedMessages() {
@@ -1755,25 +1731,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 +1748,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 +1758,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 +1934,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/facade/ok_main.cc

@@ -28,13 +28,15 @@ class OkService : public ServiceInterface {
     return DispatchResult::OK;
   }
 
-  DispatchManyResult DispatchManyCommands(absl::Span<ParsedArgs> args_lists,
+  DispatchManyResult DispatchManyCommands(std::function<ParsedArgs()> arg_gen, unsigned count,
                                           SinkReplyBuilder* builder,
                                           ConnectionContext* cntx) final {
-    for (auto args : args_lists)
+    for (unsigned i = 0; i < count; i++) {
+      ParsedArgs args = arg_gen();
       DispatchCommand(args, builder, cntx);
+    }
     DispatchManyResult result{
-        .processed = static_cast<uint32_t>(args_lists.size()),
+        .processed = static_cast<uint32_t>(count),
         .account_in_stats = true,
     };
     return result;

src/facade/service_interface.h

@@ -39,8 +39,8 @@ class ServiceInterface {
   virtual DispatchResult DispatchCommand(ParsedArgs args, SinkReplyBuilder* builder,
                                          ConnectionContext* cntx) = 0;
 
-  virtual DispatchManyResult DispatchManyCommands(absl::Span<ParsedArgs> commands,
-                                                  SinkReplyBuilder* builder,
+  virtual DispatchManyResult DispatchManyCommands(std::function<ParsedArgs()> arg_gen,
+                                                  unsigned count, SinkReplyBuilder* builder,
                                                   ConnectionContext* cntx) = 0;
 
   virtual void DispatchMC(const MemcacheParser::Command& cmd, std::string_view value,

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_};
 }
 

src/server/main_service.cc

@@ -1725,13 +1725,13 @@ DispatchResult Service::InvokeCmd(const CommandId* cid, CmdArgList tail_args,
   return res;
 }
 
-DispatchManyResult Service::DispatchManyCommands(absl::Span<facade::ParsedArgs> args_list,
-                                                 SinkReplyBuilder* builder,
+DispatchManyResult Service::DispatchManyCommands(std::function<facade::ParsedArgs()> arg_gen,
+                                                 unsigned count, SinkReplyBuilder* builder,
                                                  facade::ConnectionContext* cntx) {
   ConnectionContext* dfly_cntx = static_cast<ConnectionContext*>(cntx);
   DCHECK(!dfly_cntx->conn_state.exec_info.IsRunning());
   DCHECK_EQ(builder->GetProtocol(), Protocol::REDIS);
-  DCHECK_GT(args_list.size(), 1u);
+  DCHECK_GT(count, 1u);
 
   auto* ss = dfly::ServerState::tlocal();
   // Don't even start when paused. We can only continue if DispatchTracker is aware of us running.
@@ -1771,7 +1771,8 @@ DispatchManyResult Service::DispatchManyCommands(absl::Span<facade::ParsedArgs>
     stored_cmds.clear();
   };
 
-  for (const auto& args : args_list) {
+  for (unsigned i = 0; i < count; i++) {
+    ParsedArgs args = arg_gen();
     string cmd = absl::AsciiStrToUpper(args.Front());
     const auto [cid, tail_args] = registry_.FindExtended(cmd, args.Tail());
 
@@ -1789,7 +1790,7 @@ DispatchManyResult Service::DispatchManyCommands(absl::Span<facade::ParsedArgs>
     const bool is_blocking = cid != nullptr && cid->IsBlocking();
 
     if (!is_multi && !is_eval && !is_blocking && cid != nullptr) {
-      stored_cmds.reserve(args_list.size());
+      stored_cmds.reserve(count);
       stored_cmds.emplace_back(cid, tail_args);  // Shallow copy
       continue;
     }

src/server/main_service.h

@@ -40,8 +40,8 @@ class Service : public facade::ServiceInterface {
                                          facade::ConnectionContext* cntx) final;
 
   // Execute multiple consecutive commands, possibly in parallel by squashing
-  facade::DispatchManyResult DispatchManyCommands(absl::Span<facade::ParsedArgs> args_list,
-                                                  facade::SinkReplyBuilder* builder,
+  facade::DispatchManyResult DispatchManyCommands(std::function<facade::ParsedArgs()> arg_gen,
+                                                  unsigned count, facade::SinkReplyBuilder* builder,
                                                   facade::ConnectionContext* cntx) final;
 
   // Check VerifyCommandExecution and invoke command with args

src/server/test_utils.cc

@@ -486,13 +486,16 @@ void BaseFamilyTest::RunMany(const std::vector<std::vector<std::string>>& cmds)
   TestConnWrapper* conn_wrapper = AddFindConn(Protocol::REDIS, GetId());
   auto* context = conn_wrapper->cmd_cntx();
   context->ns = &namespaces->GetDefaultNamespace();
-  vector<facade::ParsedArgs> args_vec(cmds.size());
   vector<cmn::BackedArguments> backed_args_vec(cmds.size());
   for (size_t i = 0; i < cmds.size(); ++i) {
-    backed_args_vec[i] = cmn::BackedArguments(cmds[i].begin(), cmds[i].end());
-    args_vec[i] = facade::ParsedArgs{backed_args_vec[i]};
+    backed_args_vec[i] = cmn::BackedArguments(cmds[i].begin(), cmds[i].end(), cmds[i].size());
   }
-  service_->DispatchManyCommands(absl::MakeSpan(args_vec), conn_wrapper->builder(), context);
+  auto next_fn = [it = backed_args_vec.begin()]() mutable {
+    ParsedArgs args(*it);
+    ++it;
+    return args;
+  };
+  service_->DispatchManyCommands(next_fn, cmds.size(), conn_wrapper->builder(), context);
   DCHECK(context->transaction == nullptr);
 }