/jira-project/jira-components/jira-plugins/jira-ha-plugin/ha-container-api/src/main/java/com/atlassian/jira/plugins/ha/container/TomcatContainerSet.java

package com.atlassian.jira.plugins.ha.container;

import com.atlassian.jira.config.database.DatabaseConfig;
import com.atlassian.jira.configurator.config.JiraHomeDatabaseConfigurationLoader;
import com.atlassian.jira.functest.framework.backdoor.ZeroDowntimeControl;
import com.atlassian.jira.plugins.ha.testapi.test.JiraCluster;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import org.ofbiz.core.entity.config.DatasourceInfo;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedReader;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStreamReader;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardCopyOption;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.time.Duration;
import java.time.Instant;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.TimeUnit;
import java.util.function.Predicate;
import java.util.function.Supplier;
import java.util.stream.Stream;

public class TomcatContainerSet implements ContainerSet<TomcatContainer> {
    private static final Logger log = LoggerFactory.getLogger(TomcatContainerSet.class);

    private final List<TomcatContainer> containers;
    private final Supplier<JiraCluster> clusterMaker;
    private final Path sharedHome;
    private final Path baseArtifactDirectory;
    private final Map<JiraType, ? extends Path> warArchiveMap;
    private final Path dbConfigFile;
    private JiraCluster cluster;

    TomcatContainerSet(List<? extends TomcatContainer> containers, Supplier<JiraCluster> clusterMaker,
                              Path sharedHome, Path baseArtifactDirectory, Map<JiraType, ? extends Path> warArchiveMap,
                              Path dbConfigFile) {
        this.containers = ImmutableList.copyOf(containers);
        this.clusterMaker = clusterMaker;
        this.sharedHome = sharedHome;
        this.baseArtifactDirectory = baseArtifactDirectory;
        this.warArchiveMap = ImmutableMap.copyOf(warArchiveMap);
        this.dbConfigFile = dbConfigFile;
    }

    @Override
    public List<? extends TomcatContainer> getContainers() {
        return containers;
    }

    @Override
    public void startAll() {
        getContainers().forEach(Container::start);
    }

    @Override
    public void stopAll() {
        getContainers().forEach(Container::stop);
    }

    @Override
    public synchronized JiraCluster cluster() {
        if (cluster == null) {
            cluster = clusterMaker.get();
        }
        return cluster;
    }

    @Override
    public void installJira(JiraType warType) {
        containers.forEach(c -> c.installJira(warType));
    }

    @Override
    public Path jiraWarFile(JiraType warType) {
        return warArchiveMap.get(warType);
    }

    @Override
    public void installWar(Path warFile, String context) {
        containers.forEach(c -> c.installWar(warFile, context));
    }

    private Path installedPluginsDirectory() {
        return sharedHome.resolve("plugins").resolve("installed-plugins");
    }

    public void removeAllSharedPlugins() {
        removeSharedPlugins(pluginFile -> true);
    }

    @Override
    public void installTestingPlugins(JiraType jiraType) {
        Path pluginsDir = baseArtifactDirectory.resolve("plugins-" + jiraType.name().toLowerCase(Locale.ROOT));
        try (Stream<Path> fileList = Files.list(pluginsDir)) {
            fileList.filter(file -> file.toString().endsWith(".jar")).forEach(this::installSharedPlugin);
        } catch (IOException e) {
            throw new RuntimeException("Failed to list files in plugins directory: " + e, e);
        }
    }

    @Override
    public void installSharedPlugin(Path pluginFile) {
        Path installedPluginsDirectory = installedPluginsDirectory();
        try {
            Files.createDirectories(installedPluginsDirectory);
        } catch (IOException e) {
            throw new RuntimeException("Failed to create installed-plugins directory under shared home: " + e, e);
        }

        Path targetPluginFile = installedPluginsDirectory.resolve(pluginFile.getFileName());
        try {
            Files.copy(pluginFile, targetPluginFile, StandardCopyOption.REPLACE_EXISTING);
        } catch (IOException e) {
            throw new RuntimeException("Failed to copy plugin to shared home: " + e, e);
        }
    }

