Explicitly declare usage of TestClusterThrottle

elastic · Nov 16, 2023 · bea07c4 · bea07c4
1 parent 69e13e5
commit bea07c4
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 7 deletions.
diff --git a/...ools/src/main/java/org/elasticsearch/gradle/testclusters/StandaloneRestIntegTestTask.java b/...ools/src/main/java/org/elasticsearch/gradle/testclusters/StandaloneRestIntegTestTask.java
@@ -70,13 +70,9 @@ public Collection<ElasticsearchCluster> getClusters() {
     @Internal
     public List<ResourceLock> getSharedResources() {
         List<ResourceLock> locks = new ArrayList<>(super.getSharedResources());
-        System.out.println("locks.size() = " + getName() + " -- " + locks.size());
         BuildServiceRegistryInternal serviceRegistry = getServices().get(BuildServiceRegistryInternal.class);
         BuildServiceProvider<?, ?> serviceProvider = serviceRegistry.consume(THROTTLE_SERVICE_NAME, TestClustersThrottle.class);
-        System.out.println("locks.size() 2 = " + getName() + " -- " + super.getSharedResources().size());
-
         SharedResource resource = serviceRegistry.forService(serviceProvider);
-        System.out.println("resource.getMaxUsages() = " + getName() + " -- " + resource.getMaxUsages());
         int nodeCount = clusters.stream().mapToInt(cluster -> cluster.getNodes().size()).sum();
         if (nodeCount > 0) {
             for (int i = 0; i < Math.min(nodeCount, resource.getMaxUsages()); i++) {

diff --git a/build-tools/src/main/java/org/elasticsearch/gradle/testclusters/TestClustersPlugin.java b/build-tools/src/main/java/org/elasticsearch/gradle/testclusters/TestClustersPlugin.java
@@ -109,15 +109,15 @@ public void apply(Project project) {
         project.getGradle().getSharedServices().registerIfAbsent(REGISTRY_SERVICE_NAME, TestClustersRegistry.class, noop());
 
         // register throttle so we only run at most max-workers/2 nodes concurrently
-        project.getGradle()
+        Provider<TestClustersThrottle> testClustersThrottleProvider = project.getGradle()
             .getSharedServices()
             .registerIfAbsent(
                 THROTTLE_SERVICE_NAME,
                 TestClustersThrottle.class,
                 spec -> spec.getMaxParallelUsages().set(Math.max(1, project.getGradle().getStartParameter().getMaxWorkerCount() / 2))
             );
 
-        // register cluster hooks
+        project.getTasks().withType(TestClustersAware.class).configureEach(task -> { task.usesService(testClustersThrottleProvider); });
         project.getRootProject().getPluginManager().apply(TestClustersHookPlugin.class);
     }
 

diff --git a/build-tools/src/main/java/org/elasticsearch/gradle/testclusters/TestClustersRegistry.java b/build-tools/src/main/java/org/elasticsearch/gradle/testclusters/TestClustersRegistry.java
@@ -34,7 +34,7 @@ public void claimCluster(ElasticsearchCluster cluster) {
     }
 
     public boolean isReusedCluster(ElasticsearchCluster cluster) {
-        return claimCount.get(cluster) > 1;
+        return claimCount.getOrDefault(cluster, 0) > 0;
     }
 
     public void maybeStartCluster(ElasticsearchCluster cluster) {