From 4f214d20abeb9f2929279e01483ce1a87e019bc8 Mon Sep 17 00:00:00 2001 From: Mark Vieira Date: Thu, 23 Jan 2020 15:20:17 -0800 Subject: [PATCH] Limit max concurrency of test cluster nodes to a function of max workers (#51338) (cherry picked from commit 9a0238c7166e70e467ca61c1353157979fd1598b) --- .../testclusters/RestTestRunnerTask.java | 25 +++ .../testclusters/TestClustersPlugin.java | 154 ++++++++++-------- .../testclusters/TestClustersRegistry.java | 4 +- .../testclusters/TestClustersThrottle.java | 6 + .../gradle/tool/Boilerplate.java | 15 ++ 5 files changed, 138 insertions(+), 66 deletions(-) create mode 100644 buildSrc/src/main/java/org/elasticsearch/gradle/testclusters/TestClustersThrottle.java diff --git a/buildSrc/src/main/java/org/elasticsearch/gradle/testclusters/RestTestRunnerTask.java b/buildSrc/src/main/java/org/elasticsearch/gradle/testclusters/RestTestRunnerTask.java index f4de829009d..5cd88ea01a3 100644 --- a/buildSrc/src/main/java/org/elasticsearch/gradle/testclusters/RestTestRunnerTask.java +++ b/buildSrc/src/main/java/org/elasticsearch/gradle/testclusters/RestTestRunnerTask.java @@ -1,12 +1,22 @@ package org.elasticsearch.gradle.testclusters; +import org.elasticsearch.gradle.tool.Boilerplate; +import org.gradle.api.provider.Provider; +import org.gradle.api.services.internal.BuildServiceRegistryInternal; import org.gradle.api.tasks.CacheableTask; +import org.gradle.api.tasks.Internal; import org.gradle.api.tasks.Nested; import org.gradle.api.tasks.testing.Test; +import org.gradle.internal.resources.ResourceLock; +import org.gradle.internal.resources.SharedResource; +import java.util.ArrayList; import java.util.Collection; import java.util.Collections; import java.util.HashSet; +import java.util.List; + +import static org.elasticsearch.gradle.testclusters.TestClustersPlugin.THROTTLE_SERVICE_NAME; /** * Customized version of Gradle {@link Test} task which tracks a collection of {@link ElasticsearchCluster} as a task input. We must do this @@ -47,4 +57,19 @@ public class RestTestRunnerTask extends Test implements TestClustersAware { return clusters; } + @Override + @Internal + public List getSharedResources() { + List locks = new ArrayList<>(super.getSharedResources()); + BuildServiceRegistryInternal serviceRegistry = getServices().get(BuildServiceRegistryInternal.class); + Provider throttleProvider = Boilerplate.getBuildService(serviceRegistry, THROTTLE_SERVICE_NAME); + SharedResource resource = serviceRegistry.forService(throttleProvider); + + int nodeCount = clusters.stream().mapToInt(cluster -> cluster.getNodes().size()).sum(); + if (nodeCount > 0) { + locks.add(resource.getResourceLock(Math.min(nodeCount, resource.getMaxUsages()))); + } + + return Collections.unmodifiableList(locks); + } } diff --git a/buildSrc/src/main/java/org/elasticsearch/gradle/testclusters/TestClustersPlugin.java b/buildSrc/src/main/java/org/elasticsearch/gradle/testclusters/TestClustersPlugin.java index e45119670d9..fc6f81f3347 100644 --- a/buildSrc/src/main/java/org/elasticsearch/gradle/testclusters/TestClustersPlugin.java +++ b/buildSrc/src/main/java/org/elasticsearch/gradle/testclusters/TestClustersPlugin.java @@ -21,6 +21,7 @@ package org.elasticsearch.gradle.testclusters; import org.elasticsearch.gradle.DistributionDownloadPlugin; import org.elasticsearch.gradle.ReaperPlugin; import org.elasticsearch.gradle.ReaperService; +import org.elasticsearch.gradle.tool.Boilerplate; import org.gradle.api.NamedDomainObjectContainer; import org.gradle.api.Plugin; import org.gradle.api.Project; @@ -30,53 +31,50 @@ import org.gradle.api.execution.TaskExecutionListener; import org.gradle.api.invocation.Gradle; import org.gradle.api.logging.Logger; import org.gradle.api.logging.Logging; +import org.gradle.api.provider.Provider; import org.gradle.api.tasks.TaskState; import java.io.File; public class TestClustersPlugin implements Plugin { - private static final String LIST_TASK_NAME = "listTestClusters"; public static final String EXTENSION_NAME = "testClusters"; - private static final String REGISTRY_EXTENSION_NAME = "testClustersRegistry"; + public static final String THROTTLE_SERVICE_NAME = "testClustersThrottle"; + private static final String LIST_TASK_NAME = "listTestClusters"; + private static final String REGISTRY_SERVICE_NAME = "testClustersRegistry"; private static final Logger logger = Logging.getLogger(TestClustersPlugin.class); - private ReaperService reaper; - @Override public void apply(Project project) { project.getPlugins().apply(DistributionDownloadPlugin.class); - project.getRootProject().getPluginManager().apply(ReaperPlugin.class); - reaper = project.getRootProject().getExtensions().getByType(ReaperService.class); + + ReaperService reaper = project.getRootProject().getExtensions().getByType(ReaperService.class); // enable the DSL to describe clusters - NamedDomainObjectContainer container = createTestClustersContainerExtension(project); + NamedDomainObjectContainer container = createTestClustersContainerExtension(project, reaper); // provide a task to be able to list defined clusters. createListClustersTask(project, container); - if (project.getRootProject().getExtensions().findByName(REGISTRY_EXTENSION_NAME) == null) { - TestClustersRegistry registry = project.getRootProject() - .getExtensions() - .create(REGISTRY_EXTENSION_NAME, TestClustersRegistry.class); + // register cluster registry as a global build service + project.getGradle().getSharedServices().registerIfAbsent(REGISTRY_SERVICE_NAME, TestClustersRegistry.class, spec -> {}); - // When we know what tasks will run, we claim the clusters of those task to differentiate between clusters - // that are defined in the build script and the ones that will actually be used in this invocation of gradle - // we use this information to determine when the last task that required the cluster executed so that we can - // terminate the cluster right away and free up resources. - configureClaimClustersHook(project.getGradle(), registry); + // register throttle so we only run at most max-workers/2 nodes concurrently + project.getGradle() + .getSharedServices() + .registerIfAbsent( + THROTTLE_SERVICE_NAME, + TestClustersThrottle.class, + spec -> spec.getMaxParallelUsages().set(project.getGradle().getStartParameter().getMaxWorkerCount() / 2) + ); - // Before each task, we determine if a cluster needs to be started for that task. - configureStartClustersHook(project.getGradle(), registry); - - // After each task we determine if there are clusters that are no longer needed. - configureStopClustersHook(project.getGradle(), registry); - } + // register cluster hooks + project.getRootProject().getPluginManager().apply(TestClustersHookPlugin.class); } - private NamedDomainObjectContainer createTestClustersContainerExtension(Project project) { + private NamedDomainObjectContainer createTestClustersContainerExtension(Project project, ReaperService reaper) { // Create an extensions that allows describing clusters NamedDomainObjectContainer container = project.container( ElasticsearchCluster.class, @@ -95,52 +93,78 @@ public class TestClustersPlugin implements Plugin { ); } - private static void configureClaimClustersHook(Gradle gradle, TestClustersRegistry registry) { - // 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. - gradle.getTaskGraph().whenReady(taskExecutionGraph -> { - taskExecutionGraph.getAllTasks() - .stream() - .filter(task -> task instanceof TestClustersAware) - .map(task -> (TestClustersAware) task) - .flatMap(task -> task.getClusters().stream()) - .forEach(registry::claimCluster); - }); - } - - private static void configureStartClustersHook(Gradle gradle, TestClustersRegistry registry) { - gradle.addListener(new TaskActionListener() { - @Override - public void beforeActions(Task task) { - if (task instanceof TestClustersAware == false) { - return; - } - // we only start the cluster before the actions, so we'll not start it if the task is up-to-date - TestClustersAware awareTask = (TestClustersAware) task; - awareTask.beforeStart(); - awareTask.getClusters().forEach(registry::maybeStartCluster); + static class TestClustersHookPlugin implements Plugin { + @Override + public void apply(Project project) { + if (project != project.getRootProject()) { + throw new IllegalStateException(this.getClass().getName() + " can only be applied to the root project."); } - @Override - public void afterActions(Task task) {} - }); - } + Provider registryProvider = Boilerplate.getBuildService( + project.getGradle().getSharedServices(), + REGISTRY_SERVICE_NAME + ); + TestClustersRegistry registry = registryProvider.get(); - private static void configureStopClustersHook(Gradle gradle, TestClustersRegistry registry) { - gradle.addListener(new TaskExecutionListener() { - @Override - public void afterExecute(Task task, TaskState state) { - if (task instanceof TestClustersAware == false) { - return; + // When we know what tasks will run, we claim the clusters of those task to differentiate between clusters + // that are defined in the build script and the ones that will actually be used in this invocation of gradle + // we use this information to determine when the last task that required the cluster executed so that we can + // terminate the cluster right away and free up resources. + configureClaimClustersHook(project.getGradle(), registry); + + // Before each task, we determine if a cluster needs to be started for that task. + configureStartClustersHook(project.getGradle(), registry); + + // After each task we determine if there are clusters that are no longer needed. + configureStopClustersHook(project.getGradle(), registry); + } + + private static void configureClaimClustersHook(Gradle gradle, TestClustersRegistry registry) { + // 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. + gradle.getTaskGraph().whenReady(taskExecutionGraph -> { + taskExecutionGraph.getAllTasks() + .stream() + .filter(task -> task instanceof TestClustersAware) + .map(task -> (TestClustersAware) task) + .flatMap(task -> task.getClusters().stream()) + .forEach(registry::claimCluster); + }); + } + + private static void configureStartClustersHook(Gradle gradle, TestClustersRegistry registry) { + gradle.addListener(new TaskActionListener() { + @Override + public void beforeActions(Task task) { + if (task instanceof TestClustersAware == false) { + return; + } + // we only start the cluster before the actions, so we'll not start it if the task is up-to-date + TestClustersAware awareTask = (TestClustersAware) task; + awareTask.beforeStart(); + awareTask.getClusters().forEach(registry::maybeStartCluster); } - // always unclaim the cluster, even if _this_ task is up-to-date, as others might not have been - // and caused the cluster to start. - ((TestClustersAware) task).getClusters().forEach(cluster -> registry.stopCluster(cluster, state.getFailure() != null)); - } - @Override - public void beforeExecute(Task task) {} - }); + @Override + public void afterActions(Task task) {} + }); + } + + private static void configureStopClustersHook(Gradle gradle, TestClustersRegistry registry) { + gradle.addListener(new TaskExecutionListener() { + @Override + public void afterExecute(Task task, TaskState state) { + if (task instanceof TestClustersAware == 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. + ((TestClustersAware) task).getClusters().forEach(cluster -> registry.stopCluster(cluster, state.getFailure() != null)); + } + + @Override + public void beforeExecute(Task task) {} + }); + } } - } diff --git a/buildSrc/src/main/java/org/elasticsearch/gradle/testclusters/TestClustersRegistry.java b/buildSrc/src/main/java/org/elasticsearch/gradle/testclusters/TestClustersRegistry.java index d78aecc8218..dff0a475eb4 100644 --- a/buildSrc/src/main/java/org/elasticsearch/gradle/testclusters/TestClustersRegistry.java +++ b/buildSrc/src/main/java/org/elasticsearch/gradle/testclusters/TestClustersRegistry.java @@ -2,13 +2,15 @@ package org.elasticsearch.gradle.testclusters; import org.gradle.api.logging.Logger; import org.gradle.api.logging.Logging; +import org.gradle.api.services.BuildService; +import org.gradle.api.services.BuildServiceParameters; import java.util.HashMap; import java.util.HashSet; import java.util.Map; import java.util.Set; -public class TestClustersRegistry { +public abstract class TestClustersRegistry implements BuildService { private static final Logger logger = Logging.getLogger(TestClustersRegistry.class); private static final String TESTCLUSTERS_INSPECT_FAILURE = "testclusters.inspect.failure"; private final Boolean allowClusterToSurvive = Boolean.valueOf(System.getProperty(TESTCLUSTERS_INSPECT_FAILURE, "false")); diff --git a/buildSrc/src/main/java/org/elasticsearch/gradle/testclusters/TestClustersThrottle.java b/buildSrc/src/main/java/org/elasticsearch/gradle/testclusters/TestClustersThrottle.java new file mode 100644 index 00000000000..4c04f939560 --- /dev/null +++ b/buildSrc/src/main/java/org/elasticsearch/gradle/testclusters/TestClustersThrottle.java @@ -0,0 +1,6 @@ +package org.elasticsearch.gradle.testclusters; + +import org.gradle.api.services.BuildService; +import org.gradle.api.services.BuildServiceParameters; + +public abstract class TestClustersThrottle implements BuildService {} diff --git a/buildSrc/src/main/java/org/elasticsearch/gradle/tool/Boilerplate.java b/buildSrc/src/main/java/org/elasticsearch/gradle/tool/Boilerplate.java index 3652944e164..e10b4099b7d 100644 --- a/buildSrc/src/main/java/org/elasticsearch/gradle/tool/Boilerplate.java +++ b/buildSrc/src/main/java/org/elasticsearch/gradle/tool/Boilerplate.java @@ -19,12 +19,17 @@ package org.elasticsearch.gradle.tool; import org.gradle.api.Action; +import org.gradle.api.GradleException; import org.gradle.api.NamedDomainObjectContainer; import org.gradle.api.PolymorphicDomainObjectContainer; import org.gradle.api.Project; import org.gradle.api.Task; import org.gradle.api.UnknownTaskException; import org.gradle.api.plugins.JavaPluginConvention; +import org.gradle.api.provider.Provider; +import org.gradle.api.services.BuildService; +import org.gradle.api.services.BuildServiceRegistration; +import org.gradle.api.services.BuildServiceRegistry; import org.gradle.api.tasks.SourceSetContainer; import org.gradle.api.tasks.TaskContainer; import org.gradle.api.tasks.TaskProvider; @@ -102,4 +107,14 @@ public abstract class Boilerplate { return task; } + + @SuppressWarnings("unchecked") + public static > Provider getBuildService(BuildServiceRegistry registry, String name) { + BuildServiceRegistration registration = registry.getRegistrations().findByName(name); + if (registration == null) { + throw new GradleException("Unable to find build service with name '" + name + "'."); + } + + return (Provider) registration.getService(); + } }