diff --git a/src/common/backed_args.h b/src/common/backed_args.h index 16b51e8e8551..8c855e96dd23 100644 --- a/src/common/backed_args.h +++ b/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 GetInlineBuffer() { + return {storage_.begin(), std::max(kStorageCap, storage_.size())}; + } + protected: absl::InlinedVector offsets_; - // The capacity is chosen so that we allocate a fully utilized (128 bytes) block. + // See CommandContext for storage cap size limits absl::InlinedVector storage_; }; -static_assert(sizeof(BackedArguments) == 128); - template void BackedArguments::Assign(I begin, I end, size_t len) { offsets_.resize(len); size_t total_size = 0; diff --git a/src/facade/dragonfly_connection.cc b/src/facade/dragonfly_connection.cc index 34d7ab4de946..cb1cb60ef669 100644 --- a/src/facade/dragonfly_connection.cc +++ b/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(parsed_cmd_q_bytes_) + delta, 0); + DCHECK_GE(static_cast(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) { diff --git a/src/facade/dragonfly_connection.h b/src/facade/dragonfly_connection.h index a604f6eb919e..821a864373be 100644 --- a/src/facade/dragonfly_connection.h +++ b/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; diff --git a/src/facade/parsed_command.h b/src/facade/parsed_command.h index 8ea572cd8529..a51a47e43880 100644 --- a/src/facade/parsed_command.h +++ b/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 reply_; }; -#ifdef __linux__ -static_assert(sizeof(ParsedCommand) == 240); -#endif - } // namespace facade diff --git a/src/facade/reply_capture.cc b/src/facade/reply_capture.cc index de944f6b68d2..224cd3e32809 100644 --- a/src/facade/reply_capture.cc +++ b/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(rb)->SendBulkStringBorrowed(bs); } + void operator()(const payload::BulkStringRef& bs) { + static_cast(rb)->SendBulkString(bs); + } + void operator()(payload::Null) { static_cast(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; diff --git a/src/facade/reply_capture.h b/src/facade/reply_capture.h index b621dd83ab32..86610c7fe058 100644 --- a/src/facade/reply_capture.h +++ b/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 TryExtractError(const Payload& pl); + // Provide inline buffer to avoid allocations. Payload references it - track the lifetime! + void ProvideInlineBuffer(std::span 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 inline_buffer_; + // List of nested active collections that are being built. std::stack, int>> stack_; diff --git a/src/facade/reply_payload.h b/src/facade/reply_payload.h index 95501d6637fa..b5c431e4281a 100644 --- a/src/facade/reply_payload.h +++ b/src/facade/reply_payload.h @@ -22,8 +22,9 @@ using Error = std::unique_ptr>; 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; @@ -31,8 +32,8 @@ struct VerbatimString { }; using Payload = std::variant /* big struct */, cmn::BorrowedString, - std::unique_ptr>; + BulkStringRef, std::unique_ptr /* big struct */, + cmn::BorrowedString, std::unique_ptr>; #if defined(__linux__) && !defined(_LIBCPP_VERSION) static_assert(sizeof(Payload) == 40); diff --git a/src/server/conn_context.h b/src/server/conn_context.h index a80fb487d998..c0b811308dfe 100644 --- a/src/server/conn_context.h +++ b/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 diff --git a/src/server/http_api.cc b/src/server/http_api.cc index 910fbc60dcd4..a463812244b6 100644 --- a/src/server/http_api.cc +++ b/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& vs) { absl::StrAppend(&str, JsonEscape(vs->str)); } diff --git a/src/server/main_service.cc b/src/server/main_service.cc index b641e0aa8dd5..aaa1a7b6af7e 100644 --- a/src/server/main_service.cc +++ b/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; diff --git a/src/server/multi_command_squasher.cc b/src/server/multi_command_squasher.cc index 11b39dbff577..9922347c85b8 100644 --- a/src/server/multi_command_squasher.cc +++ b/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 &&