Add native JVMTI agent with allocation sampling

agent/native/build.gradle.kts

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2020 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import org.gradle.internal.jvm.Jvm
+
+// Note: java-library and cpp-library cannot be applied simultaneously
+// (see https://github.com/gradle/gradle-native/issues/352), which is why
+// we use two separate build.gradle files for the agent.
+plugins {
+    id("cpp-library")
+}
+
+group = "com.google.idea.perf"
+
+library {
+    baseName.set("agent")
+
+    binaries.configureEach {
+        val compileTask = compileTask.get()
+
+        // JNI and JVMTI headers.
+        val javaHome = Jvm.current().javaHome
+        val osFamily = targetPlatform.targetMachine.operatingSystemFamily
+        val osDirName = when {
+            osFamily.isLinux -> "linux"
+            osFamily.isMacOs -> "darwin"
+            osFamily.isWindows -> "win32"
+            else -> error("Unknown OS: $osFamily")
+        }
+        compileTask.includes.from(
+            "$javaHome/include",
+            "$javaHome/include/$osDirName"
+        )
+
+        // Compiler args.
+        val compilerArgs = when (toolChain) {
+            is GccCompatibleToolChain -> {
+                arrayOf("-std=c++17", "-Wall", "-Wpedantic", "-Werror")
+            }
+            is VisualCpp -> {
+                arrayOf("/std:c++17")
+            }
+            else -> error("Unknown toolchain: $toolChain")
+        }
+        compileTask.compilerArgs.addAll(*compilerArgs)
+    }
+}

agent/native/src/main/cpp/agent.cpp

@@ -0,0 +1,127 @@
+/*
+ * Copyright 2020 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <jni.h>
+#include <jvmti.h>
+#include <stdio.h>
+#include <algorithm>
+
+// This JVMTI agent gathers performance data for use in IDE Perf.
+//
+// JVMTI docs: https://docs.oracle.com/en/java/javase/11/docs/specs/jvmti.html
+// JNI docs: https://docs.oracle.com/en/java/javase/11/docs/specs/jni/index.html
+
+static bool
+HandleJvmtiError(jvmtiEnv* jvmti, jvmtiError err, const char* msg) {
+    if (err == JVMTI_ERROR_NONE) {
+        return false;
+    } else {
+        char* name = nullptr;
+        jvmti->GetErrorName(err, &name);
+        auto desc = name != nullptr ? name : "Unknown";
+        fprintf(stderr, "JVMTI error: %d(%s) %s\n", err, desc, msg);
+        jvmti->Deallocate((unsigned char*)name);
+        return true;
+    }
+}
+
+constexpr jint kHeapSamplingInterval = 256 * 1024;
+constexpr jlong kHeapSamplingIntervalLong = kHeapSamplingInterval;
+
+// A thread-local estimate of heap allocations in bytes.
+//
+// TODO: The use of thread_local linkage here assumes that Java threads have a
+// 1-to-1 correspondence with native threads. That assumption will be invalid
+// after Project Loom. One alternative is to use jvmti->GetThreadLocalStorage,
+// but unfortunately the overhead is somewhat high. Maybe a better alternative
+// is to switch to Java-land in the SampledObjectAlloc callback and just
+// store the counter in a normal ThreadLocal variable.
+static thread_local jlong thread_allocations = 0;
+
+extern "C" void JNICALL
+SampledObjectAlloc(
+    jvmtiEnv* jvmti,
+    JNIEnv* jni,
+    jthread thread,
+    jobject object,
+    jclass klass,
+    jlong size
+) {
+    // Note: in the case where the allocation overlaps with the boundary between
+    // two heap sampling intervals, JVMTI does not specify whether the
+    // "excess bytes" affect the length of the next sampling interval.
+    // So, we have to live with the possibility that those excess bytes are
+    // lost and that we slightly underestimate allocations.
+    thread_allocations += std::max(kHeapSamplingIntervalLong, size);
+}
+
+extern "C" JNIEXPORT jlong JNICALL
+Java_com_google_idea_perf_AllocationSampling_countAllocationsForCurrentThread(
+    JNIEnv* jni,
+    jobject thiz
+) {
+    return thread_allocations;
+}
+
+// This function is called by the JVM when the agent is loaded after startup.
+extern "C" JNIEXPORT jint JNICALL
+Agent_OnAttach(JavaVM* vm, char* options, void* reserved) {
+    jvmtiError err;
+
+    jvmtiEnv* jvmti;
+    if (vm->GetEnv((void**)&jvmti, JVMTI_VERSION_11) != JNI_OK) {
+        fprintf(stderr, "Error retrieving JVMTI function table\n");
+        return JNI_ERR;
+    }
+
+    // Set JVMTI capabilities.
+    jvmtiCapabilities capabilities = {};
+    capabilities.can_generate_sampled_object_alloc_events = 1;
+    err = jvmti->AddCapabilities(&capabilities);
+    if (HandleJvmtiError(jvmti, err, "Failed to add JVMTI capabilities")) {
+        return JNI_ERR;
+    }
+
+    // Set JVMTI callbacks.
+    jvmtiEventCallbacks callbacks = {};
+    callbacks.SampledObjectAlloc = SampledObjectAlloc;
+    err = jvmti->SetEventCallbacks(&callbacks, sizeof(callbacks));
+    if (HandleJvmtiError(jvmti, err, "Failed to set JVMTI callbacks")) {
+        return JNI_ERR;
+    }
+
+    // Set heap sampling interval.
+    err = jvmti->SetHeapSamplingInterval(kHeapSamplingInterval);
+    if (HandleJvmtiError(jvmti, err, "Failed to set heap sampling interval")) {
+        return JNI_ERR;
+    }
+
+    // Enable JVMTI events.
+    err = jvmti->SetEventNotificationMode(JVMTI_ENABLE,
+                                          JVMTI_EVENT_SAMPLED_OBJECT_ALLOC,
+                                          /*thread*/ nullptr);
+    if (HandleJvmtiError(jvmti, err, "Failed to enable JVMTI events")) {
+        return JNI_ERR;
+    }
+
+    return JNI_OK;
+}
+
+// This function is called by the JVM when the agent is loaded at startup.
+extern "C" JNIEXPORT jint JNICALL
+Agent_OnLoad(JavaVM* vm, char* options, void* reserved) {
+    return Agent_OnAttach(vm, options, reserved);
+}

