From e37c35c6ee1291805b593c3de5bfa8b94aad4721 Mon Sep 17 00:00:00 2001
From: "Anna.Kozlova" <anna.kozlova@jetbrains.com>
Date: Mon, 23 Jul 2018 11:10:38 +0200
Subject: [PATCH] test discovery: minimize number of forks to speadup tests

each fork leads to new application creation ~8 sec
---
 .../intellij/execution/junit/TestObject.java  | 12 ++---
 ...nitTestDiscoveryConfigurationProducer.java | 46 +++++++++++++++++--
 2 files changed, 48 insertions(+), 10 deletions(-)

diff --git a/plugins/junit/src/com/intellij/execution/junit/TestObject.java b/plugins/junit/src/com/intellij/execution/junit/TestObject.java
index c225160e076d..89061d8a4738 100644
--- a/plugins/junit/src/com/intellij/execution/junit/TestObject.java
+++ b/plugins/junit/src/com/intellij/execution/junit/TestObject.java
@@ -24,7 +24,6 @@ import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.module.Module;
 import com.intellij.openapi.module.ModuleUtilCore;
 import com.intellij.openapi.project.DumbService;
-import com.intellij.openapi.project.IndexNotReadyException;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.roots.OrderRootType;
 import com.intellij.openapi.roots.ProjectFileIndex;
@@ -124,12 +123,7 @@ public abstract class TestObject extends JavaTestFrameworkRunnableState<JUnitCon
         if (perModule != null && psiElement != null) {
           final Module module = ModuleUtilCore.findModuleForPsiElement(psiElement);
           if (module != null) {
-            List<String> list = perModule.get(module);
-            if (list == null) {
-              list = new ArrayList<>();
-              perModule.put(module, list);
-            }
-            list.add(name);
+            fillForkModule(perModule, module, name);
           }
         }
         else {
@@ -160,6 +154,10 @@ public abstract class TestObject extends JavaTestFrameworkRunnableState<JUnitCon
     }
   }
 
+  protected void fillForkModule(Map<Module, List<String>> perModule, Module module, String name) {
+    perModule.computeIfAbsent(module, elemList -> new ArrayList<>()).add(name);
+  }
+
   public Module[] getModulesToCompile() {
     final SourceScope sourceScope = getSourceScope();
     return sourceScope != null ? sourceScope.getModulesToCompile() : Module.EMPTY_ARRAY;
diff --git a/plugins/junit/src/com/intellij/execution/junit/testDiscovery/JUnitTestDiscoveryConfigurationProducer.java b/plugins/junit/src/com/intellij/execution/junit/testDiscovery/JUnitTestDiscoveryConfigurationProducer.java
index 750f3ca1cbb4..fed204186010 100644
--- a/plugins/junit/src/com/intellij/execution/junit/testDiscovery/JUnitTestDiscoveryConfigurationProducer.java
+++ b/plugins/junit/src/com/intellij/execution/junit/testDiscovery/JUnitTestDiscoveryConfigurationProducer.java
@@ -13,15 +13,14 @@ import com.intellij.execution.junit.TestsPattern;
 import com.intellij.execution.runners.ExecutionEnvironment;
 import com.intellij.execution.testDiscovery.TestDiscoveryConfigurationProducer;
 import com.intellij.openapi.module.Module;
+import com.intellij.openapi.module.ModuleUtilCore;
 import com.intellij.openapi.util.Pair;
 import com.intellij.psi.PsiClass;
 import com.intellij.psi.PsiMethod;
 import com.intellij.rt.execution.junit.JUnitStarter;
 import org.jetbrains.annotations.NotNull;
 
-import java.util.Arrays;
-import java.util.Iterator;
-import java.util.LinkedHashSet;
+import java.util.*;
 import java.util.stream.Collectors;
 
 public class JUnitTestDiscoveryConfigurationProducer extends TestDiscoveryConfigurationProducer {
@@ -63,16 +62,57 @@ public class JUnitTestDiscoveryConfigurationProducer extends TestDiscoveryConfig
             })
             .collect(Collectors.toCollection(LinkedHashSet::new)));
     data.TEST_OBJECT = JUnitConfiguration.TEST_PATTERN;
+    Map<Module, Module> toRoot = splitModulesIntoChunks(testMethods, module);
     return new TestsPattern((JUnitConfiguration)configuration, environment) {
       @Override
       protected boolean forkPerModule() {
         return module == null;
       }
 
+      @Override
+      protected void fillForkModule(Map<Module, List<String>> perModule, Module module, String name) {
+        super.fillForkModule(perModule, toRoot.get(module), name);
+      }
+
       @Override
       protected String getRunner() {
         return JUnitStarter.JUNIT4_PARAMETER;
       }
     };
   }
+
+  private static Map<Module, Module> splitModulesIntoChunks(@NotNull Location<PsiMethod>[] testMethods, Module module) {
+    Map<Module, Module> toRoot = new HashMap<>();
+    if (module == null) {
+      List<Module> usedModules = Arrays.stream(testMethods).map(Location::getModule).collect(Collectors.toList());
+      while (!usedModules.isEmpty()) {
+        Map<Module, List<Module>> allDeps = new HashMap<>();
+        for (Module usedModule : usedModules) {
+          List<Module> rootModules = ModuleUtilCore.getAllDependentModules(usedModule);
+          for (Module rootModule : rootModules) {
+            allDeps.computeIfAbsent(rootModule, __ -> new ArrayList<>()).add(usedModule);
+          }
+          allDeps.computeIfAbsent(usedModule, __ -> new ArrayList<>()).add(usedModule);
+        }
+
+        
+        Optional<Map.Entry<Module, List<Module>>> maxDependency =
+          allDeps.entrySet().stream().max(Comparator.comparingInt(e -> e.getValue().size()));
+
+        if (maxDependency.isPresent()) {
+          Map.Entry<Module, List<Module>> entry = maxDependency.get();
+          Module rootModule = entry.getKey();
+          List<Module> srcModules = entry.getValue();
+          for (Module srcModule : srcModules) {
+            toRoot.put(srcModule, rootModule);
+          }
+          usedModules.removeAll(srcModules);
+        }
+        else {
+          break;
+        }
+      }
+    }
+    return toRoot;
+  }
 }