Cluster Management

package org.elasticsearch.test;

import org.elasticsearch.client.internal.Client;
import org.elasticsearch.cluster.ClusterState;
import org.elasticsearch.cluster.node.DiscoveryNode;
import org.elasticsearch.cluster.service.ClusterService;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.core.TimeValue;
import org.elasticsearch.node.Node;
import org.elasticsearch.plugins.Plugin;

/**
 * Internal implementation of TestCluster providing comprehensive multi-node cluster management.
 * Handles node startup, shutdown, configuration, and cluster state management for integration testing.
 */
public final class InternalTestCluster extends TestCluster {
    
    /**
     * Starts a new data node with random configuration.
     * 
     * @return name of the started node
     */
    public String startNode();
    
    /**
     * Starts a new node with the specified settings.
     * 
     * @param settings node configuration settings
     * @return name of the started node
     */
    public String startNode(Settings settings);
    
    /**
     * Starts a data node with custom configuration.
     * 
     * @param settings node settings
     * @param version transport version for the node
     * @return name of the started node
     */
    public String startDataOnlyNode(Settings settings, TransportVersion version);
    
    /**
     * Starts a master-eligible node.
     * 
     * @param settings node configuration
     * @return name of the started master node
     */
    public String startMasterOnlyNode(Settings settings);
    
    /**
     * Starts a coordinating-only node (no data, no master).
     * 
     * @param settings node configuration
     * @return name of the started coordinating node
     */
    public String startCoordinatingOnlyNode(Settings settings);
    
    /**
     * Stops the specified node gracefully.
     * 
     * @param node name or node reference to stop
     * @return true if node was stopped successfully
     */
    public boolean stopRandomNode();
    
    /**
     * Stops a specific node.
     * 
     * @param node node name or predicate to identify node
     */
    public void stopNode(String node);
    
    /**
     * Restarts a node that was previously stopped.
     * 
     * @param node node name to restart
     * @param callback optional callback executed during restart
     * @return name of the restarted node
     */
    public String restartNode(String node, InternalTestCluster.RestartCallback callback);
    
    /**
     * Restarts a random node in the cluster.
     * 
     * @param callback optional restart callback
     * @return name of the restarted node
     */
    public String restartRandomNode(InternalTestCluster.RestartCallback callback);
    
    /**
     * Restarts all nodes in the cluster.
     * 
     * @param callback callback executed for each node restart
     */
    public void restartAllNodes(InternalTestCluster.RestartCallback callback);
    
    /**
     * Ensures the cluster has at least the specified number of data nodes.
     * Starts additional nodes if necessary.
     * 
     * @param n minimum number of data nodes required
     */
    public void ensureAtLeastNumDataNodes(int n);
    
    /**
     * Ensures the cluster has at most the specified number of data nodes.
     * Stops excess nodes if necessary.
     * 
     * @param n maximum number of data nodes allowed
     */
    public void ensureAtMostNumDataNodes(int n);
    
    /**
     * Returns a client connected to a specific node.
     * 
     * @param nodeName name of the node
     * @return client for the specified node
     */
    public Client client(String nodeName);
    
    /**
     * Returns a client connected to a random node.
     * 
     * @return client for a random node
     */
    public Client client();
    
    /**
     * Returns a client that performs operations on the master node.
     * 
     * @return master client
     */
    public Client masterClient();
    
    /**
     * Returns a client connected to a non-master node.
     * 
     * @return non-master client
     */
    public Client nonMasterClient();
    
    /**
     * Returns a client connected to a data node.
     * 
     * @return data node client
     */
    public Client dataNodeClient();
    
    /**
     * Gets the current master node.
     * 
     * @return name of the master node
     */
    public String getMasterName();
    
    /**
     * Gets all node names in the cluster.
     * 
     * @return set of all node names
     */
    public Set<String> getNodeNames();
    
    /**
     * Gets names of all data nodes.
     * 
     * @return set of data node names
     */
    public Set<String> getDataNodeNames();
    
