Retained size by ClassLoaders

CMakeLists.txt

@@ -36,6 +36,8 @@ add_library(memory_agent SHARED
         src/sizes/tag_info_array.cpp
         src/sizes/retained_size_by_classes.h
         src/sizes/retained_size_by_classes.cpp
+        src/sizes/retained_size_by_classloaders.h
+        src/sizes/retained_size_by_classloaders.cpp
         src/sizes/retained_size_action.h
         src/sizes/retained_size_by_objects.h
         src/sizes/retained_size_by_objects.cpp

src/agent.cpp

@@ -14,6 +14,7 @@
 #include "sizes/retained_size_and_held_objects.h"
 #include "sizes/retained_size_by_classes.h"
 #include "sizes/retained_size_by_objects.h"
+#include "sizes/retained_size_by_classloaders.h"
 #include "allocation_sampling.h"
 
 #pragma ide diagnostic ignored "OCUnusedGlobalDeclarationInspection"
@@ -140,6 +141,13 @@ JNIEXPORT jboolean JNICALL Java_com_intellij_memory_agent_IdeaNativeAgentProxy_c
     return (jboolean) 1;
 }
 
+extern "C"
+JNIEXPORT jboolean JNICALL Java_com_intellij_memory_agent_IdeaNativeAgentProxy_canGetRetainedSizeByClassLoaders(
+        JNIEnv *env,
+        jobject thisObject) {
+    return (jboolean) 1;
+}
+
 extern "C"
 JNIEXPORT jboolean JNICALL Java_com_intellij_memory_agent_IdeaNativeAgentProxy_canGetShallowSizeByClasses(
         JNIEnv *env,
@@ -196,6 +204,14 @@ JNIEXPORT jobjectArray JNICALL Java_com_intellij_memory_agent_IdeaNativeAgentPro
     return RetainedSizeByClassesAction(env, gdata->jvmti, thisObject).run(classesArray);
 }
 
