Reset CallTraceStorage counters before reporting live objects.

src/callTraceStorage.cpp

@@ -289,3 +289,19 @@ void CallTraceStorage::add(u32 call_trace_id, u64 samples, u64 counter) {
         }
     }
 }
+
+void CallTraceStorage::resetCounters() {
+     for (LongHashTable* table = _current_table; table != NULL; table = table->prev()) {
+        u64* keys = table->keys();
+        CallTraceSample* values = table->values();
+        u32 capacity = table->capacity();
+
+        for (u32 slot = 0; slot < capacity; slot++) {
+            if (keys[slot] != 0) {
+                CallTraceSample& s = values[slot];
+                storeRelease(s.samples, 0);
+                storeRelease(s.counter, 0);
+            }
+        }
+    }
+}

src/callTraceStorage.h

@@ -70,6 +70,7 @@ class CallTraceStorage {
 
     u32 put(int num_frames, ASGCT_CallFrame* frames, u64 counter);
     void add(u32 call_trace_id, u64 samples, u64 counter);
+    void resetCounters();
 };
 
 #endif // _CALLTRACESTORAGE

src/objectSampler.cpp

@@ -100,6 +100,9 @@ class LiveRefs {
         jvmtiEnv* jvmti = VM::jvmti();
         Profiler* profiler = Profiler::instance();
 
+        // Reset counters before dumping to collect live objects only.
+        profiler->tryResetCounters();
+
         for (u32 i = 0; i < MAX_REFS; i++) {
             if ((i % 32) == 0) jni->PushLocalFrame(64);
 
@@ -147,13 +150,9 @@ void ObjectSampler::recordAllocation(jvmtiEnv* jvmti, JNIEnv* jni, EventType eve
     event._instance_size = size;
     event._class_id = lookupClassId(jvmti, object_klass);
 
-    if (_live) {
-        u64 trace = Profiler::instance()->recordSample(NULL, 0, event_type, &event);
-        if (trace != 0) {
-            live_refs.add(jni, object, size, trace);
-        }
-    } else {
-        Profiler::instance()->recordSample(NULL, event._total_size, event_type, &event);
+    u64 trace = Profiler::instance()->recordSample(NULL, event._total_size, event_type, &event);
+    if (_live && trace != 0) {
+        live_refs.add(jni, object, size, trace);
     }
 }
 

src/profiler.cpp

@@ -773,6 +773,15 @@ void Profiler::recordEventOnly(EventType event_type, Event* event) {
     _locks[lock_index].unlock();
 }
 
+void Profiler::tryResetCounters() {
+    // Reset counters only for non-JFR recording, otherwise resetting may cause missing stack traces for some
+    // allocation events and skewed incorrect number of samples.
+    // In JFR recording, each sample is recorded individually, so accumulated counters are not actually used.
+    if (!_jfr.active()) {
+        _call_trace_storage.resetCounters();
+    }
+}
+
 void Profiler::writeLog(LogLevel level, const char* message) {
     _jfr.recordLog(level, message, strlen(message));
 }

src/profiler.h

@@ -205,6 +205,7 @@ class Profiler {
     void recordExternalSample(u64 counter, int tid, EventType event_type, Event* event, int num_frames, ASGCT_CallFrame* frames);
     void recordExternalSamples(u64 samples, u64 counter, int tid, u32 call_trace_id, EventType event_type, Event* event);
     void recordEventOnly(EventType event_type, Event* event);
+    void tryResetCounters();
     void writeLog(LogLevel level, const char* message);
     void writeLog(LogLevel level, const char* message, size_t len);
 

test/one/profiler/test/Assert.java

@@ -25,6 +25,12 @@ public static void isGreater(double value, double threshold, String message) {
         }
     }
 
+    public static void isGreaterOrEqual(double value, double threshold) {
+        if (value < threshold) {
+            throw new AssertionError("Expected " + value + " >= " + threshold);
+        }
+    }
+
     public static void isLess(double value, double threshold) {
         isLess(value, threshold, null);
     }
@@ -38,4 +44,10 @@ public static void isLess(double value, double threshold, String message) {
                     "Expected " + value + " < " + threshold + (message != null ? (": " + message) : ""));
         }
     }
+
+    public static void isLessOrEqual(double value, double threshold) {
+        if (value > threshold) {
+            throw new AssertionError("Expected " + value + " <= " + threshold);
+        }
+    }
 }

test/one/profiler/test/TestProcess.java

@@ -53,6 +53,7 @@ private static MethodHandle getPidHandle() {
         }
     }
 
+    private final Test test;
     private final Os currentOs;
     private final String logDir;
     private final String[] inputs;
