Fix negative table size calculation when socket tracer disabled via s…

…tirling flag (#1894) Summary: When some or all of the tables `http_events, stirling_error, probe_status, proc_exit_events` are disabled via stirling flags (`PL_STIRLING_SOURCES=...`) the `other_table_size` calculation returns a negative output if `num_tables < 4`. It also attempts division by zero when `num_tables = 4` ```cpp int64_t other_table_size = (memory_limit - http_table_size - stirling_error_table_size - probe_status_table_size - proc_exit_events_table_size) / (num_tables - 4); ``` This PR streamlines the table size calculation and fixes the edge condition. Type of change: /kind bug Test Plan: deployed to node and manually checked calculated table sizes. --------- Signed-off-by: Benjamin Kilimnik <[email protected]>
pixie-io · May 20, 2024 · d8aab71 · d8aab71
1 parent 3078d85
commit d8aab71
Show file tree

Hide file tree

Showing 2 changed files with 92 additions and 18 deletions.
diff --git a/src/experimental/standalone_pem/standalone_pem_manager.cc b/src/experimental/standalone_pem/standalone_pem_manager.cc
@@ -158,16 +158,53 @@ Status StandalonePEMManager::InitSchemas() {
   stirling_->GetPublishProto(&publish_pb);
   auto relation_info_vec = ConvertPublishPBToRelationInfo(publish_pb);
 
-  int64_t memory_limit = FLAGS_table_store_data_limit * 1024 * 1024;
-  int64_t num_tables = relation_info_vec.size();
-  int64_t http_table_size = (FLAGS_table_store_http_events_percent * memory_limit) / 100;
-  int64_t stirling_error_table_size = FLAGS_table_store_stirling_error_limit_bytes / 2;
-  int64_t probe_status_table_size = FLAGS_table_store_stirling_error_limit_bytes / 2;
-  int64_t proc_exit_events_table_size = FLAGS_table_store_proc_exit_events_limit_bytes;
-  int64_t other_table_size = (memory_limit - http_table_size - stirling_error_table_size -
-                              probe_status_table_size - proc_exit_events_table_size) /
-                             (num_tables - 4);
+  const int64_t memory_limit = FLAGS_table_store_data_limit * 1024 * 1024;
+  const int64_t num_tables = relation_info_vec.size();
+  const int64_t http_table_size = (FLAGS_table_store_http_events_percent * memory_limit) / 100;
+  const int64_t stirling_error_table_size = FLAGS_table_store_stirling_error_limit_bytes / 2;
+  const int64_t probe_status_table_size = FLAGS_table_store_stirling_error_limit_bytes / 2;
+  const int64_t proc_exit_events_table_size = FLAGS_table_store_proc_exit_events_limit_bytes;
+
+  // Determine which of the four default tables are present
+  bool has_http_events = false, has_stirling_error = false, has_probe_status = false,
+       has_proc_exit_events = false;
+  for (const auto& relation_info : relation_info_vec) {
+    if (relation_info.name == "http_events") {
+      has_http_events = true;
+    } else if (relation_info.name == "stirling_error") {
+      has_stirling_error = true;
+    } else if (relation_info.name == "probe_status") {
+      has_probe_status = true;
+    } else if (relation_info.name == "proc_exit_events") {
+      has_proc_exit_events = true;
+    }
+  }
+
+  // Calculate memory used by specific tables
+  int64_t used_memory = 0;
+  if (has_http_events) {
+    used_memory += http_table_size;
+  }
+  if (has_stirling_error) {
+    used_memory += stirling_error_table_size;
+  }
+  if (has_probe_status) {
+    used_memory += probe_status_table_size;
+  }
+  if (has_proc_exit_events) {
+    used_memory += proc_exit_events_table_size;
+  }
+
+  const int64_t remaining_memory = memory_limit - used_memory;
+  if (remaining_memory < 0) {
+    return error::Internal("Table store data limit is too low to store the tables.");
+  }
+  const int64_t other_table_count =
+      num_tables - (has_http_events + has_stirling_error + has_probe_status + has_proc_exit_events);
+  const int64_t other_table_size =
+      (other_table_count > 0) ? remaining_memory / other_table_count : 0;
 
+  // Create tables with allocated sizes
   for (const auto& relation_info : relation_info_vec) {
     std::shared_ptr<table_store::Table> table_ptr;
     if (relation_info.name == "http_events") {

diff --git a/src/vizier/services/agent/pem/pem_manager.cc b/src/vizier/services/agent/pem/pem_manager.cc
@@ -92,16 +92,53 @@ Status PEMManager::InitSchemas() {
   stirling_->GetPublishProto(&publish_pb);
   auto relation_info_vec = ConvertPublishPBToRelationInfo(publish_pb);
 
-  int64_t memory_limit = FLAGS_table_store_data_limit * 1024 * 1024;
-  int64_t num_tables = relation_info_vec.size();
-  int64_t http_table_size = (FLAGS_table_store_http_events_percent * memory_limit) / 100;
-  int64_t stirling_error_table_size = FLAGS_table_store_stirling_error_limit_bytes / 2;
-  int64_t probe_status_table_size = FLAGS_table_store_stirling_error_limit_bytes / 2;
-  int64_t proc_exit_events_table_size = FLAGS_table_store_proc_exit_events_limit_bytes;
-  int64_t other_table_size = (memory_limit - http_table_size - stirling_error_table_size -
-                              probe_status_table_size - proc_exit_events_table_size) /
-                             (num_tables - 4);
+  const int64_t memory_limit = FLAGS_table_store_data_limit * 1024 * 1024;
+  const int64_t num_tables = relation_info_vec.size();
+  const int64_t http_table_size = (FLAGS_table_store_http_events_percent * memory_limit) / 100;
+  const int64_t stirling_error_table_size = FLAGS_table_store_stirling_error_limit_bytes / 2;
+  const int64_t probe_status_table_size = FLAGS_table_store_stirling_error_limit_bytes / 2;
+  const int64_t proc_exit_events_table_size = FLAGS_table_store_proc_exit_events_limit_bytes;
+
+  // Determine which of the four default tables are present
+  bool has_http_events = false, has_stirling_error = false, has_probe_status = false,
+       has_proc_exit_events = false;
+  for (const auto& relation_info : relation_info_vec) {
+    if (relation_info.name == "http_events") {
+      has_http_events = true;
+    } else if (relation_info.name == "stirling_error") {
+      has_stirling_error = true;
+    } else if (relation_info.name == "probe_status") {
+      has_probe_status = true;
+    } else if (relation_info.name == "proc_exit_events") {
+      has_proc_exit_events = true;
+    }
+  }
+
+  // Calculate memory used by specific tables
+  int64_t used_memory = 0;
+  if (has_http_events) {
+    used_memory += http_table_size;
+  }
+  if (has_stirling_error) {
+    used_memory += stirling_error_table_size;
+  }
+  if (has_probe_status) {
+    used_memory += probe_status_table_size;
+  }
+  if (has_proc_exit_events) {
+    used_memory += proc_exit_events_table_size;
+  }
+
+  const int64_t remaining_memory = memory_limit - used_memory;
+  if (remaining_memory < 0) {
+    return error::Internal("Table store data limit is too low to store the tables.");
+  }
+  const int64_t other_table_count =
+      num_tables - (has_http_events + has_stirling_error + has_probe_status + has_proc_exit_events);
+  const int64_t other_table_size =
+      (other_table_count > 0) ? remaining_memory / other_table_count : 0;
 
+  // Create tables with allocated sizes
   for (const auto& relation_info : relation_info_vec) {
     std::shared_ptr<table_store::Table> table_ptr;
     if (relation_info.name == "http_events") {