Fix use-after-free in CallTraceStorage and optimize trace collection (#290)

* Fix use-after-free in CallTraceStorage and optimize trace collection

The original issue was a use-after-free bug in CallTraceStorage::processTraces()
where trace pointers were accessed after the memory was deallocated during
triple-buffer rotation.

Key changes:
- Fixed use-after-free by copying preserved traces immediately using putWithExistingId()
- Added optional hook parameter to CallTraceHashTable::collect() for inline processing
- Eliminated separate trace buffers (_standby_traces_buffer, _active_traces_buffer)
- Simplified processTraces() to collect directly to _traces_buffer with hooks
- Removed double iteration pattern (collect then iterate) in favor of single pass
- Added regression test to catch use-after-free with frame content access

The fix ensures memory safety while improving performance by eliminating
unnecessary buffer copying and reducing iteration overhead.

Co-Authored-By: Claude <[email protected]>
(cherry picked from commit 91fc9ec)

Author: jbachorik

.github/workflows/test_workflow.yml

@@ -124,7 +124,7 @@ jobs:
           name: failures-glibc-${{ matrix.java_version }}-${{ matrix.config }}-amd64
           path: failures_glibc-${{ matrix.java_version }}-${{ matrix.config }}-amd64.txt
       - name: Prepare reports
-        if: always()
+        if: steps.set_enabled.outputs.enabled == 'true'
         run: |
           .github/scripts/prepare_reports.sh
       - name: Upload unwinding reports
@@ -247,7 +247,7 @@ jobs:
           name: failures-musl-${{ matrix.java_version }}-${{ matrix.config }}-amd64
           path: failures_musl-${{ matrix.java_version }}-${{ matrix.config }}-amd64.txt
       - name: Prepare reports
-        if: always()
+        if: steps.set_enabled.outputs.enabled == 'true'
         run: |
           .github/scripts/prepare_reports.sh
       - name: Upload unwinding reports
@@ -378,7 +378,7 @@ jobs:
           name: failures-glibc-${{ matrix.java_version }}-${{ matrix.config }}-aarch64
           path: failures_glibc-${{ matrix.java_version }}-${{ matrix.config }}-aarch64.txt
       - name: Prepare reports
-        if: always()
+        if: steps.set_enabled.outputs.enabled == 'true'
         run: |
           .github/scripts/prepare_reports.sh
       - name: Upload unwinding reports
@@ -480,7 +480,7 @@ jobs:
            name: failures-musl-${{ matrix.java_version }}-${{ matrix.config }}-aarch64
            path: failures_musl-${{ matrix.java_version }}-${{ matrix.config }}-aarch64.txt
        - name: Prepare reports
-         if: always()
+         if: steps.set_enabled.outputs.enabled == 'true'
          run: |
            .github/scripts/prepare_reports.sh
        - name: Upload unwinding reports

ddprof-lib/src/main/cpp/callTraceHashTable.cpp

@@ -9,6 +9,7 @@
 #include "os.h"
 #include "arch_dd.h"
 #include "common.h"
+#include "primeProbing.h"
 #include <string.h>
 #include <signal.h>
 #include <pthread.h>
@@ -70,8 +71,6 @@ class LongHashTable {
 
   CallTraceSample *values() { return (CallTraceSample *)(keys() + _capacity); }
 
-  u32 nextSlot(u32 slot) const { return (slot + 1) & (_capacity - 1); }
-
   void clear() {
     memset(keys(), 0, (sizeof(u64) + sizeof(CallTraceSample)) * _capacity);
     _size = 0;
@@ -178,21 +177,23 @@ CallTrace *CallTraceHashTable::storeCallTrace(int num_frames,
 
 CallTrace *CallTraceHashTable::findCallTrace(LongHashTable *table, u64 hash) {
   u64 *keys = table->keys();
-  u32 capacity = table->capacity();
-  u32 slot = hash & (capacity - 1);
-  u32 step = 0;
+  HashProbe probe(hash, table->capacity());
 
-  while (keys[slot] != hash) {
+  u32 slot = probe.slot();
+  while (true) {
+    if (keys[slot] == hash) {
+      return table->values()[slot].trace;
+    }
     if (keys[slot] == 0) {
       return nullptr;
     }
-    if (++step >= capacity) {
-      return nullptr;
+    if (!probe.hasNext()) {
+      break;
     }
-    slot = (slot + step) & (capacity - 1);
-  }
+    slot = probe.next();
+  };
 
-  return table->values()[slot].trace;
+  return nullptr;
 }
 
 u64 CallTraceHashTable::put(int num_frames, ASGCT_CallFrame *frames,
@@ -207,10 +208,9 @@ u64 CallTraceHashTable::put(int num_frames, ASGCT_CallFrame *frames,
   }
 
   u64 *keys = table->keys();
-  u32 capacity = table->capacity();
-  u32 slot = hash & (capacity - 1);
-  u32 step = 0;
-  
+  HashProbe probe(hash, table->capacity());
+
+  u32 slot = probe.slot();
   while (true) {
     u64 key_value = __atomic_load_n(&keys[slot], __ATOMIC_RELAXED);
     if (key_value == hash) { 
@@ -278,6 +278,8 @@ u64 CallTraceHashTable::put(int num_frames, ASGCT_CallFrame *frames,
       u32 new_size = table->incSize();
       u32 capacity = table->capacity();
 
+      probe.updateCapacity(new_size);
+
       // EXPANSION LOGIC: Check if 75% capacity reached after incrementing size
       if (new_size == capacity * 3 / 4) {
         // Allocate new table with double capacity using LinearAllocator
@@ -314,18 +316,18 @@ u64 CallTraceHashTable::put(int num_frames, ASGCT_CallFrame *frames,
       return trace->trace_id;
     }
 
-    if (++step >= capacity) {
+    if (!probe.hasNext()) {
       // Table overflow - very unlikely with expansion logic
       atomicIncRelaxed(_overflow);
       return OVERFLOW_TRACE_ID;
     }
-    // Linear probing with step increment
-    slot = (slot + step) & (capacity - 1);
+    // Prime probing for better distribution
+    slot = probe.next();
   }
 }
 
-void CallTraceHashTable::collect(std::unordered_set<CallTrace *> &traces) {
-  // Lock-free collection for read-only tables (after atomic swap)
+void CallTraceHashTable::collect(std::unordered_set<CallTrace *> &traces, std::function<void(CallTrace*)> trace_hook) {
+  // Lock-free collection for read-only tables
   // No new put() operations can occur, so no synchronization needed
 
   // Collect from all tables in the chain (current and previous tables)
@@ -337,6 +339,9 @@ void CallTraceHashTable::collect(std::unordered_set<CallTrace *> &traces) {
       if (keys[slot] != 0) {
         CallTrace *trace = values[slot].acquireTrace();
         if (trace != nullptr && trace != CallTraceSample::PREPARING) {
+          if (trace_hook) {
+            trace_hook(trace);  // Call hook first if provided
+          }
           traces.insert(trace);
         }
       }
@@ -345,10 +350,14 @@ void CallTraceHashTable::collect(std::unordered_set<CallTrace *> &traces) {
 
   // Handle overflow trace
   if (_overflow > 0) {
+    if (trace_hook) {
+      trace_hook(&_overflow_trace);  // Call hook for overflow trace too
+    }
     traces.insert(&_overflow_trace);
   }
 }
 
+
 void CallTraceHashTable::putWithExistingId(CallTrace* source_trace, u64 weight) {
   // Trace preservation for standby tables (no contention with new puts)
   // This is safe because new put() operations go to the new active table
@@ -371,22 +380,22 @@ void CallTraceHashTable::putWithExistingId(CallTrace* source_trace, u64 weight)
 
   u64 *keys = table->keys();
   u32 capacity = table->capacity();
-  u32 slot = hash & (capacity - 1);
-  
+
+  HashProbe probe(hash, capacity);
+  u32 slot = probe.slot();
+
   // Look for existing entry or empty slot - no locking needed
   while (true) {
     u64 key_value = __atomic_load_n(&keys[slot], __ATOMIC_RELAXED);
-    if (key_value == hash) {
-      // Found existing entry - just use it (trace already preserved)
-      break;
-    }
     if (key_value == 0) {
       // Found empty slot - claim it atomically
       u64 expected = 0;
       if (!__atomic_compare_exchange_n(&keys[slot], &expected, hash, false, __ATOMIC_ACQ_REL, __ATOMIC_RELAXED)) {
         // Another thread claimed it, try next slot
-        slot = table->nextSlot(slot);
-        continue;
+        if (probe.hasNext()) {
+          slot = probe.next();
+          continue;
+        }
       }
 
       // Create a copy of the source trace preserving its exact ID
@@ -403,14 +412,15 @@ void CallTraceHashTable::putWithExistingId(CallTrace* source_trace, u64 weight)
         Counters::increment(CALLTRACE_STORAGE_TRACES);
 
         // Increment table size
-        table->incSize();
+        u32 new_size = table->incSize();
+        probe.updateCapacity(new_size);
       } else {
         // Allocation failure - clear the key we claimed
         __atomic_store_n(&keys[slot], 0, __ATOMIC_RELEASE);
       }
       break;
     }
 
-    slot = table->nextSlot(slot);
+    slot = probe.next();
   }
 }

ddprof-lib/src/main/cpp/callTraceHashTable.h

@@ -11,6 +11,7 @@
 #include "vmEntry.h"
 #include <unordered_set>
 #include <atomic>
+#include <functional>
 
 class LongHashTable;
 
@@ -78,7 +79,7 @@ class CallTraceHashTable {
   ~CallTraceHashTable();
 
   void clear();
-  void collect(std::unordered_set<CallTrace *> &traces);
+  void collect(std::unordered_set<CallTrace *> &traces, std::function<void(CallTrace*)> trace_hook = nullptr);
 
   u64 put(int num_frames, ASGCT_CallFrame *frames, bool truncated, u64 weight);
   void putWithExistingId(CallTrace* trace, u64 weight);  // For standby tables with no contention