    /**
     * Gets names of all master-eligible nodes.
     * 
     * @return set of master node names
     */  
    public Set<String> getMasterEligibleNodeNames();
    
    /**
     * Returns the ClusterService for a specific node.
     * 
     * @param node node name
     * @return ClusterService instance
     */
    public ClusterService clusterService(String node);
    
    /**
     * Returns the current cluster state.
     * 
     * @return current ClusterState
     */
    public ClusterState clusterState();
    
    /**
     * Waits for the cluster to reach the specified state.
     * 
     * @param statePredicate predicate to match desired state
     * @param timeout maximum time to wait
     */
    public void waitForState(Predicate<ClusterState> statePredicate, TimeValue timeout);
    
    /**
     * Waits for all nodes to agree on the cluster state.
     * 
     * @param timeout maximum time to wait
     */
    public void waitForConsensus(TimeValue timeout);
    
    /**
     * Forces a master election by stopping the current master.
     */
    public void forceMasterElection();
    
    /**
     * Validates that the cluster is properly formed and healthy.
     */
    public void validateClusterFormed();
    
    /**
     * Interface for callbacks executed during node restarts.
     */
    public static interface RestartCallback {
        /**
         * Called with the settings that will be used for the restarted node.
         * 
         * @param nodeSettings current node settings
         * @return modified settings for restart
         */
        Settings onNodeStopped(String nodeName);
        
        /**
         * Called before the node is restarted.
         * 
         * @param nodeName name of the node being restarted
         */
        boolean clearData(String nodeName);
    }
}

import org.elasticsearch.test.ESIntegTestCase;
import org.elasticsearch.test.InternalTestCluster;

public class ClusterManagementTest extends ESIntegTestCase {
    
    public void testMultiNodeCluster() {
        InternalTestCluster cluster = internalCluster();
        
        // Start with minimum nodes
        cluster.ensureAtLeastNumDataNodes(3);
        
        // Verify cluster formation
        cluster.validateClusterForformed();
        
        assertThat(cluster.size(), greaterThanOrEqualTo(3));
        assertThat(cluster.numDataNodes(), greaterThanOrEqualTo(3));
    }
    
    public void testNodeLifecycle() {
        InternalTestCluster cluster = internalCluster();
        
        // Start additional nodes
        String dataNode = cluster.startDataOnlyNode(Settings.EMPTY);
        String masterNode = cluster.startMasterOnlyNode(Settings.EMPTY);
        String coordinatingNode = cluster.startCoordinatingOnlyNode(Settings.EMPTY);
        
        // Verify nodes are running
        assertTrue(cluster.getNodeNames().contains(dataNode));
        assertTrue(cluster.getNodeNames().contains(masterNode));
        assertTrue(cluster.getNodeNames().contains(coordinatingNode));
        
        // Test node restart
        cluster.restartNode(dataNode, new InternalTestCluster.RestartCallback() {
            @Override
            public Settings onNodeStopped(String nodeName) {
                return Settings.builder()
                    .put("node.attr.restarted", true)
                    .build();
            }
            
            @Override
            public boolean clearData(String nodeName) {
                return false; // Keep data during restart
            }
        });
        
        // Verify node restarted successfully
        assertTrue(cluster.getNodeNames().contains(dataNode));
    }
    
    public void testMasterElection() {
        InternalTestCluster cluster = internalCluster();
        
        // Start master-eligible nodes
        cluster.ensureAtLeastNumDataNodes(3);
        
        String originalMaster = cluster.getMasterName();
        assertThat(originalMaster, notNullValue());
        
        // Force master election
        cluster.forceMasterElection();
        
        // Wait for new master
        cluster.waitForState(state -> 
            state.nodes().getMasterNode() != null, 
            TimeValue.timeValueSeconds(30)
        );
        
        String newMaster = cluster.getMasterName();
        assertThat(newMaster, notNullValue());
    }
    
