Update JfrReader

Author: lucko

spark-common/src/main/java/me/lucko/spark/common/sampler/async/jfr/Dictionary.java

@@ -5,6 +5,8 @@
 
 package me.lucko.spark.common.sampler.async.jfr;
 
+import java.util.Arrays;
+
 /**
  * Fast and compact long->Object map.
  */
@@ -16,26 +18,35 @@ public class Dictionary<T> {
     private int size;
 
     public Dictionary() {
-        this.keys = new long[INITIAL_CAPACITY];
-        this.values = new Object[INITIAL_CAPACITY];
+        this(INITIAL_CAPACITY);
+    }
+
+    public Dictionary(int initialCapacity) {
+        this.keys = new long[initialCapacity];
+        this.values = new Object[initialCapacity];
     }
 
     public void clear() {
-        keys = new long[INITIAL_CAPACITY];
-        values = new Object[INITIAL_CAPACITY];
+        Arrays.fill(keys, 0);
+        Arrays.fill(values, null);
         size = 0;
     }
 
-    // spark start
     public int size() {
-        return this.size;
+        return size;
     }
-    // spark end
 
-    public void put(long key, T value) {
-        if (key == 0) {
-            throw new IllegalArgumentException("Zero key not allowed");
+    // key[i]==0 is used to signal that the i-th position is unset.
+    // Thus, we map key=key+1, so the user can still use key=0.
+    private static long remapKey(long key) {
+        if (key < 0) {
+            throw new IllegalArgumentException("Negative keys not allowed");
         }
+        return key + 1;
+    }
+
+    public void put(long key, T value) {
+        key = remapKey(key);
 
         int mask = keys.length - 1;
         int i = hashCode(key) & mask;
@@ -56,6 +67,8 @@ public void put(long key, T value) {
 
     @SuppressWarnings("unchecked")
     public T get(long key) {
+        key = remapKey(key);
+
         int mask = keys.length - 1;
         int i = hashCode(key) & mask;
         while (keys[i] != key && keys[i] != 0) {
@@ -68,7 +81,8 @@ public T get(long key) {
     public void forEach(Visitor<T> visitor) {
         for (int i = 0; i < keys.length; i++) {
             if (keys[i] != 0) {
-                visitor.visit(keys[i], (T) values[i]);
+                // Map key back, see remapKey
+                visitor.visit(keys[i] - 1, (T) values[i]);
             }
         }
     }

spark-common/src/main/java/me/lucko/spark/common/sampler/async/jfr/JfrReader.java

@@ -26,6 +26,7 @@
 /**
  * Parses JFR output produced by async-profiler.
  */
+@SuppressWarnings("UnqualifiedFieldAccess")
 public class JfrReader implements Closeable {
     private static final int BUFFER_SIZE = 2 * 1024 * 1024;
     private static final int CHUNK_HEADER_SIZE = 68;
@@ -53,7 +54,7 @@ public class JfrReader implements Closeable {
 
     public final Dictionary<JfrClass> types = new Dictionary<>();
     public final Map<String, JfrClass> typesByName = new HashMap<>();
-    public final Map<Long, String> threads = new HashMap<>(); // spark - convert to map
+    public final Dictionary<String> threads = new Dictionary<>();
     public final Dictionary<ClassRef> classes = new Dictionary<>();
     public final Dictionary<String> strings = new Dictionary<>();
     public final Dictionary<byte[]> symbols = new Dictionary<>();
@@ -75,6 +76,8 @@ public class JfrReader implements Closeable {
     private int monitorEnter;
     private int threadPark;
     private int activeSetting;
+    private int malloc;
+    private int free;
 
     public JfrReader(Path path) throws IOException { // spark - Path instead of String
         this.ch = FileChannel.open(path, StandardOpenOption.READ); // spark - Path instead of String
@@ -118,6 +121,10 @@ public long durationNanos() {
         return endNanos - startNanos;
     }
 
+    public long chunkDurationNanos() {
+        return chunkEndNanos - chunkStartNanos;
+    }
+
     public <E extends Event> void registerEvent(String name, Class<E> eventClass) {
         JfrClass type = typesByName.get(name);
         if (type != null) {
@@ -176,6 +183,10 @@ public <E extends Event> E readEvent(Class<E> cls) throws IOException {
                 if (cls == null || cls == AllocationSample.class) return (E) readAllocationSample(true);
             } else if (type == allocationOutsideTLAB || type == allocationSample) {
                 if (cls == null || cls == AllocationSample.class) return (E) readAllocationSample(false);
+            } else if (type == malloc) {
+                if (cls == null || cls == MallocEvent.class) return (E) readMallocEvent(true);
+            } else if (type == free) {
+                if (cls == null || cls == MallocEvent.class) return (E) readMallocEvent(false);
             } else if (type == liveObject) {
                 if (cls == null || cls == LiveObject.class) return (E) readLiveObject();
             } else if (type == monitorEnter) {
@@ -223,6 +234,15 @@ private AllocationSample readAllocationSample(boolean tlab) {
         return new AllocationSample(time, tid, stackTraceId, classId, allocationSize, tlabSize);
     }
 
+    private MallocEvent readMallocEvent(boolean hasSize) {
+        long time = getVarlong();
+        int tid = getVarint();
+        int stackTraceId = getVarint();
+        long address = getVarlong();
+        long size = hasSize ? getVarlong() : 0;
+        return new MallocEvent(time, tid, stackTraceId, address, size);
+    }
+
     private LiveObject readLiveObject() {
         long time = getVarlong();
         int tid = getVarint();
@@ -409,7 +429,7 @@ private void readConstants(JfrClass type) {
     }
 
     private void readThreads(int fieldCount) {
-        int count = getVarint(); // spark - don't preallocate
+        int count = threads.preallocate(getVarint());
         for (int i = 0; i < count; i++) {
             long id = getVarlong();
             String osName = getString();
@@ -543,6 +563,8 @@ private void cacheEventTypes() {
         monitorEnter = getTypeId("jdk.JavaMonitorEnter");
         threadPark = getTypeId("jdk.ThreadPark");
         activeSetting = getTypeId("jdk.ActiveSetting");
+        malloc = getTypeId("profiler.Malloc");
+        free = getTypeId("profiler.Free");
 
         registerEvent("jdk.CPULoad", CPULoad.class);
         registerEvent("jdk.GCHeapSummary", GCHeapSummary.class);
@@ -966,6 +988,22 @@ public long value() {
         }
     }
 
+    static class MallocEvent extends Event {
+        public final long address;
+        public final long size;
+
+        public MallocEvent(long time, int tid, int stackTraceId, long address, long size) {
+            super(time, tid, stackTraceId);
+            this.address = address;
+            this.size = size;
+        }
+
+        @Override
+        public long value() {
+            return size;
+        }
+    }
+
     static class ObjectCount extends Event {
         public final int gcId;
         public final int classId;