    public void removeSharedPlugins(Predicate<? super Path> filter) {
        Path installedPluginsDirectory = installedPluginsDirectory();
        if (Files.notExists(installedPluginsDirectory))
            return;
        try (Stream<Path> fileList = Files.list(installedPluginsDirectory)) {
            fileList.filter(filter).forEach(file -> {
                try {
                    Files.delete(file);
                } catch (IOException e) {
                    throw new RuntimeException("Failed to delete plugin " + file.toString() + ": " + e, e);
                }
            });
        } catch (IOException e) {
            throw new RuntimeException("Failed to delete shared plugins: " + e, e);
        }
    }

    @Override
    public void resetDatabase() {
        //Use driver manager to read dbconfig.xml

        JiraHomeDatabaseConfigurationLoader loader = new JiraHomeDatabaseConfigurationLoader(dbConfigFile.getParent().toAbsolutePath().toString());

        DatabaseConfig dbConfig = loader.loadDatabaseConfiguration();
        DatasourceInfo datasourceInfo = dbConfig.getDatasourceInfo();

        String uri = datasourceInfo.getJdbcDatasource().getUri();
        String user = datasourceInfo.getJdbcDatasource().getUsername();
        String password = datasourceInfo.getJdbcDatasource().getPassword();

        try (Connection connection = DriverManager.getConnection(uri, user, password)) {
            runDbScript("reset-" + dbConfig.getDatabaseType(), connection);
        } catch (SQLException | IOException e) {
            throw new RuntimeException("Error restoring database dump: " + e, e);
        }
    }

    private static void runDbScript(String scriptName, Connection connection) throws SQLException, IOException {
        String resourceName = "/dbscript/" + scriptName + ".sql";
        URL scriptResource = TomcatContainerSet.class.getResource(resourceName);
        if (scriptResource == null) {
            throw new FileNotFoundException("Missing test resource: " + resourceName);
        }

        try (BufferedReader scriptReader = new BufferedReader(new InputStreamReader(scriptResource.openStream(), StandardCharsets.UTF_8))) {
            String line;
            do {
                line = scriptReader.readLine();
                if (line != null) {
                    try (PreparedStatement stat = connection.prepareStatement(line)) {
                        stat.executeUpdate();
                    }
                }
            } while (line != null);
        }
    }

    @Override
    public Path getSharedHome() {
        return sharedHome;
    }

    @Override
    public void waitForUpgradeState(ZeroDowntimeControl.UpgradeState targetState, Duration timeout) {
        Instant timeoutTime = Instant.now().plus(timeout);
        Set<ZeroDowntimeControl.UpgradeState> statesObserved = new LinkedHashSet<>();
        try {
            do {
                ZeroDowntimeControl.ClusterState currentState = cluster.anyNode().backdoor().zdu().currentState();
                if (currentState.getState() == targetState) {
                    return; //Got to our target
                }
                statesObserved.add(currentState.getState());
                Thread.sleep(200L);
            } while (Instant.now().isBefore(timeoutTime));
        } catch (InterruptedException e) {
            throw new RuntimeException("Interrupted waiting for target state " + targetState);
        }

        //If we get here we timed out waiting
        throw new RuntimeException("Timed out (" + timeout + ") waiting for target state " + targetState + ", observed " + statesObserved);
    }

    @Override
    public void waitForStartupOfActiveContainers(Duration timeout) {

        int startedContainers = Math.toIntExact(getContainers().stream()
                                        .map(container -> container.getCargoContainer().getState())
                                        .filter(state -> state.isStarted() || state.isStarting())
                                        .count());

        log.debug("Waiting for startup of " + startedContainers + " containers...");

        //Use first active node to perform replication operation - it doesn't actually matter which one we use
        getContainers().stream()
                .filter(container -> container.getCargoContainer().getState().isStarted())
                .findFirst()
                .orElseThrow(() -> new RuntimeException("No nodes have started up"))
                .replication().waitForClusterNodes(startedContainers, timeout.toMillis(), TimeUnit.MILLISECONDS);

        log.debug(startedContainers + " containers have now started up");
    }

    @Override
    public void waitForStartupOfActiveContainers() {
        waitForStartupOfActiveContainers(Duration.ofMinutes(5));
    }

}