feat(lua): Add Lua stats to INFO STATS and /metrics (#3008)

* feat(lua): Add Lua stats to `INFO STATS` and `/metrics`

Stats include:
1. Interpreter count
2. Interpreter bytes
3. Time spent blocked waiting for an available interpreter

Signed-off-by: Roman Gershman <roman@dragonflydb.io>

src/core/dash_test.cc

@@ -7,6 +7,7 @@
 #include "core/dash.h"
 
 #include <absl/container/flat_hash_map.h>
+#include <absl/strings/str_cat.h>
 #include <mimalloc.h>
 
 #include <functional>

src/core/interpreter.cc

@@ -474,18 +474,26 @@ int RedisLogCommand(lua_State* lua) {
 void* mimalloc_glue(void* ud, void* ptr, size_t osize, size_t nsize) {
   (void)ud;
   if (nsize == 0) {
+    InterpreterManager::tl_stats().used_bytes -= mi_usable_size(ptr);
     mi_free_size(ptr, osize);
     return nullptr;
   } else if (ptr == nullptr) {
-    return mi_malloc(nsize);
+    ptr = mi_malloc(nsize);
+    InterpreterManager::tl_stats().used_bytes += mi_usable_size(ptr);
+    return ptr;
   } else {
-    return mi_realloc(ptr, nsize);
+    InterpreterManager::tl_stats().used_bytes -= mi_usable_size(ptr);
+    ptr = mi_realloc(ptr, nsize);
+    InterpreterManager::tl_stats().used_bytes += mi_usable_size(ptr);
+    return ptr;
   }
 }
 
 }  // namespace
 
 Interpreter::Interpreter() {
+  InterpreterManager::tl_stats().interpreter_cnt++;
+
   lua_ = lua_newstate(mimalloc_glue, nullptr);
   InitLua(lua_);
   void** ptr = static_cast<void**>(lua_getextraspace(lua_));
@@ -562,6 +570,8 @@ Interpreter::Interpreter() {
 }
 
 Interpreter::~Interpreter() {
+  InterpreterManager::tl_stats().interpreter_cnt--;
+
   lua_close(lua_);
 }
 
@@ -1031,6 +1041,19 @@ int Interpreter::RedisAPCallCommand(lua_State* lua) {
   return reinterpret_cast<Interpreter*>(*ptr)->RedisGenericCommand(false, true);
 }
 
+InterpreterManager::Stats& InterpreterManager::Stats::operator+=(const Stats& other) {
+  this->used_bytes += other.used_bytes;
+  this->interpreter_cnt += other.interpreter_cnt;
+  this->blocked_cnt += other.blocked_cnt;
+
+  return *this;
+}
+
+InterpreterManager::Stats& InterpreterManager::tl_stats() {
+  static thread_local Stats stats;
+  return stats;
+}
+
 Interpreter* InterpreterManager::Get() {
   // Grow if none is available and we have unused capacity left.
   if (available_.empty() && storage_.size() < storage_.capacity()) {
@@ -1038,7 +1061,8 @@ Interpreter* InterpreterManager::Get() {
     return &storage_.back();
   }
 
-  waker_.await([this]() { return available_.size() > 0; });
+  bool blocked = waker_.await([this]() { return available_.size() > 0; });
+  tl_stats().blocked_cnt += (uint64_t)blocked;
 
   Interpreter* ir = available_.back();
   available_.pop_back();

src/core/interpreter.h

@@ -143,6 +143,15 @@ class Interpreter {
 // Manages an internal interpreter pool. This allows multiple connections residing on the same
 // thread to run multiple lua scripts in parallel.
 class InterpreterManager {
+ public:
+  struct Stats {
+    Stats& operator+=(const Stats& other);
+
+    uint64_t used_bytes = 0;
+    uint64_t interpreter_cnt = 0;
+    uint64_t blocked_cnt = 0;
+  };
+
  public:
   InterpreterManager(unsigned num) : waker_{}, available_{}, storage_{} {
     // We pre-allocate the backing storage during initialization and
@@ -157,6 +166,8 @@ class InterpreterManager {
   // Clear all interpreters, keeps capacity. Waits until all are returned.
   void Reset();
 
+  static Stats& tl_stats();
+
  private:
   util::fb2::EventCount waker_, reset_ec_;
   std::vector<Interpreter*> available_;

src/core/lru_test.cc

@@ -4,6 +4,8 @@
 
 #include "core/lru.h"
 
+#include <absl/strings/str_cat.h>
+
 #include "base/gtest.h"
 #include "base/logging.h"
 #include "core/compact_object.h"

src/server/server_family.cc

@@ -1135,6 +1135,12 @@ void PrintPrometheusMetrics(const Metrics& m, StringResponse* resp) {
                             &resp->body());
   AppendMetricWithoutLabels("keyspace_mutations_total", "", m.events.mutations, MetricType::COUNTER,
                             &resp->body());
+  AppendMetricWithoutLabels("lua_interpreter_cnt", "", m.lua_stats.interpreter_cnt,
+                            MetricType::GAUGE, &resp->body());
+  AppendMetricWithoutLabels("used_memory_lua", "", m.lua_stats.used_bytes, MetricType::GAUGE,
+                            &resp->body());
+  AppendMetricWithoutLabels("lua_blocked_total", "", m.lua_stats.blocked_cnt, MetricType::COUNTER,
+                            &resp->body());
 
   // Net metrics
   AppendMetricWithoutLabels("net_input_bytes_total", "", conn_stats.io_read_bytes,
@@ -1867,6 +1873,8 @@ Metrics ServerFamily::GetMetrics() const {
     result.tls_bytes += Listener::TLSUsedMemoryThreadLocal();
     result.refused_conn_max_clients_reached_count += Listener::RefusedConnectionMaxClientsCount();
 
+    result.lua_stats += InterpreterManager::tl_stats();
+
     service_.mutable_registry()->MergeCallStats(index, cmd_stat_cb);
   };
 
@@ -1969,6 +1977,8 @@ void ServerFamily::Info(CmdArgList args, ConnectionContext* cntx) {
     append("maxmemory", max_memory_limit);
     append("maxmemory_human", HumanReadableNumBytes(max_memory_limit));
 
+    append("used_memory_lua", m.lua_stats.used_bytes);
+
     // Blob - all these cases where the key/objects are represented by a single blob allocated on
     // heap. For example, strings or intsets. members of lists, sets, zsets etc
     // are not accounted for to avoid complex computations. In some cases, when number of members
@@ -2061,6 +2071,9 @@ void ServerFamily::Info(CmdArgList args, ConnectionContext* cntx) {
     append("blocked_on_interpreter", m.coordinator_stats.blocked_on_interpreter);
     append("ram_hits", m.events.ram_hits);
     append("ram_misses", m.events.ram_misses);
+
+    append("lua_interpreter_cnt", m.lua_stats.interpreter_cnt);
+    append("lua_blocked", m.lua_stats.blocked_cnt);
   }
 
   if (should_enter("TIERED", true)) {

src/server/server_family.h

@@ -107,6 +107,8 @@ struct Metrics {
   uint32_t worker_fiber_count = 0;
   size_t worker_fiber_stack_size = 0;
 
+  InterpreterManager::Stats lua_stats;
+
   // command call frequencies (count, aggregated latency in usec).
   std::map<std::string, std::pair<uint64_t, uint64_t>> cmd_stats_map;
   std::vector<ReplicaRoleInfo> replication_metrics;