Add debug option to specify MALLOC_ARENA_MAX (#108)

build.gradle

@@ -10,7 +10,7 @@ buildscript {
 }
 
 plugins {
-  id 'io.github.gradle-nexus.publish-plugin' version '1.1.0'
+  id 'io.github.gradle-nexus.publish-plugin' version '2.0.0'
   id "com.diffplug.spotless" version "6.11.0"
 }
 

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

@@ -254,3 +254,8 @@ extern "C" DLLEXPORT jlong JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_tscFrequency0(JNIEnv* env, jobject unused) {
     return TSC::frequency();
 }
+
+extern "C" DLLEXPORT void JNICALL
+Java_com_datadoghq_profiler_JavaProfiler_mallocArenaMax0(JNIEnv* env, jobject unused, jint maxArenas) {
+    OS::mallocArenaMax(maxArenas);
+}

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

@@ -93,6 +93,8 @@ class OS {
     static int fileSize(int fd);
     static int truncateFile(int fd);
     static void freePageCache(int fd, off_t start_offset);
+
+    static void mallocArenaMax(int arena_max);
 };
 
 #endif // _OS_H

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

@@ -35,6 +35,9 @@
 #include <unistd.h>
 #include "os.h"
 
+#ifndef __musl__
+#include <malloc.h>
+#endif
 
 #ifdef __LP64__
 #  define MMAP_SYSCALL __NR_mmap
@@ -360,4 +363,10 @@ void OS::freePageCache(int fd, off_t start_offset) {
     posix_fadvise(fd, start_offset & ~page_mask, 0, POSIX_FADV_DONTNEED);
 }
 
+void OS::mallocArenaMax(int arena_max) {
+#ifndef __musl__
+    mallopt(M_ARENA_MAX, arena_max);
+#endif
+}
+
 #endif // __linux__

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

@@ -342,4 +342,8 @@ void OS::freePageCache(int fd, off_t start_offset) {
     // Not supported on macOS
 }
 
+void OS::mallocArenaMax(int arena_max) {
+    // Not supported on macOS
+}
+
 #endif // __APPLE__

ddprof-lib/src/main/java/com/datadoghq/profiler/JavaProfiler.java

@@ -42,7 +42,7 @@
 public final class JavaProfiler {
     private static final String NATIVE_LIBS = "/META-INF/native-libs";
     private static final String LIBRARY_NAME = "libjavaProfiler." + (OperatingSystem.current() == OperatingSystem.macos ? "dylib" : "so");
-    
+
     private static final Unsafe UNSAFE;
     static {
         Unsafe unsafe = null;
@@ -122,6 +122,16 @@ public static synchronized JavaProfiler getInstance(String libLocation, String s
         System.load(libLocation);
         profiler.initializeContextStorage();
         instance = profiler;
+
+        String maxArenaValue = System.getProperty("ddprof.debug.malloc_arena_max");
+        if (maxArenaValue != null) {
+            try {
+                mallocArenaMax0(Integer.parseInt(maxArenaValue));
+            } catch (NumberFormatException e) {
+                System.out.println("[WARN] Invalid value for ddprof.debug.malloc_arena_max: " + maxArenaValue + ". Expecting an integer.");
+            }
+        }
+
         return profiler;
     }
 
@@ -139,7 +149,7 @@ public static synchronized JavaProfiler getInstance(String libLocation, String s
      */
     private static Path libraryFromClasspath(OperatingSystem os, Arch arch, String qualifier, Path tempDir) throws IOException {
         String resourcePath = NATIVE_LIBS + "/" + os.name().toLowerCase() + "-" + arch.name().toLowerCase() + ((qualifier != null && !qualifier.isEmpty()) ? "-" + qualifier : "") + "/" + LIBRARY_NAME;
-        
+
         InputStream libraryData =  JavaProfiler.class.getResourceAsStream(resourcePath);
 
         if (libraryData != null) {
@@ -433,7 +443,7 @@ public boolean isThresholdExceeded(long thresholdMillis, long startTicks, long e
      * @param origin the thread the task was submitted on
      */
     public void recordQueueTime(long startTicks, long endTicks, Class<?> task, Class<?> scheduler,
-                               Thread origin) {
+                                Thread origin) {
         recordQueueEnd0(startTicks, endTicks, task.getName(), scheduler.getName(), origin);
     }
 
@@ -499,7 +509,7 @@ static boolean isMuslProcSelfMaps() throws IOException {
      * However, if such string is missing should indicate that the system is not a musl one.
      */
     static boolean isMuslJavaExecutable() throws IOException {
-        
+
         byte[] magic = new byte[]{(byte)0x7f, (byte)'E', (byte)'L', (byte)'F'};
         byte[] prefix = new byte[]{(byte)'/', (byte)'l', (byte)'d', (byte)'-'}; // '/ld-*'
         byte[] musl = new byte[]{(byte)'m', (byte)'u', (byte)'s', (byte)'l'}; // 'musl'
@@ -573,4 +583,6 @@ private static boolean containsArray(byte[] container, int offset, byte[] contai
     private static native long currentTicks0();
 
     private static native long tscFrequency0();
+
+    private static native void mallocArenaMax0(int max);
 }

ddprof-test/src/test/java/com/datadoghq/profiler/context/TagContextTest.java

@@ -68,6 +68,8 @@ public void test() throws InterruptedException {
         long sum = 0;
         long[] weights = new long[strings.length];
         for (int i = 0; i < strings.length; i++) {
+            AtomicLong weight = weightsByTagValue.get(strings[i]);
+            assertNotNull(weight, "Weight for " + strings[i] + " not found");
             weights[i] = weightsByTagValue.get(strings[i]).get();
             sum += weights[i];
         }