TestClustersPlugin.java

/*
 * SPDX-License-Identifier: Apache-2.0
 *
 * The OpenSearch Contributors require contributions made to
 * this file be licensed under the Apache-2.0 license or a
 * compatible open source license.
 */

/*
 * Licensed to Elasticsearch under one or more contributor
 * license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright
 * ownership. Elasticsearch licenses this file to you under
 * the Apache License, Version 2.0 (the "License"); you may
 * not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
/*
 * Modifications Copyright OpenSearch Contributors. See
 * GitHub history for details.
 */

package org.opensearch.gradle.testclusters;

import org.opensearch.gradle.DistributionDownloadPlugin;
import org.opensearch.gradle.JdkDownloadPlugin;
import org.opensearch.gradle.ReaperPlugin;
import org.opensearch.gradle.ReaperService;
import org.opensearch.gradle.VersionProperties;
import org.opensearch.gradle.info.BuildParams;
import org.opensearch.gradle.info.GlobalBuildInfoPlugin;
import org.opensearch.gradle.internal.InternalDistributionDownloadPlugin;
import org.opensearch.gradle.util.GradleUtils;
import org.gradle.api.NamedDomainObjectContainer;
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;
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 org.gradle.api.tasks.testing.Test;

import javax.inject.Inject;

import java.io.File;
import java.io.FileOutputStream;
import java.io.InputStream;
import java.net.URI;
import java.util.Arrays;
import java.util.List;

import static org.opensearch.gradle.util.GradleUtils.noop;

public class TestClustersPlugin implements Plugin<Project> {