    public void testClientTypes() {
        InternalTestCluster cluster = internalCluster();
        cluster.ensureAtLeastNumDataNodes(2);
        
        // Test different client types
        Client masterClient = cluster.masterClient();
        Client dataClient = cluster.dataNodeClient(); 
        Client randomClient = cluster.client();
        
        // All clients should be able to perform operations
        masterClient.admin().cluster().prepareHealth().get();
        dataClient.admin().cluster().prepareHealth().get();
        randomClient.admin().cluster().prepareHealth().get();
    }
}

package org.elasticsearch.test;

import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.plugins.Plugin;

/**
 * Interface for customizing node configuration in test clusters.
 * Allows injection of custom settings, plugins, and node-specific configuration.
 */
public interface NodeConfigurationSource {
    
    /**
     * Returns settings that should be applied to all nodes.
     * 
     * @param nodeOrdinal ordinal number of the node being configured
     * @return settings for the node
     */
    Settings nodeSettings(int nodeOrdinal);
    
    /**
     * Returns plugins that should be installed on cluster nodes.
     * 
     * @return collection of plugin classes
     */
    Collection<Class<? extends Plugin>> nodePlugins();
    
    /**
     * Returns settings that should be applied to transport clients.
     * 
     * @return transport client settings
     */
    Settings transportClientSettings();
}

package org.elasticsearch.test;

import org.elasticsearch.cluster.service.ClusterService;
import org.elasticsearch.cluster.ClusterState;
import org.elasticsearch.cluster.ClusterStateUpdateTask;

/**
 * Utilities for working with ClusterService in tests.
 */
public class ClusterServiceUtils {
    
    /**
     * Sets a new cluster state directly (bypassing normal update mechanisms).
     * Use with caution - primarily for testing specific cluster states.
     * 
     * @param clusterService target cluster service
     * @param clusterState new cluster state to set
     */
    public static void setState(ClusterService clusterService, ClusterState clusterState);
    
    /**
     * Adds a high priority cluster state update task.
     * 
     * @param clusterService target cluster service
     * @param reason reason for the update
     * @param task update task to execute
     */
    public static void submitStateUpdateTask(ClusterService clusterService, 
                                           String reason, 
                                           ClusterStateUpdateTask task);
    
    /**
     * Creates a ClusterService instance for testing with the specified initial state.
     * 
     * @param initialState initial cluster state
     * @return new ClusterService instance
     */
    public static ClusterService createClusterService(ClusterState initialState);
    
    /**
     * Creates a minimal cluster state for testing.
     * 
     * @param clusterName name of the test cluster
     * @return basic cluster state
     */
    public static ClusterState createClusterState(String clusterName);
}

package org.elasticsearch.cluster.node;

import org.elasticsearch.cluster.node.DiscoveryNode;
import org.elasticsearch.common.transport.TransportAddress;

/**
 * Utilities for creating and managing discovery nodes in tests.
 */
public class DiscoveryNodeUtils {
    
    /**
     * Creates a new discovery node with random configuration.
     * 
     * @param id node identifier
     * @return configured DiscoveryNode
     */
    public static DiscoveryNode create(String id);
    
    /**
     * Creates a discovery node with specific roles.
     * 
     * @param id node identifier
     * @param address transport address
     * @param roles set of node roles
     * @return configured DiscoveryNode
     */
    public static DiscoveryNode create(String id, 
                                     TransportAddress address, 
                                     Set<DiscoveryNodeRole> roles);
    
    /**
     * Creates a master-eligible node.
     * 
     * @param id node identifier
     * @return master-eligible DiscoveryNode
     */
    public static DiscoveryNode createMasterNode(String id);
    
    /**
     * Creates a data-only node.
     * 
     * @param id node identifier
     * @return data-only DiscoveryNode
     */
    public static DiscoveryNode createDataNode(String id);
    