+extern "C"
+JNIEXPORT jobjectArray JNICALL Java_com_intellij_memory_agent_IdeaNativeAgentProxy_getRetainedSizeByClassLoaders(
+        JNIEnv *env,
+        jobject thisObject,
+        jobjectArray classLoadersArray) {
+    return RetainedSizeByClassLoadersAction(env, gdata->jvmti, thisObject).run(classLoadersArray);
+}
+
 extern "C"
 JNIEXPORT jobjectArray JNICALL Java_com_intellij_memory_agent_IdeaNativeAgentProxy_getShallowAndRetainedSizeByClasses(
         JNIEnv *env,

src/sizes/retained_size_action.h

@@ -72,6 +72,24 @@ class RetainedSizeAction : public MemoryAgentAction<RESULT_TYPE, jobjectArray> {
         return JVMTI_ERROR_NONE;
     }
 
+    jvmtiError createTagsForClassLoadersClasses(JNIEnv *env, jvmtiEnv *jvmti, jobjectArray classesArray, jsize classLoaderIndex) {
+        for (jsize i = 0; i < this->env->GetArrayLength(classesArray); i++) {
+            jobject classObject = this->env->GetObjectArrayElement(classesArray, i);
+            jvmtiError err = tagClass(env, jvmti, classObject, [classLoaderIndex](jlong oldTag) -> jlong {
+                ClassTag *classTag = tagToClassTagPointer(oldTag);
+                if (classTag != nullptr) {
+                    classTag->ids.push_back(classLoaderIndex);
+                } else {
+                    return pointerToTag(ClassTag::create(static_cast<query_size_t>(classLoaderIndex)));
+                }
+
+                return 0;
+            });
+            if (err != JVMTI_ERROR_NONE) return err;
+        }
+        return JVMTI_ERROR_NONE;
+    }
+
     jvmtiError tagObjectsOfClasses(jobjectArray classesArray) {
         debug("tag objects of classes");
         jvmtiError err = createTagsForClasses(this->env, this->jvmti, classesArray);
@@ -80,6 +98,14 @@ class RetainedSizeAction : public MemoryAgentAction<RESULT_TYPE, jobjectArray> {
         return this->IterateThroughHeap(0, nullptr, tagObjectOfTaggedClass, nullptr);
     }
 
+    jvmtiError tagObjectsOfClassLoaderClasses(jobjectArray classesArray, jsize classLoaderIndex) {
+        debug("tag objects of classes");
+        jvmtiError err = createTagsForClassLoadersClasses(this->env, this->jvmti, classesArray, classLoaderIndex);
+        if (err != JVMTI_ERROR_NONE) return err;
+
+        return this->IterateThroughHeap(0, nullptr, tagObjectOfTaggedClass, nullptr);
+    }
+
     jvmtiError tagHeap() {
         jvmtiError err = this->FollowReferences(0, nullptr, nullptr, getTagsWithNewInfo, nullptr, "find objects with new info");
         if (err != JVMTI_ERROR_NONE) return err;

src/sizes/retained_size_by_classes.cpp

@@ -21,6 +21,11 @@ jint JNICALL visitObject(jlong classTag, jlong size, jlong *tagPtr, jint length,
     return JVMTI_ITERATION_CONTINUE;
 }
 
+jint JNICALL countHeapSize(jlong classTag, jlong size, jlong *tagPtr, jint length, void *userData) {
+    *reinterpret_cast<jlong*>(userData) += size;
+    return JVMTI_ITERATION_CONTINUE;
+}
+
 
 jint JNICALL visitObjectForShallowAndRetainedSize(jlong classTag, jlong size, jlong *tagPtr, jint length, void *userData) {
     if (*tagPtr == 0) {

src/sizes/retained_size_by_classes.h

@@ -26,5 +26,6 @@ class RetainedAndShallowSizeByClassesAction : public RetainedSizeAction<jobjectA
 };
 
 jint JNICALL visitObject(jlong classTag, jlong size, jlong *tagPtr, jint length, void *userData);
+jint JNICALL countHeapSize(jlong classTag, jlong size, jlong *tagPtr, jint length, void *userData);
 
 #endif //MEMORY_AGENT_RETAINED_SIZE_BY_CLASSES_H

src/sizes/retained_size_by_classloaders.cpp

@@ -0,0 +1,102 @@
+// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+
+#include <vector>
+#include "retained_size_by_classloaders.h"
+#include <algorithm>
+#include <numeric>
+
+
+jvmtiError RetainedSizeByClassLoadersAction::tagObjectsByClassLoader(jclass *classes, jint* cnt, jobject classLoader, jsize classLoaderIndex) {
+    jvmtiError err = JVMTI_ERROR_NONE;
+    std::vector<jclass> filteredClasses;
+    for (jsize i = 0; i < *cnt; i++) {
+        jobject rootClassLoader = getClassLoader(env, classes[i]);
+        if (!env->IsSameObject(rootClassLoader, NULL)) {
+            if (isEqual(env, classLoader, rootClassLoader)) {
+                filteredClasses.emplace_back(classes[i]);
+            }
+        }
+    }
+
+    err = tagObjectsOfClassLoaderClasses(toJavaArray(env, filteredClasses), classLoaderIndex);
+    if (!isOk(err)) return err;
+    if (shouldStopExecution()) return MEMORY_AGENT_INTERRUPTED_ERROR;
+    return err;
+}
+
+jlong RetainedSizeByClassLoadersAction::tagOtherObjects(jclass *classes, jint* cnt, jobjectArray classLoadersArray) {
+    jvmtiError err = JVMTI_ERROR_NONE;
+    std::vector<jclass> filteredClasses;
+    for (jsize i = 0; i < *cnt; i++) {
+        jobject rootClassLoader = getClassLoader(env, classes[i]);
+        if (!env->IsSameObject(rootClassLoader, NULL)) {
+            for (jsize j = 0; j < env->GetArrayLength(classLoadersArray); j++) {
+                if (isEqual(env, env->GetObjectArrayElement(classLoadersArray, j), rootClassLoader)) {
+                    filteredClasses.emplace_back(classes[i]);
+                    break;
+                }
+            }
+        }
+    }
+
+    std::vector<jclass> nonfilteredClasses;
+    for (jsize i = 0; i < *cnt; i++) {
+        if (std::find(filteredClasses.begin(), filteredClasses.end(), classes[i]) == filteredClasses.end()) {
+            nonfilteredClasses.emplace_back(classes[i]);
+        }
+    }
+    std::cout << "SIZE: " << nonfilteredClasses.size() << std::endl;
+
+    err = tagObjectsOfClassLoaderClasses(toJavaArray(env, nonfilteredClasses), env->GetArrayLength(classLoadersArray));
+}
+
+jlong RetainedSizeByClassLoadersAction::getHeapSize(){
+    jlong totalSize = 0;
+    IterateThroughHeap(0, nullptr, countHeapSize, &totalSize,
+                       "calculate heap size");
+    return totalSize;
+}
+
+RetainedSizeByClassLoadersAction::RetainedSizeByClassLoadersAction(JNIEnv *env, jvmtiEnv *jvmti, jobject object)
+        : RetainedSizeAction(env, jvmti, object) {
+
+}
+
+jvmtiError RetainedSizeByClassLoadersAction::getRetainedSizeByClassLoaders(jobjectArray classLoadersArray,
+                                                                           std::vector <jlong> &result) {
+    jvmtiError err = JVMTI_ERROR_NONE;
+    jclass *classes;
+    jint cnt;
+    err = jvmti->GetLoadedClasses(&cnt, &classes);
+    std::cout << cnt << std::endl;
+    for (jsize i = 0; i < env->GetArrayLength(classLoadersArray); i++) {
+        jvmtiError err = tagObjectsByClassLoader(classes, &cnt, env->GetObjectArrayElement(classLoadersArray, i), i);
+    }
+    tagOtherObjects(classes, &cnt, classLoadersArray);
+
+    if (!isOk(err)) return err;
+    if (shouldStopExecution()) return MEMORY_AGENT_INTERRUPTED_ERROR;
+
+    err = tagHeap();
+    if (!isOk(err)) return err;
+    if (shouldStopExecution()) return MEMORY_AGENT_INTERRUPTED_ERROR;
+
+    result.resize(env->GetArrayLength(classLoadersArray) + 1);
+    err = IterateThroughHeap(JVMTI_HEAP_FILTER_UNTAGGED, nullptr, visitObject, result.data(),
+                             "calculate retained size");
+    return err;
+}
+
+jlongArray RetainedSizeByClassLoadersAction::executeOperation(jobjectArray classLoadersArray) {
+    std::vector <jlong> result;
+    jvmtiError err = getRetainedSizeByClassLoaders(classLoadersArray, result);
+    if (!isOk(err)) {
+        handleError(jvmti, err, "Could not estimate retained size by classLoaders");
+        return env->NewLongArray(0);
+    }
+    std::cout << "HEAP SIZE: " << getHeapSize() << std::endl;
+    std::cout << "CLASSLOADERS CLASSES: " << std::accumulate(result.begin(), result.end()-1, 0) <<std::endl;
+    std::cout << "OTHER CLASSES: " << result.back() <<std::endl;
+    result.pop_back();
+    return toJavaArray(env, result);
+}

src/sizes/retained_size_by_classloaders.h

@@ -0,0 +1,27 @@
+// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+
+#ifndef MEMORY_AGENT_RETAINED_SIZE_BY_CLASSLOADERS_H
+#define MEMORY_AGENT_RETAINED_SIZE_BY_CLASSLOADERS_H
+
+#include "../memory_agent_action.h"
+#include "retained_size_action.h"
+#include "retained_size_by_classes.h"
+
+
+class RetainedSizeByClassLoadersAction : public RetainedSizeAction<jlongArray> {
+public:
+    RetainedSizeByClassLoadersAction(JNIEnv *env, jvmtiEnv *jvmti, jobject object);
+
+private:
+    jlongArray executeOperation(jobjectArray classLoadersArray) override;
+    jvmtiError getRetainedSizeByClassLoaders(jobjectArray classLoadersArray, std::vector<jlong> &result);
+    jvmtiError tagObjectsByClassLoader(jclass *classes, jint* cnt, jobject classLoader, jsize classLoaderIndex);
+    jlong getHeapSize();
+    jlong tagOtherObjects(jclass *classes, jint* cnt, jobjectArray classLoadersArray);
+};
+
+jint JNICALL visitObject(jlong classTag, jlong size, jlong *tagPtr, jint length, void *userData);
+
+
+
+#endif //MEMORY_AGENT_RETAINED_SIZE_BY_CLASSLOADERS_H

src/utils.cpp

@@ -49,6 +49,15 @@ jobjectArray toJavaArray(JNIEnv *env, std::vector<jobject> &objects) {
     return res;
 }
 
+jobjectArray toJavaArray(JNIEnv *env, std::vector <jclass> &objects) {
+    auto count = static_cast<jsize>(objects.size());
+    jobjectArray res = env->NewObjectArray(count, env->FindClass("java/lang/Class"), nullptr);
+    for (auto i = 0; i < count; ++i) {
+        env->SetObjectArrayElement(res, i, objects[i]);
+    }
+    return res;
+}
+
 jlongArray toJavaArray(JNIEnv *env, std::vector<jlong> &items) {
     auto count = static_cast<jsize>(items.size());
     jlongArray result = env->NewLongArray(count);
@@ -260,6 +269,18 @@ jvmtiError tagClassAndItsInheritors(JNIEnv *env, jvmtiEnv *jvmti, jobject classO
     return err;
 }
 
+jvmtiError tagClass(JNIEnv *env, jvmtiEnv *jvmti, jobject classObject, std::function<jlong(jlong)> &&createTag) {
+    jlong oldTag;
+    jvmtiError err = jvmti->GetTag(classObject, &oldTag);
+    if (err != JVMTI_ERROR_NONE) return err;
+    jlong newTag = createTag(oldTag);
+    if (newTag != 0) {
+        err = jvmti->SetTag(classObject, newTag);
+        if (err != JVMTI_ERROR_NONE) return err;
+    }
+    return err;
+}
+
 jmethodID getIsAssignableFromMethod(JNIEnv *env) {
     jclass langClass = env->FindClass("java/lang/Class");
     return env->GetMethodID(langClass, "isAssignableFrom", "(Ljava/lang/Class;)Z");
@@ -270,6 +291,17 @@ std::string getToString(JNIEnv *env, jobject object) {
     return jstringTostring(env, reinterpret_cast<jstring>(name));
 }
 
+jobject getClassLoader(JNIEnv *env, jclass clazz) {
+    return env->CallObjectMethod(clazz, env->GetMethodID(env->GetObjectClass(clazz), "getClassLoader",
+                                                         "()Ljava/lang/ClassLoader;"));
+}
+
+bool isEqual(JNIEnv *env, jobject obj, jobject otherObj) {
+    jclass objClass = env->GetObjectClass(obj);
+    jmethodID equalsID = env->GetMethodID(objClass, "equals", "(Ljava/lang/Object;)Z");
+    return env->CallBooleanMethod(obj, equalsID, otherObj);
+}
+
 ThreadSuspender::ThreadSuspender(jvmtiEnv *jvmti) : jvmti(jvmti) {
     jint threadCnt;
     jthread *threads;

src/utils.h

@@ -27,6 +27,8 @@ jbooleanArray toJavaArray(JNIEnv *env, std::vector<jboolean> &items);
 
 jobjectArray toJavaArray(JNIEnv *env, std::vector<jobject> &objects);
 
+jobjectArray toJavaArray(JNIEnv *env, std::vector<jclass> &objects);
+
 jintArray toJavaArray(JNIEnv *env, std::vector<jint> &items);
 
 jlongArray toJavaArray(JNIEnv *env, std::vector<jlong> &items);
@@ -65,6 +67,8 @@ jvmtiError cleanHeapAndGetObjectsByTags(jvmtiEnv *jvmti, std::vector<jlong> &tag
 
 jvmtiError tagClassAndItsInheritors(JNIEnv *env, jvmtiEnv *jvmti, jobject classObject, std::function<jlong (jlong)> &&createTag);
 
+jvmtiError tagClass(JNIEnv *env, jvmtiEnv *jvmti, jobject classObject, std::function<jlong (jlong)> &&createTag);
+
 bool isOk(jvmtiError error);
 
 bool fileExists(const std::string &fileName);
@@ -73,6 +77,10 @@ std::string jstringTostring(JNIEnv *env, jstring jStr);
 
 jmethodID getIsAssignableFromMethod(JNIEnv *env);
 
+jobject getClassLoader(JNIEnv *env, jclass obj);
+
+bool isEqual(JNIEnv *env, jobject obj1, jobject obj2);
+
 std::string getToString(JNIEnv *env, jobject klass);
 
 template<typename T>

test_data/proxy/src/com/intellij/memory/agent/IdeaNativeAgentProxy.java

@@ -28,6 +28,8 @@ public IdeaNativeAgentProxy(Object cancellationFileName, Object progressFileName
 
   public native boolean canGetRetainedSizeByClasses();
 
+  public native boolean canGetRetainedSizeByClassLoaders();
+
   public native boolean canGetShallowSizeByClasses();
 
   public native boolean canEstimateObjectsSizes();
@@ -38,6 +40,8 @@ public IdeaNativeAgentProxy(Object cancellationFileName, Object progressFileName
 
   public native Object[] getRetainedSizeByClasses(Object[] classes);
 
+  public native Object[] getRetainedSizeByClassLoaders(ClassLoader[] classloaders);
+
   public native Object[] getShallowAndRetainedSizeByClasses(Object[] classes);
 
   public native Object[] findPathsToClosestGcRoots(Object object, int pathsNumber, int objectsNumber);

test_data/proxy/src/com/intellij/memory/agent/MemoryAgent.java

@@ -97,6 +97,14 @@ public synchronized <T> T[] getAllReachableObjects(Object startObject, Class<T>
         return resultArray;
     }
 
+    /**
+         * TODO
+         */
+        @SuppressWarnings("unchecked")
+        public synchronized long[] getRetainedSizeByClassLoaders(ClassLoader[] classloaders) throws MemoryAgentExecutionException {
+            return (long [])getResult(callProxyMethod(() -> proxy.getRetainedSizeByClassLoaders(classloaders)));
+        }
+
     /**
      * Adds an allocation listener that catches allocation sampling events for specified classes.
      *