From cf2e94f3742985595947802ca584ac069e176150 Mon Sep 17 00:00:00 2001
From: Shahar Mike <chakaz@users.noreply.github.com>
Date: Sun, 13 Oct 2024 12:50:05 +0300
Subject: [PATCH] chore: Add `--allocator_tracker` for default tracking (#3901)

* chore: Add `--allocator_tracker` for default tracking

Before, in order to use allocation tracker, one had to issue a `MEMORY
TRACK` command. This flag is identical to that, but allows starting
Dragonfly with certain ranges without issuing a command.

While here, fix a bug. Apparently, `absl::InlinedVector<>` has a bug in
the implementation of `max_size()` and so in practice we did not limit
the number of trackers. I switched to use `capacity()` instead, which I
tested and it works well.

Notes:
1. Currently the flag always add 100% "sampling", we can extend that in
   the future if need be
2. I added the flag in `dfly_main.cc` with custom initialization,
   because it's low level, and I couldn't get it reasonably working with
   changes only to `allocation_tracker.cc`

* fixes
---
 src/core/allocation_tracker.cc |  2 +-
 src/server/dfly_main.cc        | 41 ++++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/src/core/allocation_tracker.cc b/src/core/allocation_tracker.cc
index e5a7d392c04b..e00afe2907bb 100644
--- a/src/core/allocation_tracker.cc
+++ b/src/core/allocation_tracker.cc
@@ -19,7 +19,7 @@ AllocationTracker& AllocationTracker::Get() {
 }
 
 bool AllocationTracker::Add(const TrackingInfo& info) {
-  if (tracking_.size() >= tracking_.max_size()) {
+  if (tracking_.size() >= tracking_.capacity()) {
     return false;
   }
 
diff --git a/src/server/dfly_main.cc b/src/server/dfly_main.cc
index 19f6318c9ea3..5caa7cdc1bf9 100644
--- a/src/server/dfly_main.cc
+++ b/src/server/dfly_main.cc
@@ -73,6 +73,9 @@ ABSL_FLAG(string, unixsocketperm, "", "Set permissions for unixsocket, in octal
 ABSL_FLAG(bool, force_epoll, false,
           "If true - uses linux epoll engine underneath. "
           "Can fit for kernels older than 5.10.");
+ABSL_FLAG(
+    string, allocation_tracker, "",
+    "Logs stack trace of memory allocation within these ranges. Format is min:max,min:max,....");
 
 ABSL_FLAG(bool, version_check, true,
           "If true, Will monitor for new releases on Dragonfly servers once a day.");
@@ -561,6 +564,42 @@ bool UpdateResourceLimitsIfInsideContainer(io::MemInfoData* mdata, size_t* max_t
 
 #endif
 
+void SetupAllocationTracker(ProactorPool* pool) {
+#ifdef DFLY_ENABLE_MEMORY_TRACKING
+  string flag = absl::GetFlag(FLAGS_allocation_tracker);
+  vector<pair<size_t, size_t>> track_ranges;
+  for (string_view entry : absl::StrSplit(flag, ",", absl::SkipEmpty())) {
+    auto separator = entry.find(":");
+    if (separator == entry.npos) {
+      LOG(ERROR) << "Can't find ':' in element";
+      exit(-1);
+    }
+
+    pair<size_t, size_t> p;
+    if (!absl::SimpleAtoi(entry.substr(0, separator), &p.first)) {
+      LOG(ERROR) << "Can't parse first number in pair";
+      exit(-1);
+    }
+    if (!absl::SimpleAtoi(entry.substr(separator + 1), &p.second)) {
+      LOG(ERROR) << "Can't parse second number in pair";
+      exit(-1);
+    }
+
+    track_ranges.push_back(p);
+  }
+
+  pool->AwaitBrief([&](unsigned, ProactorBase*) {
+    for (auto range : track_ranges) {
+      if (!AllocationTracker::Get().Add(
+              {.lower_bound = range.first, .upper_bound = range.second, .sample_odds = 1.0})) {
+        LOG(ERROR) << "Unable to track allocation range";
+        exit(-1);
+      }
+    }
+  });
+#endif
+}
+
 }  // namespace
 }  // namespace dfly
 
@@ -747,6 +786,8 @@ Usage: dragonfly [FLAGS]
 
   pool->Run();
 
+  SetupAllocationTracker(pool.get());
+
   AcceptServer acceptor(pool.get(), &fb2::std_malloc_resource, true);
   acceptor.set_back_log(absl::GetFlag(FLAGS_tcp_backlog));