Refactor retained and shallow sizes by class and by objects actions

Author: Nikita Nazarov

src/agent.cpp

@@ -221,20 +221,20 @@ JNIEXPORT jobjectArray JNICALL Java_com_intellij_memory_agent_IdeaNativeAgentPro
 }
 
 extern "C"
-JNIEXPORT jobjectArray JNICALL Java_com_intellij_memory_agent_IdeaNativeAgentProxy_getShallowAndRetainedSizeByObjects(
+JNIEXPORT jobjectArray JNICALL Java_com_intellij_memory_agent_IdeaNativeAgentProxy_getShallowAndRetainedSizesByObjects(
         JNIEnv *env,
         jobject thisObject,
         jobjectArray objects) {
-    return RetainedSizeViaDominatorTreeAction(env, gdata->jvmti, thisObject).run(objects);
+    return RetainedSizesViaDominatorTreeAction(env, gdata->jvmti, thisObject).run(objects);
 }
 
 extern "C"
-JNIEXPORT jobjectArray JNICALL Java_com_intellij_memory_agent_IdeaNativeAgentProxy_getShallowAndRetainedSizeByClass(
+JNIEXPORT jobjectArray JNICALL Java_com_intellij_memory_agent_IdeaNativeAgentProxy_getSortedShallowAndRetainedSizesByClass(
         JNIEnv *env,
         jobject thisObject,
         jobject classRef,
         jlong objectsLimit) {
-    return RetainedSizeByClassViaDominatorTreeAction(env, gdata->jvmti, thisObject).run(classRef, objectsLimit);
+    return RetainedSizesByClassViaDominatorTreeAction(env, gdata->jvmti, thisObject).run(classRef, objectsLimit);
 }
 
 extern "C"

src/sizes/dominator_tree.cpp

@@ -4,11 +4,14 @@
 #include "dominator_tree.h"
 
 namespace {
-    void link(long v, long w, std::vector<long> &ancestor) {
+    using jlongs = std::vector<jlong>;
+    using graph_t = std::vector<jlongs>;
+
+    void link(jlong v, jlong w, jlongs &ancestor) {
         ancestor[w] = v;
     }
 
-    void compress(long v, std::vector<long> &ancestor, std::vector<long> &label, std::vector<long> &semi) {
+    void compress(jlong v, jlongs &ancestor, jlongs &label, jlongs &semi) {
         if (ancestor[ancestor[v]] != -1) {
             compress(ancestor[v], ancestor, label, semi);
             if (semi[label[ancestor[v]]] < semi[label[v]]) {
@@ -18,7 +21,7 @@ namespace {
         }
     }
 
-    long eval(long v, std::vector<long> &ancestor, std::vector<long> &label, std::vector<long> &semi) {
+    jlong eval(jlong v, jlongs &ancestor, jlongs &label, jlongs &semi) {
         if (ancestor[v] == -1) {
             return v;
         }
@@ -27,14 +30,8 @@ namespace {
         return label[v];
     }
 
-    void dfs(
-        jlong v, jlong &n,
-        const std::vector<std::vector<jlong>> &graph,
-        std::vector <jlong> &semi,
-        std::vector <jlong> &parent,
-        std::vector <jlong> &vertex,
-        std::vector<std::vector<jlong>> &pred
-    ) {
+    void dfs(jlong v, jlong &n, const graph_t &graph, jlongs &semi,
+             jlongs &parent, jlongs &vertex, graph_t &pred) {
         semi[v] = n;
         vertex[n] = v;
         n++;
@@ -48,33 +45,25 @@ namespace {
     }
 }
 
-std::vector<jlong> calculateRetainedSizes(const std::vector<std::vector<jlong>> &graph, const std::vector<jlong> &sizes) {
-    size_t numOfVertices = graph.size();
-
-    std::vector<jlong> semi(numOfVertices);
-    std::vector<jlong> parent(numOfVertices);
-    std::vector<jlong> vertex(numOfVertices);
-    std::vector<std::vector<jlong>> pred(numOfVertices);
-    std::vector<std::vector<jlong>> bucket(numOfVertices);
-    std::vector<jlong> ancestor(numOfVertices);
-    std::vector<jlong> label(numOfVertices);
-    std::vector<jlong> dom(numOfVertices);
-    std::vector<jlong> retained_sizes(numOfVertices);
-    std::vector<jlong> child(numOfVertices);
-    std::vector<jlong> link_size(numOfVertices);
-
-    for (jlong i = 0; i < numOfVertices; i++) {
-        retained_sizes[i] = sizes[i];
+jlongs calculateRetainedSizesViaDominatorTree(const graph_t &graph, const jlongs &sizes) {
+    auto n = static_cast<jlong>(graph.size());
+    jlongs semi(n);
+    jlongs parent(n);
+    jlongs vertex(n);
+    jlongs ancestor(n);
+    jlongs label(n);
+    jlongs dom(n);
+    jlongs retainedSizes(n);
+    graph_t pred(n);
+    graph_t bucket(n);
+
+    for (jlong i = 0; i < n; i++) {
+        retainedSizes[i] = sizes[i];
         label[i] = i;
-        link_size[i] = 1;
         ancestor[i] = -1;
     }
 
-    link_size[0] = 0;
-    label[0] = 0;
-    semi[0] = 0;
-
-    jlong n = 0;
+    n = 0;
     dfs(0, n, graph, semi, parent, vertex, pred);
 
     for (jlong i = n - 1; i > 0; i--) {
@@ -100,7 +89,7 @@ std::vector<jlong> calculateRetainedSizes(const std::vector<std::vector<jlong>>
     }
 
     dom[0] = -1;
-    std::vector<jlong> childCount(numOfVertices);
+    jlongs childCount(n);
     for (jlong i = 1; i < n; i++) {
         jlong w = vertex[i];
         if (dom[w] != vertex[semi[w]]) {
@@ -124,12 +113,12 @@ std::vector<jlong> calculateRetainedSizes(const std::vector<std::vector<jlong>>
         leaves.pop();
         jlong d = dom[leaf];
         if (d >= 0) {
-            retained_sizes[d] += retained_sizes[leaf];
+            retainedSizes[d] += retainedSizes[leaf];
             if (--childCount[d] == 0) {
                 leaves.push(d);
             }
         }
     }
 
-    return retained_sizes;
+    return retainedSizes;
 }

src/sizes/dominator_tree.h

@@ -6,6 +6,7 @@
 #include <vector>
 #include <jni.h>
 
-std::vector<jlong> calculateRetainedSizes(const std::vector<std::vector<jlong>> &graph, const std::vector<jlong> &sizes);
+std::vector<jlong> calculateRetainedSizesViaDominatorTree(const std::vector<std::vector<jlong>> &graph,
+                                                          const std::vector<jlong> &sizes);
 
 #endif //MEMORY_AGENT_DOMINATOR_TREE_H

src/sizes/retained_size_via_dominator_tree.cpp

@@ -128,110 +128,113 @@ namespace {
 
         return JVMTI_VISIT_OBJECTS;
     }
+
+    jobjectArray getObjectArrayOfSize(JNIEnv *env, size_t size) {
+        return env->NewObjectArray(static_cast<jsize>(size), env->FindClass("java/lang/Object"), nullptr);
+    }
 }
 
-jobjectArray RetainedSizeViaDominatorTreeAction::executeOperation(jobjectArray objects) {
-    SizesViaDominatorTreeHeapDumpInfo info;
-    jvmtiError err = info.initAndSetTagsForObjects(env, jvmti, objects);
-    if (err != JVMTI_ERROR_NONE) return nullptr;
+template<typename RESULT_TYPE, typename... ARGS_TYPES>
+jvmtiError RetainedSizesAction<RESULT_TYPE, ARGS_TYPES...>::calculateRetainedSizes(jobjectArray objects,
+                                                                                   std::vector<jlong> &retainedSizes,
+                                                                                   SizesViaDominatorTreeHeapDumpInfo &info) {
+    jvmtiError err = info.initAndSetTagsForObjects(this->env, this->jvmti, objects);
+    if (!isOk(err)) return err;
 
     // We set a tag for the input array to ignore it during traversal
-    err = jvmti->SetTag(objects, MOCK_REFERRER_TAG);
-    if (err != JVMTI_ERROR_NONE) return nullptr;
+    err = this->jvmti->SetTag(objects, MOCK_REFERRER_TAG);
+    if (!isOk(err)) return err;
 
-    progressManager.updateProgress(10, "Traversing heap for the first time...");
+    this->progressManager.updateProgress(10, "Traversing heap for the first time...");
     logger::resetTimer();
-    err = FollowReferences(0, nullptr, nullptr, firstTraversal, &info, "first traversal");
+    err = this->FollowReferences(0, nullptr, nullptr, firstTraversal, &info);
     logger::logPassedTime();
-    if (err != JVMTI_ERROR_NONE || shouldStopExecution()) return nullptr;
+    if (!isOk(err) || this->shouldStopExecution()) return err;
 
     logger::resetTimer();
-    progressManager.updateProgress(60, "Traversing heap for the second time...");
-    err = FollowReferences(0, nullptr, objects, secondTraversal, &info, "second traversal");
+    this->progressManager.updateProgress(60, "Traversing heap for the second time...");
+    err = this->FollowReferences(0, nullptr, objects, secondTraversal, &info);
     logger::logPassedTime();
-    if (err != JVMTI_ERROR_NONE || shouldStopExecution()) return nullptr;
+    if (!isOk(err) || this->shouldStopExecution()) return err;
 
-    progressManager.updateProgress(80, "Calculating retained size...");
+    this->progressManager.updateProgress(80, "Calculating retained size...");
     info.setUpNeighboursForMasterNode();
-    std::vector<jlong> retained_sizes = calculateRetainedSizes(info.graph, info.sizes);
+    retainedSizes = calculateRetainedSizesViaDominatorTree(info.graph, info.sizes);
+
+    return err;
+}
+
+template<typename RESULT_TYPE, typename... ARGS_TYPES>
+RetainedSizesAction<RESULT_TYPE, ARGS_TYPES...>::RetainedSizesAction(JNIEnv *env, jvmtiEnv *jvmti, jobject object) :
+    MemoryAgentAction<RESULT_TYPE, ARGS_TYPES...>(env, jvmti, object) {
+
+}
+
+jobjectArray RetainedSizesViaDominatorTreeAction::executeOperation(jobjectArray objects) {
+    SizesViaDominatorTreeHeapDumpInfo info;
+    std::vector<jlong> retainedSizes;
+    jvmtiError err = calculateRetainedSizes(objects, retainedSizes, info);
+    if (!isOk(err) || shouldStopExecution()) return nullptr;
 
     progressManager.updateProgress(95, "Extracting answer...");
+    return constructResultObject(objects, retainedSizes, info);
+}
+
+jobjectArray RetainedSizesViaDominatorTreeAction::constructResultObject(jobjectArray objects,
+                                                                        const std::vector<jlong> &retainedSizes,
+                                                                        const SizesViaDominatorTreeHeapDumpInfo &info) {
     jsize size = env->GetArrayLength(objects);
     std::vector<jlong> shallowSizes;
-    std::vector<jlong> retainedSizes;
-    retainedSizes.reserve(size);
+    std::vector<jlong> resultingRetainedSizes;
+    resultingRetainedSizes.reserve(size);
     shallowSizes.reserve(size);
     for (jsize i = 0; i < size; i++) {
         jobject object = env->GetObjectArrayElement(objects, i);
         jlong tag;
         jvmti->GetTag(object, &tag);
-        retainedSizes.push_back(retained_sizes[tag]);
+        resultingRetainedSizes.push_back(retainedSizes[tag]);
         shallowSizes.push_back(info.sizes[tag]);
     }
-
-    jclass langObject = env->FindClass("java/lang/Object");
-    jobjectArray result = env->NewObjectArray(2, langObject, nullptr);
+    jobjectArray result = getObjectArrayOfSize(env, 2);
     env->SetObjectArrayElement(result, 0, toJavaArray(env, shallowSizes));
-    env->SetObjectArrayElement(result, 1, toJavaArray(env, retainedSizes));
-    if (!isOk(err)) {
-        handleError(jvmti, err, "Could not estimate retained size by classes");
-        return nullptr;
-    }
+    env->SetObjectArrayElement(result, 1, toJavaArray(env, resultingRetainedSizes));
+
     return result;
 }
 
-jvmtiError RetainedSizeViaDominatorTreeAction::cleanHeap() {
+jvmtiError RetainedSizesViaDominatorTreeAction::cleanHeap() {
     return removeAllTagsFromHeap(jvmti, nullptr);
 }
 
-RetainedSizeViaDominatorTreeAction::RetainedSizeViaDominatorTreeAction(JNIEnv *env, jvmtiEnv *jvmti, jobject object) : 
-	MemoryAgentAction(env, jvmti, object) {
+RetainedSizesViaDominatorTreeAction::RetainedSizesViaDominatorTreeAction(JNIEnv *env, jvmtiEnv *jvmti, jobject object) :
+    RetainedSizesAction(env, jvmti, object) {
 }
 
-jobjectArray RetainedSizeByClassViaDominatorTreeAction::executeOperation(jobject classRef, jlong objectsLimit) {
+jobjectArray RetainedSizesByClassViaDominatorTreeAction::executeOperation(jobject classRef, jlong objectsLimit) {
     jvmtiError err = jvmti->SetTag(classRef, CLASS_TAG);
-    if (err != JVMTI_ERROR_NONE) return nullptr;
+    if (!isOk(err)) return nullptr;
 
-    err = FollowReferences(0, nullptr, nullptr, collectObjects, nullptr, "collect objects of class");
-    if (err != JVMTI_ERROR_NONE || shouldStopExecution()) return nullptr;
+    progressManager.updateProgress(5, "Collecting objects of class");
+    err = FollowReferences(0, nullptr, nullptr, collectObjects, nullptr);
+    if (!isOk(err) || shouldStopExecution()) return nullptr;
 
     std::vector<jobject> objectsOfClass;
     err = getObjectsByTags(jvmti, std::vector<jlong>{OBJECT_OF_CLASS_TAG}, objectsOfClass);
-    if (err != JVMTI_ERROR_NONE || shouldStopExecution()) return nullptr;
+    if (!isOk(err) || shouldStopExecution()) return nullptr;
 
     // Constructing the master node
-    jclass langObject = env->FindClass("java/lang/Object");
-    jobjectArray objects = env->NewObjectArray(objectsOfClass.size(), langObject, nullptr);
+    jobjectArray objects = getObjectArrayOfSize(env, objectsOfClass.size());
     for (int i = 0; i < objectsOfClass.size(); i++) {
         env->SetObjectArrayElement(objects, i, objectsOfClass[i]);
     }
 
     err = jvmti->SetTag(classRef, 0);
-    if (err != JVMTI_ERROR_NONE) return nullptr;
+    if (!isOk(err)) return nullptr;
 
     SizesViaDominatorTreeHeapDumpInfo info;
-    err = info.initAndSetTagsForObjects(env, jvmti, objects);
-    if (err != JVMTI_ERROR_NONE) return nullptr;
-
-    // We set a tag for the input array to ignore it during traversal
-    err = jvmti->SetTag(objects, MOCK_REFERRER_TAG);
-    if (err != JVMTI_ERROR_NONE) return nullptr;
-
-    progressManager.updateProgress(10, "Traversing heap for the first time...");
-    logger::resetTimer();
-    err = FollowReferences(0, nullptr, nullptr, firstTraversal, &info, "first traversal");
-    logger::logPassedTime();
-    if (err != JVMTI_ERROR_NONE || shouldStopExecution()) return nullptr;
-
-    logger::resetTimer();
-    progressManager.updateProgress(60, "Traversing heap for the second time...");
-    err = FollowReferences(0, nullptr, objects, secondTraversal, &info, "second traversal");
-    logger::logPassedTime();
-    if (err != JVMTI_ERROR_NONE || shouldStopExecution()) return nullptr;
-
-    progressManager.updateProgress(80, "Calculating retained size...");
-    info.setUpNeighboursForMasterNode();
-    std::vector<jlong> retainedSizes = calculateRetainedSizes(info.graph, info.sizes);
+    std::vector<jlong> retainedSizes;
+    err = calculateRetainedSizes(objects, retainedSizes, info);
+    if (!isOk(err) || shouldStopExecution()) return nullptr;
 
     // Sort objects by their retained size
     std::sort(objectsOfClass.begin(), objectsOfClass.end(), [&](jobject a, jobject b) {
@@ -243,7 +246,14 @@ jobjectArray RetainedSizeByClassViaDominatorTreeAction::executeOperation(jobject
     });
 
     progressManager.updateProgress(95, "Extracting answer...");
-    size_t size = std::min((size_t)objectsLimit, objectsOfClass.size());
+    return constructResultObject(objectsOfClass, retainedSizes, info, objectsLimit);
+}
+
+jobjectArray RetainedSizesByClassViaDominatorTreeAction::constructResultObject(const std::vector<jobject> &objectsOfClass,
+                                                                               const std::vector<jlong> &retainedSizes,
+                                                                               const SizesViaDominatorTreeHeapDumpInfo &info,
+                                                                               jlong objectsLimit) {
+    size_t size = std::min(static_cast<size_t>(objectsLimit), objectsOfClass.size());
     std::vector<jlong> shallowSizes;
     std::vector<jlong> sortedRetainedSizes;
     sortedRetainedSizes.reserve(size);
@@ -256,25 +266,22 @@ jobjectArray RetainedSizeByClassViaDominatorTreeAction::executeOperation(jobject
         shallowSizes.push_back(info.sizes[tag]);
     }
 
-    jobjectArray result = env->NewObjectArray(3, langObject, nullptr);
-    objects = env->NewObjectArray(size, langObject, nullptr);
+    jobjectArray result = getObjectArrayOfSize(env, 3);
+    jobjectArray objects = getObjectArrayOfSize(env, size);
     for (int i = 0; i < size; i++) {
         env->SetObjectArrayElement(objects, i, objectsOfClass[i]);
     }
     env->SetObjectArrayElement(result, 0, objects);
     env->SetObjectArrayElement(result, 1, toJavaArray(env, shallowSizes));
     env->SetObjectArrayElement(result, 2, toJavaArray(env, sortedRetainedSizes));
-    if (!isOk(err)) {
-        handleError(jvmti, err, "Could not estimate retained size by classes");
-        return nullptr;
-    }
+
     return result;
 }
 
-jvmtiError RetainedSizeByClassViaDominatorTreeAction::cleanHeap() {
+jvmtiError RetainedSizesByClassViaDominatorTreeAction::cleanHeap() {
     return removeAllTagsFromHeap(jvmti, nullptr);
 }
 
-RetainedSizeByClassViaDominatorTreeAction::RetainedSizeByClassViaDominatorTreeAction(JNIEnv *env, jvmtiEnv *jvmti, jobject object) :
-        MemoryAgentAction(env, jvmti, object) {
+RetainedSizesByClassViaDominatorTreeAction::RetainedSizesByClassViaDominatorTreeAction(JNIEnv *env, jvmtiEnv *jvmti, jobject object) :
+    RetainedSizesAction(env, jvmti, object) {
 }