feat(memory_cmd): Add WITHOUTKEY option for the MEMORY USAGE command. SECOND PR (#5068)

feat(memory_cmd): Add WITHOUTKEY option for the MEMORY USAGE command

Signed-off-by: Stepan Bagritsevich <stefan@dragonflydb.io>
This commit is contained in:
Stepan Bagritsevich 2025-05-06 10:44:16 +02:00 committed by GitHub
parent a8b19c9b88
commit 3da7e49712
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 30 additions and 27 deletions

View file

@ -82,8 +82,8 @@ std::string MallocStatsCb(bool backing, unsigned tid) {
return str; return str;
} }
size_t MemoryUsage(PrimeIterator it) { size_t MemoryUsage(PrimeIterator it, bool account_key_memory_usage) {
size_t key_size = it->first.MallocUsed(); size_t key_size = account_key_memory_usage ? it->first.MallocUsed() : 0;
return key_size + it->second.MallocUsed(true); return key_size + it->second.MallocUsed(true);
} }
@ -95,9 +95,9 @@ MemoryCmd::MemoryCmd(ServerFamily* owner, facade::SinkReplyBuilder* builder,
} }
void MemoryCmd::Run(CmdArgList args) { void MemoryCmd::Run(CmdArgList args) {
string sub_cmd = absl::AsciiStrToUpper(ArgS(args, 0)); CmdArgParser parser(args);
if (sub_cmd == "HELP") { if (parser.Check("HELP")) {
string_view help_arr[] = { string_view help_arr[] = {
"MEMORY <subcommand> [<arg> ...]. Subcommands are:", "MEMORY <subcommand> [<arg> ...]. Subcommands are:",
"STATS", "STATS",
@ -110,8 +110,9 @@ void MemoryCmd::Run(CmdArgList args) {
"ARENA SHOW", "ARENA SHOW",
" Prints the arena summary report for the entire process.", " Prints the arena summary report for the entire process.",
" Requires MIMALLOC_VERBOSE=1 environment to be set. The output goes to stdout", " Requires MIMALLOC_VERBOSE=1 environment to be set. The output goes to stdout",
"USAGE <key>", "USAGE <key> [WITHOUTKEY]",
" Show memory usage of a key.", " Show memory usage of a key.",
" If WITHOUTKEY is specified, the key itself is not accounted.",
"DECOMMIT", "DECOMMIT",
" Force decommit the memory freed by the server back to OS.", " Force decommit the memory freed by the server back to OS.",
"TRACK", "TRACK",
@ -137,35 +138,36 @@ void MemoryCmd::Run(CmdArgList args) {
return rb->SendSimpleStrArr(help_arr); return rb->SendSimpleStrArr(help_arr);
}; };
if (sub_cmd == "STATS") { if (parser.Check("STATS")) {
return Stats(); return Stats();
} }
if (sub_cmd == "USAGE" && args.size() > 1) { if (parser.Check("USAGE") && args.size() > 1) {
string_view key = ArgS(args, 1); string_view key = parser.Next();
return Usage(key); bool account_key_memory_usage = !parser.Check("WITHOUTKEY");
return Usage(key, account_key_memory_usage);
} }
if (sub_cmd == "DECOMMIT") { if (parser.Check("DECOMMIT")) {
shard_set->pool()->AwaitBrief( shard_set->pool()->AwaitBrief(
[](unsigned, auto* pb) { ServerState::tlocal()->DecommitMemory(ServerState::kAllMemory); }); [](unsigned, auto* pb) { ServerState::tlocal()->DecommitMemory(ServerState::kAllMemory); });
return builder_->SendSimpleString("OK"); return builder_->SendSimpleString("OK");
} }
if (sub_cmd == "MALLOC-STATS") { if (parser.Check("MALLOC-STATS")) {
return MallocStats(); return MallocStats();
} }
if (sub_cmd == "ARENA") { if (parser.Check("ARENA")) {
return ArenaStats(args); return ArenaStats(args);
} }
if (sub_cmd == "TRACK") { if (parser.Check("TRACK")) {
args.remove_prefix(1); args.remove_prefix(1);
return Track(args); return Track(args);
} }
if (sub_cmd == "DEFRAGMENT") { if (parser.Check("DEFRAGMENT")) {
shard_set->pool()->DispatchOnAll([](util::ProactorBase*) { shard_set->pool()->DispatchOnAll([](util::ProactorBase*) {
if (auto* shard = EngineShard::tlocal(); shard) if (auto* shard = EngineShard::tlocal(); shard)
shard->ForceDefrag(); shard->ForceDefrag();
@ -173,7 +175,7 @@ void MemoryCmd::Run(CmdArgList args) {
return builder_->SendSimpleString("OK"); return builder_->SendSimpleString("OK");
} }
string err = UnknownSubCmd(sub_cmd, "MEMORY"); string err = UnknownSubCmd(parser.Next(), "MEMORY");
return builder_->SendError(err, kSyntaxErrType); return builder_->SendError(err, kSyntaxErrType);
} }
@ -346,18 +348,19 @@ void MemoryCmd::ArenaStats(CmdArgList args) {
return rb->SendVerbatimString(mi_malloc_info); return rb->SendVerbatimString(mi_malloc_info);
} }
void MemoryCmd::Usage(std::string_view key) { void MemoryCmd::Usage(std::string_view key, bool account_key_memory_usage) {
ShardId sid = Shard(key, shard_set->size()); ShardId sid = Shard(key, shard_set->size());
ssize_t memory_usage = shard_set->pool()->at(sid)->AwaitBrief([key, this, sid]() -> ssize_t { ssize_t memory_usage = shard_set->pool()->at(sid)->AwaitBrief(
auto& db_slice = cntx_->ns->GetDbSlice(sid); [key, account_key_memory_usage, this, sid]() -> ssize_t {
auto [pt, exp_t] = db_slice.GetTables(cntx_->db_index()); auto& db_slice = cntx_->ns->GetDbSlice(sid);
PrimeIterator it = pt->Find(key); auto [pt, exp_t] = db_slice.GetTables(cntx_->db_index());
if (IsValid(it)) { PrimeIterator it = pt->Find(key);
return MemoryUsage(it); if (IsValid(it)) {
} else { return MemoryUsage(it, account_key_memory_usage);
return -1; } else {
} return -1;
}); }
});
auto* rb = static_cast<RedisReplyBuilder*>(builder_); auto* rb = static_cast<RedisReplyBuilder*>(builder_);
if (memory_usage < 0) if (memory_usage < 0)

View file

@ -20,7 +20,7 @@ class MemoryCmd {
void Stats(); void Stats();
void MallocStats(); void MallocStats();
void ArenaStats(CmdArgList args); void ArenaStats(CmdArgList args);
void Usage(std::string_view key); void Usage(std::string_view key, bool account_key_memory_usage);
void Track(CmdArgList args); void Track(CmdArgList args);
ConnectionContext* cntx_; ConnectionContext* cntx_;