Provide configurable memory limit for Conn tracker buffer usage (#2196)

Summary: Provide configurable memory limit for Conn tracker buffer usage

Some memory conscious users want to prevent PEM's from having unbounded
memory use. The socket tracer's conn tracker component allows unbounded
memory growth due to storing an uncapped amount of socket data. From
profiling the [PEM's heap
use](https://github.com/pixie-io/pixie/blob/301198f2af6739f0de8b01e1cbcb3451b3075e28/scripts/collect_heap_pprofs.sh),
this memory growth is a significant contributor to the PEM's memory
footprint.

By leveraging the existing [table store
limit](https://github.com/pixie-io/pixie/blob/301198f2af6739f0de8b01e1cbcb3451b3075e28/src/table_store/table/table.cc#L44)
and this new configuration option, these memory conscious users run
Pixie with a smaller and more predictable steady state memory usage.

Relevant Issues: N/A

Type of change: /kind feature

Test Plan: Ran perf_tool and k9s based load tests to verify the
following:
- [x] Verified results by running PEM's with a 0.5-1 GiB memory limit
- [x] Verified Conn tracker buffers no longer showed up as significant
memory contributors in heap profiles.

Changelog Message: PEM's can now limit the total memory used by Pixie's
socket tracer buffers through the `--total_conn_tracker_mem_usage`
command line flag or the `PX_TOTAL_CONN_TRACKER_MEM_USAGE` env var. This
allows memory conscious users to have more predictable memory use when
paired with `PL_TABLE_STORE_TABLE_SIZE_LIMIT`.

Signed-off-by: Dom Del Nano <[email protected]>

Author: ddelnano

src/stirling/source_connectors/socket_tracer/socket_trace_connector.cc

@@ -170,6 +170,11 @@ DEFINE_uint32(datastream_buffer_retention_size,
               gflags::Uint32FromEnv("PL_DATASTREAM_BUFFER_SIZE", 1024 * 1024),
               "The maximum size of a data stream buffer retained between cycles.");
 
+DEFINE_uint64(total_conn_tracker_mem_usage,
+              gflags::Uint64FromEnv("PX_TOTAL_CONN_TRACKER_MEM_USAGE", 0),
+              "The maximum size in bytes of the collective connection tracker buffers. Defaults to "
+              "0, which corresponds to no limit. When data is beyond this limit, new data is "
+              "dropped until the limit is no longer exceeded.");
 DEFINE_uint64(max_body_bytes, gflags::Uint64FromEnv("PL_STIRLING_MAX_BODY_BYTES", 512),
               "The maximum number of bytes in the body of protocols like HTTP");
 
@@ -834,6 +839,7 @@ void SocketTraceConnector::TransferDataImpl(ConnectorContext* ctx) {
     }
   }
 
+  size_t current_conn_tracker_buffer_size = 0;
   for (const auto& conn_tracker : conn_trackers_mgr_.active_trackers()) {
     const auto& transfer_spec = protocol_transfer_specs_[conn_tracker->protocol()];
 
@@ -857,7 +863,8 @@ void SocketTraceConnector::TransferDataImpl(ConnectorContext* ctx) {
                                    socket_info_mgr_.get());
 
     if (transfer_spec.transfer_fn != nullptr) {
-      transfer_spec.transfer_fn(*this, ctx, conn_tracker, data_table);
+      current_conn_tracker_buffer_size +=
+          transfer_spec.transfer_fn(*this, ctx, conn_tracker, data_table);
     } else {
       // If there's no transfer function, then the tracker should not be holding any data.
       // http::ProtocolTraits is used as a placeholder; the frames deque is expected to be
@@ -868,6 +875,7 @@ void SocketTraceConnector::TransferDataImpl(ConnectorContext* ctx) {
 
     conn_tracker->IterationPostTick();
   }
+  total_conn_tracker_mem_usage_ = current_conn_tracker_buffer_size;
 
   CheckTracerState();
 
@@ -1094,6 +1102,19 @@ void SocketTraceConnector::AcceptDataEvent(std::unique_ptr<SocketDataEvent> even
     WriteDataEvent(*event);
   }
 
+  auto msg_size = event->msg.size();
+  auto protocol = event->attr.protocol;
+  total_conn_tracker_mem_usage_ += msg_size;
+  if (FLAGS_total_conn_tracker_mem_usage != 0 &&
+      total_conn_tracker_mem_usage_ > FLAGS_total_conn_tracker_mem_usage) {
+    VLOG_EVERY_N(1, 1000) << absl::Substitute(
+        "Total buffer size of all active ConnTrackers $0 exceeds the limit $1. "
+        "Dropping data event of size $2 for protocol $3",
+        total_conn_tracker_mem_usage_, FLAGS_total_conn_tracker_mem_usage, msg_size, protocol);
+    stats_.Increment(StatKey::kDroppedSocketDataEvent);
+    return;
+  }
+
   stats_.Increment(StatKey::kPollSocketDataEventCount);
   stats_.Increment(StatKey::kPollSocketDataEventAttrSize, sizeof(event->attr));
   stats_.Increment(StatKey::kPollSocketDataEventDataSize, event->msg.size());
@@ -1113,11 +1134,32 @@ void SocketTraceConnector::AcceptConnStatsEvent(conn_stats_event_t event) {
 }
 
 void SocketTraceConnector::AcceptHTTP2Header(std::unique_ptr<HTTP2HeaderEvent> event) {
+  if (FLAGS_total_conn_tracker_mem_usage != 0 &&
+      total_conn_tracker_mem_usage_ > FLAGS_total_conn_tracker_mem_usage) {
+    VLOG_EVERY_N(1, 1000) << absl::Substitute(
+        "Total buffer size of all active ConnTrackers $0 exceeds the limit $1. "
+        "Dropping header event",
+        total_conn_tracker_mem_usage_, FLAGS_total_conn_tracker_mem_usage);
+    stats_.Increment(StatKey::kDroppedSocketDataEvent);
+    return;
+  }
+
   ConnTracker& tracker = GetOrCreateConnTracker(event->attr.conn_id);
   tracker.AddHTTP2Header(std::move(event));
 }
 
 void SocketTraceConnector::AcceptHTTP2Data(std::unique_ptr<HTTP2DataEvent> event) {
+  auto payload_size = event->payload.size();
+  total_conn_tracker_mem_usage_ += payload_size;
+  if (FLAGS_total_conn_tracker_mem_usage != 0 &&
+      total_conn_tracker_mem_usage_ > FLAGS_total_conn_tracker_mem_usage) {
+    VLOG_EVERY_N(1, 1000) << absl::Substitute(
+        "Total buffer size of all active ConnTrackers $0 exceeds the limit $1. "
+        "Dropping data event of size $2",
+        total_conn_tracker_mem_usage_, FLAGS_total_conn_tracker_mem_usage, payload_size);
+    stats_.Increment(StatKey::kDroppedSocketDataEvent);
+    return;
+  }
   ConnTracker& tracker = GetOrCreateConnTracker(event->attr.conn_id);
   tracker.AddHTTP2Data(std::move(event));
 }
@@ -1791,8 +1833,8 @@ void SocketTraceConnector::WriteDataEvent(const SocketDataEvent& event) {
 //-----------------------------------------------------------------------------
 
 template <typename TProtocolTraits>
-void SocketTraceConnector::TransferStream(ConnectorContext* ctx, ConnTracker* tracker,
-                                          DataTable* data_table) {
+size_t SocketTraceConnector::TransferStream(ConnectorContext* ctx, ConnTracker* tracker,
+                                            DataTable* data_table) {
   using TFrameType = typename TProtocolTraits::frame_type;
   using TKey = typename TProtocolTraits::key_type;
 
@@ -1821,6 +1863,7 @@ void SocketTraceConnector::TransferStream(ConnectorContext* ctx, ConnTracker* tr
   tracker->Cleanup<TProtocolTraits>(FLAGS_messages_size_limit_bytes,
                                     FLAGS_datastream_buffer_retention_size,
                                     message_expiry_timestamp, buffer_expiry_timestamp);
+  return tracker->MemUsage<TProtocolTraits>();
 }
 
 void SocketTraceConnector::TransferConnStats(ConnectorContext* ctx, DataTable* data_table) {

src/stirling/source_connectors/socket_tracer/socket_trace_connector.h

@@ -213,7 +213,7 @@ class SocketTraceConnector : public BCCSourceConnector {
       /* OUT */ struct go_grpc_http2_header_event_t* header_event_data_go_style);
 
   template <typename TProtocolTraits>
-  void TransferStream(ConnectorContext* ctx, ConnTracker* tracker, DataTable* data_table);
+  size_t TransferStream(ConnectorContext* ctx, ConnTracker* tracker, DataTable* data_table);
   void TransferConnStats(ConnectorContext* ctx, DataTable* data_table);
 
   void set_iteration_time(std::chrono::time_point<std::chrono::steady_clock> time) {
@@ -256,7 +256,7 @@ class SocketTraceConnector : public BCCSourceConnector {
     int32_t trace_mode = TraceMode::Off;
     uint32_t table_num = 0;
     std::vector<endpoint_role_t> trace_roles;
-    std::function<void(SocketTraceConnector&, ConnectorContext*, ConnTracker*, DataTable*)>
+    std::function<size_t(SocketTraceConnector&, ConnectorContext*, ConnTracker*, DataTable*)>
         transfer_fn = nullptr;
     bool enabled = false;
   };
@@ -300,6 +300,8 @@ class SocketTraceConnector : public BCCSourceConnector {
 
   UProbeManager uprobe_mgr_;
 
+  size_t total_conn_tracker_mem_usage_ = 0;
+
   enum class StatKey {
     kLossSocketDataEvent,
     kLossSocketControlEvent,
@@ -314,6 +316,8 @@ class SocketTraceConnector : public BCCSourceConnector {
     kPollSocketDataEventAttrSize,
     kPollSocketDataEventDataSize,
     kPollSocketDataEventSize,
+
+    kDroppedSocketDataEvent,
   };
 
   utils::StatCounter<StatKey> stats_;