Skip to content

Commit

Permalink
Make TestClusterPlugin configuration cache compatible
Browse files Browse the repository at this point in the history
  • Loading branch information
@breskeby
breskeby committed Oct 18, 2023
1 parent e3cb876 commit 17517cc
Show file tree
Hide file tree
Showing 2 changed files with 77 additions and 20 deletions.
89 changes: 69 additions & 20 deletions build-tools/src/main/java/org/elasticsearch/gradle/testclusters/TestClustersPlugin.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
import org.gradle.api.Plugin;
import org.gradle.api.Project;
import org.gradle.api.Task;
import org.gradle.api.execution.TaskActionListener;
import org.gradle.api.execution.TaskExecutionListener;
import org.gradle.api.file.ArchiveOperations;
import org.gradle.api.file.FileSystemOperations;
Expand All @@ -26,11 +25,20 @@
import org.gradle.api.logging.Logging;
import org.gradle.api.provider.Provider;
import org.gradle.api.provider.ProviderFactory;
import org.gradle.api.services.BuildService;
import org.gradle.api.services.BuildServiceParameters;
import org.gradle.api.tasks.TaskState;
import org.gradle.build.event.BuildEventsListenerRegistry;
import org.gradle.internal.jvm.Jvm;
import org.gradle.process.ExecOperations;
import org.gradle.tooling.events.FinishEvent;
import org.gradle.tooling.events.OperationCompletionListener;
import org.gradle.tooling.events.task.TaskFailureResult;
import org.gradle.tooling.events.task.TaskFinishEvent;

import java.io.File;
import java.util.HashMap;
import java.util.Map;
import java.util.function.Function;

import javax.inject.Inject;
Expand Down Expand Up @@ -153,17 +161,27 @@ private void createListClustersTask(Project project, NamedDomainObjectContainer<

}

static class TestClustersHookPlugin implements Plugin<Project> {
@Override
static abstract class TestClustersHookPlugin implements Plugin<Project> {
@Inject
public abstract BuildEventsListenerRegistry getEventsListenerRegistry();

@Inject
public TestClustersHookPlugin() {
}
public void apply(Project project) {
if (project != project.getRootProject()) {
throw new IllegalStateException(this.getClass().getName() + " can only be applied to the root project.");
}

Provider<TestClustersRegistry> registryProvider = GradleUtils.getBuildService(
project.getGradle().getSharedServices(),
REGISTRY_SERVICE_NAME
);

Provider<TaskEventsService> testClusterTasksService =
project.getGradle().getSharedServices().registerIfAbsent(
"testClusterTasksService", TaskEventsService.class, spec -> {
});

TestClustersRegistry registry = registryProvider.get();

// When we know what tasks will run, we claim the clusters of those task to differentiate between clusters
Expand All @@ -173,10 +191,12 @@ public void apply(Project project) {
configureClaimClustersHook(project.getGradle(), registry);

// Before each task, we determine if a cluster needs to be started for that task.
configureStartClustersHook(project.getGradle(), registry);
configureStartClustersHook(project.getGradle(), registry, testClusterTasksService);

getEventsListenerRegistry().onTaskCompletion(testClusterTasksService);

// After each task we determine if there are clusters that are no longer needed.
configureStopClustersHook(project.getGradle(), registry);
//configureStopClustersHook(project.getGradle(), registry, testClusterTasksService;
}

private static void configureClaimClustersHook(Gradle gradle, TestClustersRegistry registry) {
Expand All @@ -192,25 +212,26 @@ private static void configureClaimClustersHook(Gradle gradle, TestClustersRegist
});
}

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);
}
private void configureStartClustersHook(Gradle gradle, TestClustersRegistry registry, Provider<TaskEventsService> testClusterTasksService) {
testClusterTasksService.get().registry(registry);

@Override
public void afterActions(Task task) {}
gradle.getTaskGraph().whenReady(taskExecutionGraph -> {
taskExecutionGraph.getAllTasks()
.stream()
.filter(task -> task instanceof TestClustersAware)
.map(task -> (TestClustersAware) task)
.forEach(awareTask -> {
testClusterTasksService.get().register(awareTask.getPath(), awareTask);
awareTask.doFirst(task -> {
awareTask.beforeStart();
awareTask.getClusters().forEach(registry::maybeStartCluster);
});
});
});
}

private static void configureStopClustersHook(Gradle gradle, TestClustersRegistry registry) {

gradle.addListener(new TaskExecutionListener() {
@Override
public void afterExecute(Task task, TaskState state) {
Expand All @@ -227,4 +248,32 @@ public void beforeExecute(Task task) {}
});
}
}

static public abstract class TaskEventsService implements BuildService<BuildServiceParameters.None>, OperationCompletionListener {

Map<String,TestClustersAware> tasksMap = new HashMap<>();
private TestClustersRegistry registryProvider;

public void register(String path, TestClustersAware task) {
tasksMap.put(path, task);
}

public void registry(TestClustersRegistry registry) {
registryProvider = registry;
}

@Override
public void onFinish(FinishEvent finishEvent) {
if (finishEvent instanceof TaskFinishEvent) {
// Handle task finish event...
TaskFinishEvent taskFinishEvent = (TaskFinishEvent)finishEvent;
String taskPath = taskFinishEvent.getDescriptor().getTaskPath();
TestClustersAware task = tasksMap.get(taskPath);
// always unclaim the cluster, even if _this_ task is up-to-date, as others might not have been
// and caused the cluster to start.
task.getClusters().forEach(cluster -> registryProvider.stopCluster(cluster, taskFinishEvent.getResult() instanceof TaskFailureResult));
}
}

}
}
8 changes: 8 additions & 0 deletions build-tools/src/main/java/org/elasticsearch/gradle/testclusters/TestClustersRegistry.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,14 @@ public void maybeStartCluster(ElasticsearchCluster cluster) {
cluster.start();
}

public void registerStartClusterAction(ElasticsearchCluster cluster) {
if (runningClusters.contains(cluster)) {
return;
}
runningClusters.add(cluster);
cluster.start();
}

public void stopCluster(ElasticsearchCluster cluster, boolean taskFailed) {
if (taskFailed) {
// If the task fails, and other tasks use this cluster, the other task will likely never be
Expand Down

0 comments on commit 17517cc

Please sign in to comment.