    /**
     * Creates a coordinating-only node.
     * 
     * @param id node identifier  
     * @return coordinating-only DiscoveryNode
     */
    public static DiscoveryNode createCoordinatingNode(String id);
    
    /**
     * Generates a random transport address for testing.
     * 
     * @return random TransportAddress
     */
    public static TransportAddress randomAddress();
}

package org.elasticsearch.test;

/**
 * Configuration options for test cluster scope and behavior.
 */
public class ClusterScope {
    
    /**
     * Defines when cluster instances are created and destroyed.
     */
    public enum Scope {
        /** Single cluster instance for entire test suite */
        SUITE,
        /** New cluster instance for each test method */
        TEST
    }
    
    /**
     * Transport client usage ratio for testing client behavior.
     * 0.0 = always use node clients, 1.0 = always use transport clients
     */
    public static final double TRANSPORT_CLIENT_RATIO = 0.3;
    
    /**
     * Default minimum number of master-eligible nodes for test clusters.
     */
    public static final int DEFAULT_MIN_MASTER_NODES = 1;
    
    /**
     * Default settings applied to all test clusters.
     */
    public static final Settings DEFAULT_SETTINGS = Settings.builder()
        .put("discovery.type", "zen")
        .put("transport.type", "netty4")
        .put("http.type", "netty4") 
        .put("cluster.routing.rebalance.enable", "all")
        .build();
}

import org.elasticsearch.test.ESIntegTestCase;
import org.elasticsearch.cluster.node.DiscoveryNodeRole;

@ESIntegTestCase.ClusterScope(
    scope = ESIntegTestCase.Scope.TEST,
    numDataNodes = 5, 
    numClientNodes = 1
)
public class AdvancedClusterTest extends ESIntegTestCase {
    
    public void testLargeClusterOperations() {
        InternalTestCluster cluster = internalCluster();
        
        // Verify cluster size
        assertThat(cluster.numDataNodes(), equalTo(5));
        
        // Test rolling restart of all data nodes
        Set<String> dataNodes = cluster.getDataNodeNames();
        for (String node : dataNodes) {
            cluster.restartNode(node, new InternalTestCluster.RestartCallback() {
                @Override
                public Settings onNodeStopped(String nodeName) {
                    return Settings.builder()
                        .put("node.attr.generation", "2")
                        .build();
                }
                
                @Override  
                public boolean clearData(String nodeName) {
                    return false;
                }
            });
            
            // Wait for cluster to stabilize after each restart
            ensureGreen();
        }
        
        // Verify all nodes have new attribute
        NodesInfoResponse nodesInfo = client().admin().cluster().prepareNodesInfo().get();
        for (NodeInfo nodeInfo : nodesInfo.getNodes()) {
            if (nodeInfo.getNode().isDataNode()) {
                assertThat(nodeInfo.getNode().getAttributes().get("generation"), equalTo("2"));
            }
        }
    }
    
    public void testMasterFailover() {
        InternalTestCluster cluster = internalCluster();
        
        // Ensure we have multiple master-eligible nodes
        cluster.ensureAtLeastNumDataNodes(3);
        
        String originalMaster = cluster.getMasterName();
        
        // Stop the current master
        cluster.stopNode(originalMaster);
        
        // Wait for new master election
        cluster.waitForState(state -> {
            DiscoveryNode master = state.nodes().getMasterNode();
            return master != null && !master.getName().equals(originalMaster);
        }, TimeValue.timeValueSeconds(30));
        
        // Verify cluster is still functional
        ensureGreen();
        
        String newMaster = cluster.getMasterName();
        assertThat(newMaster, not(equalTo(originalMaster)));
    }
    
    @Override
    protected Settings nodeSettings(int nodeOrdinal) {
        return Settings.builder()
            .put(super.nodeSettings(nodeOrdinal))
            .put("node.name", "test-node-" + nodeOrdinal)
            .put("cluster.routing.allocation.disk.threshold_enabled", false)
            .build();
    }
}