    public static final String EXTENSION_NAME = "testClusters";
    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);

    @Inject
    protected FileSystemOperations getFileSystemOperations() {
        throw new UnsupportedOperationException();
    }

    @Inject
    protected ArchiveOperations getArchiveOperations() {
        throw new UnsupportedOperationException();
    }

    @Override
    public void apply(Project project) {
        project.getPluginManager().apply(JdkDownloadPlugin.class);
        project.getRootProject().getPluginManager().apply(GlobalBuildInfoPlugin.class);
        if (BuildParams.isInternal()) {
            project.getPlugins().apply(InternalDistributionDownloadPlugin.class);
        } else {
            project.getPlugins().apply(DistributionDownloadPlugin.class);
        }
        project.getRootProject().getPluginManager().apply(ReaperPlugin.class);

        ReaperService reaper = project.getRootProject().getExtensions().getByType(ReaperService.class);

        // enable the DSL to describe clusters
        NamedDomainObjectContainer<OpenSearchCluster> container = createTestClustersContainerExtension(project, reaper);

        // provide a task to be able to list defined clusters.
        createListClustersTask(project, container);

        // register cluster registry as a global build service
        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()
            .getSharedServices()
            .registerIfAbsent(
                THROTTLE_SERVICE_NAME,
                TestClustersThrottle.class,
                spec -> spec.getMaxParallelUsages().set(Math.max(1, project.getGradle().getStartParameter().getMaxWorkerCount() / 2))
            );

        // register cluster hooks
        project.getRootProject().getPluginManager().apply(TestClustersHookPlugin.class);

        // configure installed plugins from -PinstalledPlugins property
        configureInstalledPlugins(project, container);
    }

    @SuppressWarnings("unchecked")
    private void configureInstalledPlugins(Project project, NamedDomainObjectContainer<OpenSearchCluster> container) {
        String installedPluginsProp = System.getProperty("installedPlugins");
        boolean securityEnabled = "true".equals(System.getProperty("security.enabled"));

        if (installedPluginsProp == null && !securityEnabled) {
            return;
        }

        project.afterEvaluate(p -> {
            container.configureEach(cluster -> {
                List<String> plugins = new java.util.ArrayList<>();

                // Parse explicit plugin list if provided
                if (installedPluginsProp != null) {
                    try {
                        plugins.addAll((List<String>) groovy.util.Eval.me(installedPluginsProp));
                    } catch (Exception e) {
                        throw new IllegalArgumentException("Invalid installedPlugins format. Expected list like: ['plugin1','plugin2']", e);
                    }
                }

                // Auto-add security plugin if security is enabled and not already in list
                if (securityEnabled && plugins.stream().noneMatch(p2 -> p2.contains("opensearch-security"))) {
                    plugins.add("opensearch-security");
                }

                for (String pluginName : plugins) {
                    // Check if it's a local project plugin
                    if (project.findProject(":plugins:" + pluginName) != null) {
                        cluster.plugin("plugins:" + pluginName);
                    } else {
                        // Resolve from Maven
                        project.getRepositories().mavenLocal();
                        project.getRepositories().maven(repo -> {
                            repo.setName("OpenSearch Snapshots");
                            repo.setUrl("https://ci.opensearch.org/ci/dbc/snapshots/maven/");
                        });

                        String version = VersionProperties.getOpenSearch().replace("-SNAPSHOT", ".0-SNAPSHOT");
                        String coords = pluginName.contains(":") ? pluginName : (pluginName + ":" + version);
                        if (!coords.contains("@")) {
                            coords += "@zip";
                        }
                        String fullCoords = coords.split(":").length < 3 ? "org.opensearch.plugin:" + coords : coords;

                        org.gradle.api.artifacts.Configuration config = project.getConfigurations()
                            .detachedConfiguration(project.getDependencies().create(fullCoords));
                        config.getResolutionStrategy().cacheChangingModulesFor(0, "seconds");
                        cluster.plugin(project.getLayout().file(project.provider(() -> config.getSingleFile())));

                        // Configure security if opensearch-security plugin is installed
                        if (pluginName.contains("opensearch-security")) {
                            configureSecurityPlugin(project, cluster);
                        }
                    }
                }
            });
        });
    }

    private static final String[] SECURITY_CERT_FILES = {
        "esnode.pem",
        "esnode-key.pem",
        "kirk.pem",
        "kirk-key.pem",
        "root-ca.pem",
        "sample.pem",
        "test-kirk.jks" };
    private static final String SECURITY_CERTS_URL =
        "https://raw.githubusercontent.com/opensearch-project/security/refs/heads/main/bwc-test/src/test/resources/security/";

    private void configureSecurityPlugin(Project project, OpenSearchCluster cluster) {
        File buildDir = project.getBuildDir();
        File securityResourcesDir = new File(buildDir, "security-resources");

        // Download certificates
        for (String certFile : SECURITY_CERT_FILES) {
            File localFile = new File(securityResourcesDir, certFile);
            if (!localFile.exists()) {
                try {
                    localFile.getParentFile().mkdirs();
                    try (
                        InputStream in = URI.create(SECURITY_CERTS_URL + certFile).toURL().openStream();
                        FileOutputStream out = new FileOutputStream(localFile)
                    ) {
                        byte[] buffer = new byte[8192];
                        int bytesRead;
                        while ((bytesRead = in.read(buffer)) != -1) {
                            out.write(buffer, 0, bytesRead);
                        }
                    }
                } catch (Exception e) {
                    throw new RuntimeException("Failed to download security certificate: " + certFile, e);
                }
            }
        }

        // Add security-resources directory to test classpath
        project.getTasks().withType(Test.class).configureEach(testTask -> {
            testTask.setClasspath(testTask.getClasspath().plus(project.files(securityResourcesDir)));
        });

        // Configure each node with certificates and security settings
        cluster.getNodes().forEach(node -> {
            // Add certificate files
            Arrays.stream(SECURITY_CERT_FILES).forEach(cert -> node.extraConfigFile(cert, new File(securityResourcesDir, cert)));

            // Security settings
            node.setting("plugins.security.ssl.transport.pemcert_filepath", "esnode.pem");
            node.setting("plugins.security.ssl.transport.pemkey_filepath", "esnode-key.pem");
            node.setting("plugins.security.ssl.transport.pemtrustedcas_filepath", "root-ca.pem");
            node.setting("plugins.security.ssl.transport.enforce_hostname_verification", "false");
            node.setting("plugins.security.ssl.http.enabled", "true");
            node.setting("plugins.security.ssl.http.pemcert_filepath", "esnode.pem");
            node.setting("plugins.security.ssl.http.pemkey_filepath", "esnode-key.pem");
            node.setting("plugins.security.ssl.http.pemtrustedcas_filepath", "root-ca.pem");
            node.setting("plugins.security.allow_unsafe_democertificates", "true");
            node.setting("plugins.security.allow_default_init_securityindex", "true");
            node.setting("plugins.security.authcz.admin_dn", "\n - CN=kirk,OU=client,O=client,L=test,C=de");
            node.setting("plugins.security.audit.type", "internal_opensearch");
            node.setting("plugins.security.enable_snapshot_restore_privilege", "true");
            node.setting("plugins.security.check_snapshot_restore_write_privileges", "true");
            node.setting("plugins.security.restapi.roles_enabled", "[\"all_access\", \"security_rest_api_access\"]");
            node.setting("plugins.security.system_indices.enabled", "true");
        });
    }

    private NamedDomainObjectContainer<OpenSearchCluster> createTestClustersContainerExtension(Project project, ReaperService reaper) {
        // Create an extensions that allows describing clusters
        NamedDomainObjectContainer<OpenSearchCluster> container = project.container(
            OpenSearchCluster.class,
            name -> new OpenSearchCluster(
                name,
                project,
                reaper,
                new File(project.getBuildDir(), "testclusters"),
                getFileSystemOperations(),
                getArchiveOperations()
            )
        );
        project.getExtensions().add(EXTENSION_NAME, container);
        return container;
    }

    private void createListClustersTask(Project project, NamedDomainObjectContainer<OpenSearchCluster> container) {
        // Task is never up to date so we can pass an lambda for the task action
        project.getTasks().register(LIST_TASK_NAME, task -> {
            task.setGroup("OpenSearch cluster formation");
            task.setDescription("Lists all OpenSearch clusters configured for this project");
            task.doLast(
                (Task t) -> container.forEach(cluster -> logger.lifecycle("   * {}: {}", cluster.getName(), cluster.getNumberOfNodes()))
            );
        });

    }

    static class TestClustersHookPlugin implements Plugin<Project> {
        @Override
        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
            );
            TestClustersRegistry registry = registryProvider.get();

            // 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);
                }

                @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) {}
            });
        }
    }
}