@@ -61,6 +62,7 @@ private static MethodHandle getPidHandle() {
     private final int timeout = 30;
 
     public TestProcess(Test test, Os currentOs, String logDir) throws Exception {
+        this.test = test;
         this.currentOs = currentOs;
         this.logDir = logDir;
         this.inputs = test.inputs();
@@ -83,6 +85,10 @@ public TestProcess(Test test, Os currentOs, String logDir) throws Exception {
         }
     }
 
+    public Test test() {
+        return this.test;
+    }
+
     public String[] inputs() {
         return this.inputs;
     }

test/test/alloc/AllocTests.java

@@ -5,16 +5,16 @@
 
 package test.alloc;
 
-import one.profiler.test.Jvm;
-import one.profiler.test.Os;
-import one.profiler.test.Output;
-import one.profiler.test.Test;
-import one.profiler.test.TestProcess;
+import one.jfr.JfrReader;
+import one.jfr.StackTrace;
+import one.jfr.event.AllocationSample;
+import one.profiler.test.*;
 
 import java.io.File;
 import java.nio.file.Files;
 import java.nio.file.Paths;
 import java.nio.file.StandardCopyOption;
+import java.util.List;
 
 public class AllocTests {
 
@@ -72,4 +72,44 @@ public void objectSamplerWtihDifferentAsprofs(TestProcess p) throws Exception {
         Output outWithCopy = p.profile(String.format("--libpath %s -e alloc -d 3 -o collapsed", asprofCopy.getAbsolutePath()));
         assert !outWithCopy.contains("_\\[k\\]"); // first instance of profiler has not properly relinquished the can_generate_sampled_object_alloc_events capability.
     }
+
+    // Without liveness tracking, results won't change except for the sampling
+    // error.
+    @Test(mainClass = RandomBlockRetainer.class, jvmVer = {11, Integer.MAX_VALUE}, args = "1.0", agentArgs = "start,alloc=1k,total,file=%f,collapsed")
+    @Test(mainClass = RandomBlockRetainer.class, jvmVer = {11, Integer.MAX_VALUE}, args = "0.0", agentArgs = "start,alloc=1k,total,file=%f,collapsed,live")
+    @Test(mainClass = RandomBlockRetainer.class, jvmVer = {11, Integer.MAX_VALUE}, args = "0.1", agentArgs = "start,alloc=1k,total,file=%f,collapsed,live")
+    @Test(mainClass = RandomBlockRetainer.class, jvmVer = {11, Integer.MAX_VALUE}, args = "1.0", agentArgs = "start,alloc=1k,total,file=%f,collapsed,live")
+    public void liveness(TestProcess p) throws Exception {
+        final long TOTAL_BYTES = 50000000;
+        final double tolerance = 0.10;
+
+        // keepChance = live ? args() : 1.0, which is equal to args().
+        final double keepChance = Double.parseDouble(p.test().args());
+
+        Output out = p.waitForExit("%f");
+        long totalBytes = out.filter("RandomBlockRetainer\\.alloc").samples("byte\\[\\]");
+
+        final double lowerLimit = (keepChance - tolerance) * TOTAL_BYTES;
+        final double upperLimit = (keepChance + tolerance) * TOTAL_BYTES;
+
+        Assert.isLessOrEqual(lowerLimit, totalBytes);
+        Assert.isGreaterOrEqual(upperLimit, totalBytes);
+    }
+
+    @Test(mainClass = RandomBlockRetainer.class, jvmVer = {11, Integer.MAX_VALUE}, args = "1.0", agentArgs = "start,alloc=1k,total,file=%f.jfr")
+    @Test(mainClass = RandomBlockRetainer.class, jvmVer = {11, Integer.MAX_VALUE}, args = "0.0", agentArgs = "start,alloc=1k,total,file=%f.jfr,live")
+    @Test(mainClass = RandomBlockRetainer.class, jvmVer = {11, Integer.MAX_VALUE}, args = "0.1", agentArgs = "start,alloc=1k,total,file=%f.jfr,live")
+    @Test(mainClass = RandomBlockRetainer.class, jvmVer = {11, Integer.MAX_VALUE}, args = "1.0", agentArgs = "start,alloc=1k,total,file=%f.jfr,live")
+    public void livenessJfrHasStacks(TestProcess p) throws Exception {
+        p.waitForExit();
+        String filename = p.getFile("%f").toPath().toString();
+        try (JfrReader r = new JfrReader(filename)) {
+            List<AllocationSample> events = r.readAllEvents(AllocationSample.class);
+            assert !events.isEmpty() : "No AllocationSample events found in JFR output";
+            for (AllocationSample event : events) {
+                StackTrace trace = r.stackTraces.get(event.stackTraceId);
+                assert trace != null : "Stack trace missing for id " + event.stackTraceId;
+            }
+        }
+    }
 }

test/test/alloc/RandomBlockRetainer.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright The async-profiler authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package test.alloc;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class RandomBlockRetainer {
+    public static void main(String[] args) throws Exception {
+        double keepChance = 0.5;
+        if (args.length > 0) {
+            try {
+                keepChance = Double.parseDouble(args[0]);
+            } catch (NumberFormatException e) {
+                System.err.println("Invalid keepChance value. Using default value of 0.5.");
+            }
+        }
+
+        // Set up a list to hold large objects and keep them in memory.
+        List<byte[]> rooter = new ArrayList<>();
+
+        final int TOTAL_BLOCKS = 500; // Has to be less than LiveRefs::MAX_REFS for testing purposes.
+        final int BLOCK_SIZE = 100 * 1000;
+
+        for (int i = 0; i < TOTAL_BLOCKS; i++) {
+            byte[] block = i % 2 == 0 ? alloc1(BLOCK_SIZE) : alloc2(BLOCK_SIZE);
+
+            if (Math.random() <= keepChance) {
+                // Keep a reference to prevent the block from being garbage collected
+                rooter.add(block);
+            }
+        }
+
+        System.gc();
+    }
+
+    private static byte[] alloc1(int size) {
+        return new byte[size];
+    }
+
+    private static byte[] alloc2(int size) {
+        return new byte[size];
+    }
+}