Restrict which tasks can use testclusters

buildSrc/src/main/groovy/org/elasticsearch/gradle/test/RestIntegTestTask.groovy

@@ -47,8 +47,6 @@ class RestIntegTestTask extends DefaultTask {
 
     protected Test runner
 
-    protected Task clusterInit
-
     /** Info about nodes in the integ test cluster. Note this is *not* available until runtime. */
     List<NodeInfo> nodes
 
@@ -59,8 +57,6 @@ class RestIntegTestTask extends DefaultTask {
     RestIntegTestTask() {
         runner = project.tasks.create("${name}Runner", RestTestRunnerTask.class)
         super.dependsOn(runner)
-        clusterInit = project.tasks.create(name: "${name}Cluster#init", dependsOn: project.testClasses)
-        runner.dependsOn(clusterInit)
         boolean usesTestclusters = project.plugins.hasPlugin(TestClustersPlugin.class)
         if (usesTestclusters == false) {
             clusterConfig = project.extensions.create("${name}Cluster", ClusterConfiguration.class, project)
@@ -73,8 +69,6 @@ class RestIntegTestTask extends DefaultTask {
             runner.useCluster project.testClusters."$name"
         }
 
-        // override/add more for rest tests
-        runner.maxParallelForks = 1
         runner.include('**/*IT.class')
         runner.systemProperty('tests.rest.load_packaged', 'false')
 
@@ -135,7 +129,6 @@ class RestIntegTestTask extends DefaultTask {
         project.gradle.projectsEvaluated {
             if (enabled == false) {
                 runner.enabled = false
-                clusterInit.enabled = false
                 return // no need to add cluster formation tasks if the task won't run!
             }
             if (usesTestclusters == false) {
@@ -186,11 +179,6 @@ class RestIntegTestTask extends DefaultTask {
         }
     }
 
-    @Override
-    public Task mustRunAfter(Object... tasks) {
-        clusterInit.mustRunAfter(tasks)
-    }
-
     public void runner(Closure configure) {
         project.tasks.getByName("${name}Runner").configure(configure)
     }

buildSrc/src/main/java/org/elasticsearch/gradle/test/RestTestRunnerTask.java

@@ -1,12 +1,13 @@
 package org.elasticsearch.gradle.test;
 
 import org.elasticsearch.gradle.testclusters.ElasticsearchCluster;
+import org.elasticsearch.gradle.testclusters.TestClustersTask;
 import org.gradle.api.tasks.CacheableTask;
 import org.gradle.api.tasks.Nested;
 import org.gradle.api.tasks.testing.Test;
 
-import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashSet;
 
 import static org.elasticsearch.gradle.testclusters.TestDistribution.INTEG_TEST;
 
@@ -16,21 +17,28 @@
  * {@link Nested} inputs.
  */
 @CacheableTask
-public class RestTestRunnerTask extends Test {
+public class RestTestRunnerTask extends Test implements TestClustersTask {
 
-    private Collection<ElasticsearchCluster> clusters = new ArrayList<>();
+    private Collection<ElasticsearchCluster> clusters = new HashSet<>();
 
     public RestTestRunnerTask() {
         super();
         this.getOutputs().doNotCacheIf("Build cache is only enabled for tests against clusters using the 'integ-test' distribution",
             task -> clusters.stream().flatMap(c -> c.getNodes().stream()).anyMatch(n -> n.getTestDistribution() != INTEG_TEST));
     }
 
+    @Override
+    public int getMaxParallelForks() {
+        return 1;
+    }
+
     @Nested
+    @Override
     public Collection<ElasticsearchCluster> getClusters() {
         return clusters;
     }
 
+    @Override
     public void testCluster(ElasticsearchCluster cluster) {
         this.clusters.add(cluster);
     }

buildSrc/src/main/java/org/elasticsearch/gradle/testclusters/DefaultTestClustersTask.java

@@ -0,0 +1,22 @@
+package org.elasticsearch.gradle.testclusters;
+
+import org.gradle.api.DefaultTask;
+
+import java.util.Collection;
+import java.util.HashSet;
+
+public class DefaultTestClustersTask extends DefaultTask implements TestClustersTask {
+
+    private Collection<ElasticsearchCluster> clusters = new HashSet<>();
+
+    @Override
+    public Collection<ElasticsearchCluster> getClusters() {
+        return clusters;
+    }
+
+    @Override
+    public void testCluster(ElasticsearchCluster cluster) {
+        this.clusters.add(cluster);
+    }
+
+}

buildSrc/src/main/java/org/elasticsearch/gradle/testclusters/TestClustersPlugin.java

@@ -23,7 +23,6 @@
 import org.elasticsearch.gradle.ElasticsearchDistribution;
 import org.elasticsearch.gradle.ReaperPlugin;
 import org.elasticsearch.gradle.ReaperService;
-import org.elasticsearch.gradle.test.RestTestRunnerTask;
 import org.gradle.api.NamedDomainObjectContainer;
 import org.gradle.api.Plugin;
 import org.gradle.api.Project;
@@ -36,8 +35,7 @@
 import org.gradle.api.tasks.TaskState;
 
 import java.io.File;
-import java.util.ArrayList;
-import java.util.Collections;
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -53,7 +51,6 @@ public class TestClustersPlugin implements Plugin<Project> {
     private static final Logger logger =  Logging.getLogger(TestClustersPlugin.class);
     private static final String TESTCLUSTERS_INSPECT_FAILURE = "testclusters.inspect.failure";
 
-    private final Map<Task, List<ElasticsearchCluster>> usedClusters = new HashMap<>();
     private final Map<ElasticsearchCluster, Integer> claimsInventory = new HashMap<>();
     private final Set<ElasticsearchCluster> runningClusters = new HashSet<>();
     private final Boolean allowClusterToSurvive = Boolean.valueOf(System.getProperty(TESTCLUSTERS_INSPECT_FAILURE, "false"));
@@ -123,7 +120,7 @@ private void createListClustersTask(Project project, NamedDomainObjectContainer<
     private void createUseClusterTaskExtension(Project project, NamedDomainObjectContainer<ElasticsearchCluster> container) {
         // register an extension for all current and future tasks, so that any task can declare that it wants to use a
         // specific cluster.
-        project.getTasks().all((Task task) ->
+        project.getTasks().withType(TestClustersTask.class, (Task task) ->
             task.getExtensions().findByType(ExtraPropertiesExtension.class)
                 .set(
                     "useCluster",
@@ -140,13 +137,11 @@ public void doCall(ElasticsearchCluster cluster) {
                                 throw new AssertionError("Expected " + thisObject + " to be an instance of " +
                                     "Task, but got: " + thisObject.getClass());
                             }
-                            usedClusters.computeIfAbsent(task, k -> new ArrayList<>()).add(cluster);
                             for (ElasticsearchNode node : cluster.getNodes()) {
                                 ((Task) thisObject).dependsOn(node.getDistribution().getExtracted());
                             }
-                            if (thisObject instanceof RestTestRunnerTask) {
-                                ((RestTestRunnerTask) thisObject).testCluster(cluster);
-                            }
+
+                            ((TestClustersTask) thisObject).testCluster(cluster);
                         }
                     })
         );
@@ -156,17 +151,15 @@ private void configureClaimClustersHook(Project project) {
         // Once we know all the tasks that need to execute, we claim all the clusters that belong to those and count the
         // claims so we'll know when it's safe to stop them.
         project.getGradle().getTaskGraph().whenReady(taskExecutionGraph -> {
-            Set<String> forExecution = taskExecutionGraph.getAllTasks().stream()
-                .map(Task::getPath)
-                .collect(Collectors.toSet());
+            taskExecutionGraph.getAllTasks().stream()
+                .filter(task -> task instanceof TestClustersTask)
+                .map(task -> (TestClustersTask) task)
+                .flatMap(task -> task.getClusters().stream())
+                .forEach(cluster -> {
+                    cluster.freeze();
+                    claimsInventory.put(cluster, claimsInventory.getOrDefault(cluster, 0) + 1);
+                });
 
-            usedClusters.forEach((task, listOfClusters) ->
-                listOfClusters.forEach(elasticsearchCluster -> {
-                    if (forExecution.contains(task.getPath())) {
-                        elasticsearchCluster.freeze();
-                        claimsInventory.put(elasticsearchCluster, claimsInventory.getOrDefault(elasticsearchCluster, 0) + 1);
-                    }
-                }));
             if (claimsInventory.isEmpty() == false) {
                 logger.info("Claims inventory: {}", claimsInventory);
             }
@@ -178,11 +171,11 @@ private void configureStartClustersHook(Project project) {
             new TaskActionListener() {
                 @Override
                 public void beforeActions(Task task) {
+                    if (task instanceof TestClustersTask == false) {
+                        return;
+                    }
                     // we only start the cluster before the actions, so we'll not start it if the task is up-to-date
-                    List<ElasticsearchCluster> neededButNotRunning = usedClusters.getOrDefault(
-                        task,
-                        Collections.emptyList()
-                    )
+                    List<ElasticsearchCluster> neededButNotRunning = ((TestClustersTask) task).getClusters()
                         .stream()
                         .filter(cluster -> runningClusters.contains(cluster) == false)
                         .collect(Collectors.toList());
@@ -204,12 +197,12 @@ private void configureStopClustersHook(Project project) {
             new TaskExecutionListener() {
                 @Override
                 public void afterExecute(Task task, TaskState state) {
+                    if (task instanceof TestClustersTask == false) {
+                        return;
+                    }
                     // always unclaim the cluster, even if _this_ task is up-to-date, as others might not have been
                     // and caused the cluster to start.
-                    List<ElasticsearchCluster> clustersUsedByTask = usedClusters.getOrDefault(
-                        task,
-                        Collections.emptyList()
-                    );
+                    Collection<ElasticsearchCluster> clustersUsedByTask = ((TestClustersTask) task).getClusters();
                     if (clustersUsedByTask.isEmpty()) {
                         return;
                     }
@@ -220,7 +213,7 @@ public void afterExecute(Task task, TaskState state) {
                         // executed at all, so we will never be called again to un-claim and terminate it.
                         clustersUsedByTask.forEach(cluster -> stopCluster(cluster, true));
                         permitsToRelease = clustersUsedByTask.stream()
-                            .map(cluster -> cluster.getNumberOfNodes())
+                            .map(ElasticsearchCluster::getNumberOfNodes)
                             .reduce(Integer::sum).get();
                     } else {
                         clustersUsedByTask.forEach(

buildSrc/src/main/java/org/elasticsearch/gradle/testclusters/TestClustersTask.java

@@ -0,0 +1,14 @@
+package org.elasticsearch.gradle.testclusters;
+
+import org.gradle.api.Task;
+import org.gradle.api.tasks.Nested;
+
+import java.util.Collection;
+
+public interface TestClustersTask extends Task {
+
+    @Nested
+    Collection<ElasticsearchCluster> getClusters();
+
+    void testCluster(ElasticsearchCluster cluster);
+}

qa/multi-cluster-search/build.gradle

@@ -39,8 +39,8 @@ testClusters.'remote-cluster' {
 }
 
 task mixedClusterTest(type: RestIntegTestTask) {
-  useCluster testClusters.'remote-cluster'
   runner {
+    useCluster testClusters.'remote-cluster'
     dependsOn 'remote-cluster'
     systemProperty 'tests.rest.suite', 'multi_cluster'
   }

qa/wildfly/build.gradle

@@ -1,6 +1,8 @@
 import org.elasticsearch.gradle.LoggedExec
 import org.elasticsearch.gradle.VersionProperties
 import org.apache.tools.ant.taskdefs.condition.Os
+import org.elasticsearch.gradle.test.RestTestRunnerTask
+import org.elasticsearch.gradle.testclusters.DefaultTestClustersTask
 
 import java.nio.charset.StandardCharsets
 import java.nio.file.Files
@@ -87,7 +89,7 @@ task deploy(type: Copy) {
     into "${wildflyInstall}/standalone/deployments"
 }
 
-task writeElasticsearchProperties {
+task writeElasticsearchProperties(type: DefaultTestClustersTask) {
     onlyIf { !Os.isFamily(Os.FAMILY_WINDOWS) }
     useCluster testClusters.integTest
     dependsOn deploy

x-pack/plugin/ccr/qa/downgrade-to-basic-license/build.gradle

@@ -41,8 +41,8 @@ task writeJavaPolicy {
 
 task "follow-cluster"(type: RestIntegTestTask) {
     dependsOn 'writeJavaPolicy', "leader-cluster"
-    useCluster testClusters."leader-cluster"
     runner {
+        useCluster testClusters."leader-cluster"
         systemProperty 'java.security.policy', "file://${policyFile}"
         systemProperty 'tests.target_cluster', 'follow'
         nonInputProperties.systemProperty 'tests.leader_host', "${-> testClusters."leader-cluster".getAllHttpSocketURI().get(0)}"

x-pack/plugin/ccr/qa/multi-cluster/build.gradle

@@ -27,8 +27,8 @@ testClusters."leader-cluster" {
 
 task "middle-cluster"(type: RestIntegTestTask) {
     dependsOn "leader-cluster"
-    useCluster testClusters."leader-cluster"
     runner {
+        useCluster testClusters."leader-cluster"
         systemProperty 'tests.target_cluster', 'middle'
         nonInputProperties.systemProperty 'tests.leader_host',
                 "${-> testClusters."leader-cluster".getAllHttpSocketURI().get(0)}"
@@ -43,9 +43,9 @@ testClusters."middle-cluster" {
 
 task 'follow-cluster'(type: RestIntegTestTask) {
     dependsOn "leader-cluster", "middle-cluster"
-    useCluster testClusters."leader-cluster"
-    useCluster testClusters."middle-cluster"
     runner {
+        useCluster testClusters."leader-cluster"
+        useCluster testClusters."middle-cluster"
         systemProperty 'tests.target_cluster', 'follow'
         nonInputProperties.systemProperty 'tests.leader_host',
                 "${-> testClusters."leader-cluster".getAllHttpSocketURI().get(0)}"

x-pack/plugin/ccr/qa/non-compliant-license/build.gradle

@@ -24,8 +24,8 @@ testClusters.'leader-cluster' {
 
 task 'follow-cluster'(type: RestIntegTestTask) {
     dependsOn 'leader-cluster'
-    useCluster testClusters.'leader-cluster'
     runner {
+        useCluster testClusters.'leader-cluster'
         systemProperty 'tests.target_cluster', 'follow'
         nonInputProperties.systemProperty 'tests.leader_host',
                 { "${testClusters.'follow-cluster'.getAllHttpSocketURI().get(0)}" }

x-pack/plugin/ccr/qa/restart/build.gradle

@@ -1,4 +1,5 @@
 import org.elasticsearch.gradle.test.RestIntegTestTask
+import org.elasticsearch.gradle.test.RestTestRunnerTask
 
 apply plugin: 'elasticsearch.testclusters'
 apply plugin: 'elasticsearch.standalone-test'
@@ -20,8 +21,8 @@ testClusters.'leader-cluster' {
 
 task 'follow-cluster'(type: RestIntegTestTask) {
     dependsOn 'leader-cluster'
-    useCluster testClusters.'leader-cluster'
     runner {
+        useCluster testClusters.'leader-cluster'
         systemProperty 'tests.target_cluster', 'follow'
         nonInputProperties.systemProperty 'tests.leader_host',
                 "${-> testClusters.'leader-cluster'.getAllHttpSocketURI().get(0)}"
@@ -36,20 +37,18 @@ testClusters.'follow-cluster' {
     nameCustomization = { 'follow'  }
 }
 
-task followClusterRestartTest(type: Test) {
+task followClusterRestartTest(type: RestTestRunnerTask) {
     dependsOn tasks.'follow-cluster'
     useCluster testClusters.'leader-cluster'
     useCluster testClusters.'follow-cluster'
 
-    maxParallelForks = 1
     systemProperty 'tests.rest.load_packaged', 'false'
     systemProperty 'tests.target_cluster', 'follow-restart'
     doFirst {
         testClusters.'follow-cluster'.restart()
         nonInputProperties.systemProperty 'tests.leader_host', "${-> testClusters.'leader-cluster'.getAllHttpSocketURI().get(0)}"
         nonInputProperties.systemProperty 'tests.rest.cluster', "${-> testClusters.'follow-cluster'.getAllHttpSocketURI().join(",")}"
     }
-    outputs.doNotCacheIf "Caching of REST tests not implemented yet", { false }
 }
 
 check.dependsOn followClusterRestartTest

x-pack/plugin/ccr/qa/security/build.gradle

@@ -31,8 +31,8 @@ testClusters.'leader-cluster' {
 
 task 'follow-cluster'(type: RestIntegTestTask) {
     dependsOn 'leader-cluster'
-    useCluster testClusters.'leader-cluster'
     runner {
+        useCluster testClusters.'leader-cluster'
         systemProperty 'tests.target_cluster', 'follow'
         nonInputProperties.systemProperty 'tests.leader_host', "${-> testClusters.'leader-cluster'.getAllHttpSocketURI().get(0)}"
     }

x-pack/plugin/ilm/qa/multi-cluster/build.gradle

@@ -35,8 +35,8 @@ testClusters.'leader-cluster' {
 
 task 'follow-cluster'(type: RestIntegTestTask) {
     dependsOn 'leader-cluster'
-    useCluster testClusters.'leader-cluster'
     runner {
+        useCluster testClusters.'leader-cluster'
         systemProperty 'tests.target_cluster', 'follow'
         nonInputProperties.systemProperty 'tests.leader_host',
                 "${-> testClusters."leader-cluster".getAllHttpSocketURI().get(0)}"

x-pack/plugin/security/qa/basic-enable-security/build.gradle

@@ -1,3 +1,5 @@
+import org.elasticsearch.gradle.test.RestTestRunnerTask
+
 apply plugin: 'elasticsearch.testclusters'
 apply plugin: 'elasticsearch.standalone-rest-test'
 apply plugin: 'elasticsearch.rest-test'
@@ -24,13 +26,11 @@ testClusters.integTest {
     setting 'xpack.security.enabled', 'false'
 }
 
-task integTestSecurity(type: Test) {
+task integTestSecurity(type: RestTestRunnerTask) {
     description = "Run tests against a cluster that has security"
     useCluster testClusters.integTest
     dependsOn integTest
     systemProperty 'tests.has_security', 'true'
-    maxParallelForks = 1
-    outputs.cacheIf "Caching of REST tests not implemented yet", { false }
 
     doFirst {
         testClusters.integTest {