chore(server): Some memory code updates (#5379)

Update memory management code

Author: dranikpg

src/server/common.cc

@@ -110,19 +110,13 @@ std::string_view LockTagOptions::Tag(std::string_view key) const {
   return key.substr(start + 1, end - start - 1);
 }
 
-atomic_uint64_t used_mem_peak(0);
+size_t max_memory_limit = 0;
 atomic_uint64_t used_mem_current(0);
 atomic_uint64_t rss_mem_current(0);
-atomic_uint64_t rss_mem_peak(0);
 
 unsigned kernel_version = 0;
-size_t max_memory_limit = 0;
 Namespaces* namespaces = nullptr;
 
-size_t FetchRssMemory(const io::StatusData& sdata) {
-  return sdata.vm_rss + sdata.hugetlb_pages;
-}
-
 const char* GlobalStateName(GlobalState s) {
   switch (s) {
     case GlobalState::ACTIVE:

src/server/common.h

@@ -128,15 +128,11 @@ bool ParseDouble(std::string_view src, double* value);
 
 const char* RdbTypeName(unsigned type);
 
-// Cached values, updated frequently to represent the correct state of the system.
-extern std::atomic_uint64_t used_mem_peak;
+extern size_t max_memory_limit;  // Value of maxmemory flag
+
+// Globally used atomics for memory readings
 extern std::atomic_uint64_t used_mem_current;
 extern std::atomic_uint64_t rss_mem_current;
-extern std::atomic_uint64_t rss_mem_peak;
-
-extern size_t max_memory_limit;
-
-size_t FetchRssMemory(const io::StatusData& sdata);
 
 extern Namespaces* namespaces;
 

src/server/db_slice.h

@@ -253,9 +253,7 @@ class DbSlice {
     expire_base_[generation & 1] = now;
   }
 
-  // From time to time DbSlice is set with a new set of params needed to estimate its
-  // memory usage.
-  void SetCachedParams(int64_t budget, size_t bytes_per_object) {
+  void UpdateMemoryParams(int64_t budget, size_t bytes_per_object) {
     memory_budget_ = budget;
     bytes_per_object_ = bytes_per_object;
   }
@@ -632,8 +630,12 @@ class DbSlice {
 
   uint64_t version_ = 1;  // Used to version entries in the PrimeTable.
   uint64_t next_moved_id_ = 1;
+
+  // Estimation of available memory dedicated to this shard.
+  // Recalculated periodically by dividing free memory left among all shards equally
   ssize_t memory_budget_ = SSIZE_MAX / 2;
   size_t bytes_per_object_ = 0;
+
   size_t soft_budget_limit_ = 0;
   size_t table_memory_ = 0;
   uint64_t entries_count_ = 0;

src/server/engine_shard.cc

@@ -476,11 +476,11 @@ uint32_t EngineShard::DefragTask() {
 }
 
 EngineShard::EngineShard(util::ProactorBase* pb, mi_heap_t* heap)
-    : queue_(kQueueLen, 1, 1),
+    : txq_([](const Transaction* t) { return t->txid(); }),
+      queue_(kQueueLen, 1, 1),
       queue2_(kQueueLen / 2, 2, 2),
-      txq_([](const Transaction* t) { return t->txid(); }),
-      mi_resource_(heap),
-      shard_id_(pb->GetPoolIndex()) {
+      shard_id_(pb->GetPoolIndex()),
+      mi_resource_(heap) {
   queue_.Start(absl::StrCat("shard_queue_", shard_id()));
   queue2_.Start(absl::StrCat("l2_queue_", shard_id()));
 }
@@ -859,30 +859,26 @@ void EngineShard::RetireExpiredAndEvict() {
 
 void EngineShard::CacheStats() {
   uint64_t now = fb2::ProactorBase::GetMonotonicTimeNs();
-  if (cache_stats_time_ + 1000000 > now)  // 1ms
+  if (last_mem_params_.updated_at + 1000000 > now)  // 1ms
     return;
 
-  cache_stats_time_ = now;
-  // Used memory for this shard.
   size_t used_mem = UsedMemory();
   DbSlice& db_slice = namespaces->GetDefaultNamespace().GetDbSlice(shard_id());
 
-  // delta can wrap if used_memory is smaller than last_cached_used_memory_ and it's fine.
-  size_t delta = used_mem - last_cached_used_memory_;
-  last_cached_used_memory_ = used_mem;
+  // Reflect local memory change on global value
+  size_t delta = used_mem - last_mem_params_.used_mem;  // negative value wraps safely
   size_t current = used_mem_current.fetch_add(delta, memory_order_relaxed) + delta;
   ssize_t free_mem = max_memory_limit - current;
 
+  // Estimate bytes per object, excluding table memory
   size_t entries = db_slice.entries_count();
-  size_t table_memory = db_slice.table_memory();
-
-  if (tiered_storage_) {
-    table_memory += tiered_storage_->CoolMemoryUsage();
-  }
+  size_t table_memory =
+      db_slice.table_memory() + (tiered_storage_ ? tiered_storage_->CoolMemoryUsage() : 0);
   size_t obj_memory = table_memory <= used_mem ? used_mem - table_memory : 0;
-
   size_t bytes_per_obj = entries > 0 ? obj_memory / entries : 0;
-  db_slice.SetCachedParams(free_mem / shard_set->size(), bytes_per_obj);
+
+  db_slice.UpdateMemoryParams(free_mem / shard_set->size(), bytes_per_obj);
+  last_mem_params_ = {now, used_mem};
 }
 
 size_t EngineShard::UsedMemory() const {

src/server/engine_shard.h

@@ -122,7 +122,7 @@ class EngineShard {
     return stats_;
   }
 
-  // Returns used memory for this shard.
+  // Calculate memory used by shard by summing multiple sources
   size_t UsedMemory() const;
 
   TieredStorage* tiered_storage() {
@@ -247,19 +247,22 @@ class EngineShard {
   // --------------------------------------------------------------------------
   uint32_t DefragTask();
 
+  TxQueue txq_;
   TaskQueue queue_, queue2_;
 
-  TxQueue txq_;
-  MiMemoryResource mi_resource_;
   ShardId shard_id_;
-
   Stats stats_;
 
   // Become passive if replica: don't automatially evict expired items.
   bool is_replica_ = false;
 
-  size_t last_cached_used_memory_ = 0;
-  uint64_t cache_stats_time_ = 0;  // monotonic, set by ProactorBase::GetMonotonicTimeNs.
+  // Precise tracking of used memory by persistent shard local values and structures
+  MiMemoryResource mi_resource_;
+
+  struct {
+    uint64_t updated_at = 0;  // from GetMonotonicTimeNs
+    size_t used_mem = 0;
+  } last_mem_params_;
 
   // Logical ts used to order distributed transactions.
   TxId committed_txid_ = 0;

src/server/server_family.cc

@@ -251,6 +251,16 @@ inline CommandId::Handler3 HandlerFunc(ServerFamily* se, EngineFunc f) {
   return [=](CmdArgList args, const CommandContext& cntx) { return (se->*f)(args, cntx); };
 }
 
+// Captured memory peaks
+struct {
+  std::atomic<size_t> used = 0;
+  std::atomic<size_t> rss = 0;
+} glob_memory_peaks;
+
+size_t FetchRssMemory(const io::StatusData& sdata) {
+  return sdata.vm_rss + sdata.hugetlb_pages;
+}
+
 using CI = CommandId;
 
 struct CmdArgListFormatter {
@@ -670,11 +680,11 @@ bool ReadProcStats(io::StatusData* sdata) {
     return false;
   }
 
-  size_t total_rss = FetchRssMemory(sdata_res.value());
+  size_t total_rss = FetchRssMemory(*sdata_res);
   rss_mem_current.store(total_rss, memory_order_relaxed);
-  if (rss_mem_peak.load(memory_order_relaxed) < total_rss) {
-    rss_mem_peak.store(total_rss, memory_order_relaxed);
-  }
+  if (total_rss > glob_memory_peaks.rss.load(memory_order_relaxed))
+    glob_memory_peaks.rss.store(total_rss, memory_order_relaxed);
+
   *sdata = *sdata_res;
   return true;
 }
@@ -1028,51 +1038,35 @@ void ServerFamily::Shutdown() {
 }
 
 bool ServerFamily::HasPrivilegedInterface() {
-  for (auto* listener : listeners_) {
-    if (listener->IsPrivilegedInterface()) {
-      return true;
-    }
-  }
-  return false;
+  return any_of(listeners_.begin(), listeners_.end(),
+                [](auto* l) { return l->IsPrivilegedInterface(); });
 }
 
 void ServerFamily::UpdateMemoryGlobalStats() {
-  ShardId sid = EngineShard::tlocal()->shard_id();
-  if (sid != 0) {  // This function is executed periodicaly on all shards. To ensure the logic
-                   // bellow runs only on one shard we return is the shard is not 0.
+  // Called from all shards, but one updates global stats below
+  if (EngineShard::tlocal()->shard_id() > 0)
     return;
-  }
 
+  // Update used memory peak
   uint64_t mem_current = used_mem_current.load(std::memory_order_relaxed);
-  if (mem_current > used_mem_peak.load(memory_order_relaxed)) {
-    used_mem_peak.store(mem_current, memory_order_relaxed);
-  }
+  if (mem_current > glob_memory_peaks.used.load(memory_order_relaxed))
+    glob_memory_peaks.used.store(mem_current, memory_order_relaxed);
 
   io::StatusData status_data;
-  bool success = ReadProcStats(&status_data);
+  bool success = ReadProcStats(&status_data);  // updates glob_memory_peaks.rss
   if (!success)
     return;
 
   size_t total_rss = FetchRssMemory(status_data);
+
+  // Decide on stopping or accepting new connections based on oom deny ratio
   double rss_oom_deny_ratio = ServerState::tlocal()->rss_oom_deny_ratio;
   if (rss_oom_deny_ratio > 0) {
     size_t memory_limit = max_memory_limit * rss_oom_deny_ratio;
-    if (total_rss > memory_limit && accepting_connections_ && HasPrivilegedInterface()) {
-      for (auto* listener : listeners_) {
-        if (!listener->IsPrivilegedInterface()) {
-          listener->socket()->proactor()->Await([listener]() { listener->pause_accepting(); });
-        }
-      }
-      accepting_connections_ = false;
-
-    } else if (total_rss < memory_limit && !accepting_connections_) {
-      for (auto* listener : listeners_) {
-        if (!listener->IsPrivilegedInterface()) {
-          listener->socket()->proactor()->Await([listener]() { listener->resume_accepting(); });
-        }
-      }
-      accepting_connections_ = true;
-    }
+    if (total_rss > memory_limit && accepting_connections_ && HasPrivilegedInterface())
+      ChangeConnectionAccept(false);
+    else if (total_rss < memory_limit && !accepting_connections_)
+      ChangeConnectionAccept(true);
   }
 }
 
@@ -1409,8 +1403,8 @@ void PrintPrometheusMetrics(uint64_t uptime, const Metrics& m, DflyCmd* dfly_cmd
   bool success = ReadProcStats(&sdata);
   AppendMetricWithoutLabels("memory_used_bytes", "", m.heap_used_bytes, MetricType::GAUGE,
                             &resp->body());
-  AppendMetricWithoutLabels("memory_used_peak_bytes", "", used_mem_peak.load(memory_order_relaxed),
-                            MetricType::GAUGE, &resp->body());
+  AppendMetricWithoutLabels("memory_used_peak_bytes", "", m.used_mem_peak, MetricType::GAUGE,
+                            &resp->body());
   AppendMetricWithoutLabels("memory_fiberstack_vms_bytes",
                             "virtual memory size used by all the fibers", m.worker_fiber_stack_size,
                             MetricType::GAUGE, &resp->body());
@@ -2061,6 +2055,14 @@ void ServerFamily::ClientUnPauseCmd(CmdArgList args, SinkReplyBuilder* builder)
   builder->SendOk();
 }
 
+void ServerFamily::ChangeConnectionAccept(bool accept) {
+  DCHECK_NE(accept, accepting_connections_);
+  auto h = accept ? &ListenerInterface::resume_accepting : &ListenerInterface::pause_accepting;
+  for (auto* listener : GetNonPriviligedListeners())
+    listener->socket()->proactor()->Await([listener, h]() { (listener->*h)(); });
+  accepting_connections_ = accept;
+}
+
 void ClientHelp(SinkReplyBuilder* builder) {
   string_view help_arr[] = {
       "CLIENT <subcommand> [<arg> [value] [opt] ...]. Subcommands are:",
@@ -2452,13 +2454,19 @@ Metrics ServerFamily::GetMetrics(Namespace* ns) const {
 
   // Update peak stats. We rely on the fact that GetMetrics is called frequently enough to
   // update peak_stats_ from it.
-  util::fb2::LockGuard lk{peak_stats_mu_};
-  UpdateMax(&peak_stats_.conn_dispatch_queue_bytes,
-            result.facade_stats.conn_stats.dispatch_queue_bytes);
-  UpdateMax(&peak_stats_.conn_read_buf_capacity, result.facade_stats.conn_stats.read_buf_capacity);
+  {
+    util::fb2::LockGuard lk{peak_stats_mu_};
+    UpdateMax(&peak_stats_.conn_dispatch_queue_bytes,
+              result.facade_stats.conn_stats.dispatch_queue_bytes);
+    UpdateMax(&peak_stats_.conn_read_buf_capacity,
+              result.facade_stats.conn_stats.read_buf_capacity);
+    result.peak_stats = peak_stats_;
+  }
 
   result.peak_stats = peak_stats_;
   result.cmd_latency_map = service_.mutable_registry()->LatencyMap();
+  result.used_mem_peak = glob_memory_peaks.used.load(memory_order_relaxed);
+  result.used_mem_rss_peak = glob_memory_peaks.rss.load(memory_order_relaxed);
 
   uint64_t delta_ms = (absl::GetCurrentTimeNanos() - start) / 1'000'000;
   if (delta_ms > 30) {
@@ -2547,9 +2555,8 @@ string ServerFamily::FormatInfoMetrics(const Metrics& m, std::string_view sectio
   auto add_mem_info = [&] {
     append("used_memory", m.heap_used_bytes);
     append("used_memory_human", HumanReadableNumBytes(m.heap_used_bytes));
-    const auto ump = used_mem_peak.load(memory_order_relaxed);
-    append("used_memory_peak", ump);
-    append("used_memory_peak_human", HumanReadableNumBytes(ump));
+    append("used_memory_peak", m.used_mem_peak);
+    append("used_memory_peak_human", HumanReadableNumBytes(m.used_mem_peak));
 
     // Virtual memory size, upper bound estimation on the RSS memory used by the fiber stacks.
     append("fibers_stack_vms", m.worker_fiber_stack_size);
@@ -2562,7 +2569,7 @@ string ServerFamily::FormatInfoMetrics(const Metrics& m, std::string_view sectio
       append("used_memory_rss", rss);
       append("used_memory_rss_human", HumanReadableNumBytes(rss));
     }
-    append("used_memory_peak_rss", rss_mem_peak.load(memory_order_relaxed));
+    append("used_memory_peak_rss", glob_memory_peaks.used.load(memory_order_relaxed));
 
     append("maxmemory", max_memory_limit);
     append("maxmemory_human", HumanReadableNumBytes(max_memory_limit));

src/server/server_family.h

@@ -87,6 +87,9 @@ struct Metrics {
 
   size_t qps = 0;
 
+  size_t used_mem_peak = 0;
+  size_t used_mem_rss_peak = 0;
+
   size_t heap_used_bytes = 0;
   size_t small_string_bytes = 0;
   uint32_t traverse_ttl_per_sec = 0;
@@ -374,14 +377,17 @@ class ServerFamily {
   void ClientPauseCmd(CmdArgList args, SinkReplyBuilder* builder, ConnectionContext* cntx);
   void ClientUnPauseCmd(CmdArgList args, SinkReplyBuilder* builder);
 
+  // Set accepting_connections_ and update listners according to it
+  void ChangeConnectionAccept(bool accept);
+
   util::fb2::Fiber snapshot_schedule_fb_;
   util::fb2::Fiber load_fiber_;
 
   Service& service_;
 
   util::AcceptServer* acceptor_ = nullptr;
   std::vector<facade::Listener*> listeners_;
-  bool accepting_connections_ = true;
+  bool accepting_connections_ = true;  // reject connections near oom
   util::ProactorBase* pb_task_ = nullptr;
 
   mutable util::fb2::Mutex replicaof_mu_, save_mu_;

src/server/tiered_storage.cc

@@ -286,7 +286,7 @@ void TieredStorage::ShardOpManager::RetireColdEntries(size_t additional_memory)
 
     // Update memory_budget directly since we know that gained bytes were released.
     // We will overwrite the budget correctly in the next Hearbeat.
-    db_slice_.SetCachedParams(gained + db_slice_.memory_budget(), db_slice_.bytes_per_object());
+    db_slice_.UpdateMemoryParams(gained + db_slice_.memory_budget(), db_slice_.bytes_per_object());
   }
 }
 

src/server/tiered_storage_test.cc

@@ -324,7 +324,8 @@ TEST_F(TieredStorageTest, MemoryPressure) {
     ThisFiber::SleepFor(500us);
   }
 
-  EXPECT_LT(used_mem_peak.load(), 20_MB);
+  auto metrics = GetMetrics();
+  EXPECT_LT(metrics.used_mem_peak, 20_MB);
 }
 
 TEST_F(TieredStorageTest, Expiry) {