build.gradle.kts

@@ -68,11 +68,14 @@ tasks.patchPluginXml {
 }
 
 val javaAgent: Configuration by configurations.creating
+val nativeAgent: Configuration by configurations.creating
 
 configureEach(tasks.prepareSandbox, tasks.prepareTestingSandbox) {
     // Copy agent artifacts into the plugin home directory.
     val agentDir = "$pluginName/agent"
     from(javaAgent) { into(agentDir) }
+    from(nativeAgent) { into(agentDir) }
+    if (isRelease) TODO("Figure out how to bundle native agent binaries for all three platforms")
 }
 
 tasks.publishPlugin {
@@ -115,9 +118,11 @@ tasks.test {
 fun JavaForkOptions.enableAgent() {
     val atStartup = project.findProperty("loadAgentAtStartup") == "true"
     if (atStartup) {
-        // Add the -javaagent startup flag.
+        // Add the -javaagent and -agentpath startup flags.
         jvmArgumentProviders.add(CommandLineArgumentProvider {
-            javaAgent.map { file -> "-javaagent:$file" }
+            val javaAgentFlags = javaAgent.map { file -> "-javaagent:$file" }
+            val nativeAgentFlags = nativeAgent.map { file -> "-agentpath:$file" }
+            javaAgentFlags + nativeAgentFlags
         })
     } else {
         // Let the agent load itself later.
@@ -127,11 +132,12 @@ fun JavaForkOptions.enableAgent() {
 
 dependencies {
     // Bundle the agent artifacts.
-    javaAgent(project(":agent", "runtimeElements"))
+    javaAgent(project(":agent:java", "runtimeElements"))
+    nativeAgent(project(":agent:native", "releaseRuntimeElements"))
 
     // Using 'compileOnly' because the agent is loaded in the boot classpath.
-    compileOnly(project(":agent"))
-    testCompileOnly(project(":agent"))
+    compileOnly(project(":agent:java"))
+    testCompileOnly(project(":agent:java"))
 
     implementation("org.ow2.asm:asm:8.0.1")
     implementation("org.ow2.asm:asm-util:8.0.1")

settings.gradle.kts

@@ -15,4 +15,4 @@
  */
 
 rootProject.name = "ide-perf"
-include("agent")
+include("agent:java", "agent:native")

src/main/java/com/google/idea/perf/AgentLoader.kt

@@ -28,6 +28,7 @@ import com.intellij.notification.Notification
 import com.intellij.notification.NotificationType
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.extensions.PluginId
+import com.intellij.openapi.util.SystemInfo
 import com.intellij.util.io.isFile
 import com.sun.tools.attach.VirtualMachine
 import java.lang.instrument.Instrumentation
@@ -54,9 +55,10 @@ object AgentLoader {
 
     val ensureJavaAgentLoaded: Boolean by lazy { doLoadJavaAgent() }
 
+    val ensureNativeAgentLoaded: Boolean by lazy { doLoadNativeAgent() }
+
     val ensureTracerHooksInstalled: Boolean by lazy { doInstallTracerHooks() }
 
-    // Note: this method can take around 200 ms or so.
     private fun doLoadJavaAgent(): Boolean {
         val agentLoadedAtStartup = try {
             // Until the agent is loaded, we cannot trigger symbol resolution for its
@@ -73,45 +75,79 @@ object AgentLoader {
         }
         else {
             try {
-                val overhead = measureTimeMillis { tryLoadAgentAfterStartup() }
+                val overhead = measureTimeMillis {
+                    tryLoadAgent("agent.jar", native = false)
+                }
                 LOG.info("Java agent was loaded on demand in $overhead ms")
             }
             catch (e: Throwable) {
-                val msg = """
-                    [Tracer] Failed to attach the instrumentation agent after startup.
-                    On JDK 9+, make sure jdk.attach.allowAttachSelf is set to true.
-                    Alternatively, you can attach the agent at startup via the -javaagent flag.
-                    """.trimIndent()
-                Notification("Tracer", "", msg, NotificationType.ERROR).notify(null)
-                LOG.warn(e)
+                var msg = "Failed to load the Java agent"
+                if (System.getProperty("jdk.attach.allowAttachSelf") == null) {
+                    msg += ". Please set the VM option jdk.attach.allowAttachSelf to true."
+                }
+                LOG.warn(msg, e)
                 return false
             }
         }
 
         // Disable tracing entirely if class retransformation is not supported.
         val instrumentation = checkNotNull(AgentMain.savedInstrumentationInstance)
         if (!instrumentation.isRetransformClassesSupported) {
-            val msg = "[Tracer] The current JVM configuration does not allow class retransformation"
-            Notification("Tracer", "", msg, NotificationType.ERROR).notify(null)
-            LOG.warn(msg)
+            LOG.warn("This JVM does not support class retransformations")
             return false
         }
 
         return true
     }
 
-    // Note: this method can throw a variety of exceptions.
-    private fun tryLoadAgentAfterStartup() {
+    private fun doLoadNativeAgent(): Boolean {
+        val agentLoadedAtStartup = try {
+            AllocationSampling.countAllocationsForCurrentThread()
+            true
+        }
+        catch (e: LinkageError) {
+            false
+        }
+
+        if (agentLoadedAtStartup) {
+            LOG.info("Native agent was loaded at startup")
+        }
+        else {
+            try {
+                val binary = when {
+                    SystemInfo.isMac -> "libagent.dylib"
+                    SystemInfo.isWindows -> "agent.dll"
+                    else -> "libagent.so"
+                }
+                val overhead = measureTimeMillis {
+                    tryLoadAgent(binary, native = true)
+                }
+                LOG.info("Native agent was loaded on demand in $overhead ms")
+            }
+            catch (e: Throwable) {
+                LOG.warn("Failed to load the native agent", e)
+                return false
+            }
+        }
+
+        return true
+    }
+
+    // Throws exceptions on failure.
+    private fun tryLoadAgent(fileName: String, native: Boolean) {
         val plugin = PluginManagerCore.getPlugin(PluginId.getId("com.google.ide-perf"))
             ?: error("Failed to find our own plugin")
-        val agentDir = plugin.pluginPath.resolve("agent")
 
-        val javaAgent = agentDir.resolve("agent.jar")
-        check(javaAgent.isFile()) { "Could not find agent.jar at $javaAgent" }
+        val path = plugin.pluginPath.resolve("agent").resolve(fileName)
+        check(path.isFile()) { "Could not find agent at $path" }
 
+        val absolutePath = path.toAbsolutePath().toString()
         val vm = VirtualMachine.attach(OSProcessUtil.getApplicationPid())
         try {
-            vm.loadAgent(javaAgent.toAbsolutePath().toString())
+            when {
+                native -> vm.loadAgentPath(absolutePath)
+                else -> vm.loadAgent(absolutePath)
+            }
         }
         finally {
             vm.detach()

src/main/java/com/google/idea/perf/AllocationSampling.kt

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2020 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.idea.perf
+
+object AllocationSampling {
+    /**
+     * Estimates total memory allocations for the current thread, in bytes.
+     *
+     * Important: you must invoke [AgentLoader.ensureNativeAgentLoaded] before calling
+     * this method, otherwise there maybe be linkage errors at runtime.
+     *
+     * Allocations are tracked using a native JVMTI agent subscribing to the
+     * JVMTI_EVENT_SAMPLED_OBJECT_ALLOC event. The allocation count is only an estimate;
+     * its accuracy depends on the heap sampling rate set by the agent.
+     */
+    external fun countAllocationsForCurrentThread(): Long
+}