diff --git a/src/main/java/io/clustercontroller/models/SearchUnitGoalState.java b/src/main/java/io/clustercontroller/models/SearchUnitGoalState.java index ba657e04..26c5f77c 100644 --- a/src/main/java/io/clustercontroller/models/SearchUnitGoalState.java +++ b/src/main/java/io/clustercontroller/models/SearchUnitGoalState.java @@ -98,4 +98,39 @@ public String toString() { ", version=" + version + '}'; } + + /** + * Compares goal states based on localShards content only. + * lastUpdated and version are metadata fields that don't affect the actual goal state. + * This is used to avoid unnecessary writes to etcd when the goal state hasn't actually changed. + */ + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + SearchUnitGoalState that = (SearchUnitGoalState) o; + + // Compare localShards content - this is the actual goal state + if (localShards == null && that.localShards == null) return true; + if (localShards == null || that.localShards == null) return false; + + return localShards.equals(that.localShards); + } + + @Override + public int hashCode() { + return localShards != null ? localShards.hashCode() : 0; + } + + /** + * Check if a specific index/shard with role already exists in the goal state. + * This is used to determine if an update is actually needed. + */ + public boolean hasShardWithRole(String indexName, String shardId, String role) { + if (localShards == null) return false; + Map shards = localShards.get(indexName); + if (shards == null) return false; + String existingRole = shards.get(shardId); + return role.equals(existingRole); + } } diff --git a/src/main/java/io/clustercontroller/orchestration/ImmediateOrchestrationStrategy.java b/src/main/java/io/clustercontroller/orchestration/ImmediateOrchestrationStrategy.java index 2cfc5a88..44d94d62 100644 --- a/src/main/java/io/clustercontroller/orchestration/ImmediateOrchestrationStrategy.java +++ b/src/main/java/io/clustercontroller/orchestration/ImmediateOrchestrationStrategy.java @@ -107,19 +107,28 @@ private void updateNodeGoalState(String nodeId, String indexName, String shardId // Get current goal state for the node SearchUnitGoalState currentGoalState = metadataStore.getSearchUnitGoalState(clusterId, nodeId); + // Determine role based on whether this node is in IngestSUs or SearchSUs + String role = planned.getIngestSUs().contains(nodeId) ? NodeRole.PRIMARY.getValue() : NodeRole.REPLICA.getValue(); + + // Check if the shard with this role already exists in the goal state + // If so, skip the update to avoid unnecessary etcd writes that trigger data node watchers + if (currentGoalState != null && currentGoalState.hasShardWithRole(indexName, shardId, role)) { + log.debug("Goal state for node {} already has shard {}/{} with role {}, skipping update", + nodeId, indexName, shardId, role); + return; + } + // If null, create new one if (currentGoalState == null) { currentGoalState = new SearchUnitGoalState(); } - // Determine role based on whether this node is in IngestSUs or SearchSUs - String role = planned.getIngestSUs().contains(nodeId) ? NodeRole.PRIMARY.getValue() : NodeRole.REPLICA.getValue(); - // Update goal state with new shard allocation SearchUnitGoalState newGoalState = updateGoalStateForIndexShard(currentGoalState, indexName, shardId, role); // Set new goal state in etcd metadataStore.setSearchUnitGoalState(clusterId, nodeId, newGoalState); + log.info("Updated goal state for node {} with shard {}/{} role {}", nodeId, indexName, shardId, role); } catch (Exception e) { log.error("Failed to update goal state for node {}: {}", nodeId, e.getMessage(), e); diff --git a/src/main/java/io/clustercontroller/orchestration/RollingUpdateOrchestrationStrategy.java b/src/main/java/io/clustercontroller/orchestration/RollingUpdateOrchestrationStrategy.java index be5ec2e2..33923957 100644 --- a/src/main/java/io/clustercontroller/orchestration/RollingUpdateOrchestrationStrategy.java +++ b/src/main/java/io/clustercontroller/orchestration/RollingUpdateOrchestrationStrategy.java @@ -1,6 +1,7 @@ package io.clustercontroller.orchestration; import io.clustercontroller.enums.NodeRole; +import io.clustercontroller.enums.ShardState; import io.clustercontroller.metrics.MetricsProvider; import io.clustercontroller.models.Index; import io.clustercontroller.models.IndexShardProgress; @@ -45,6 +46,10 @@ public void orchestrate(String clusterId) { cleanupStaleGoalStates(clusterId); // PHASE 2: Orchestrate new goal states (rolling update) + // Batch all goal state updates per node to minimize etcd writes + // Key: nodeId, Value: accumulated goal state for that node + Map pendingGoalStateUpdates = new HashMap<>(); + // Get all index configs to iterate over indexes and shards List indexConfigs = metadataStore.getAllIndexConfigs(clusterId); @@ -71,7 +76,7 @@ public void orchestrate(String clusterId) { } // Process this index-shard with rolling update logic - orchestrateIndexShard(indexName, shardId, planned, clusterId); + orchestrateIndexShard(indexName, shardId, planned, clusterId, pendingGoalStateUpdates); } catch (Exception e) { log.error("Failed to orchestrate shard {}/{}: {}", indexName, shardId, e.getMessage(), e); @@ -80,6 +85,9 @@ public void orchestrate(String clusterId) { } } + // PHASE 3: Flush all pending goal state updates to etcd (one write per node) + flushPendingGoalStates(clusterId, pendingGoalStateUpdates); + // Cleanup stale gauges for deleted indices/shards metricsProvider.cleanupStaleGauges(); @@ -89,19 +97,20 @@ public void orchestrate(String clusterId) { } } - private void orchestrateIndexShard(String indexName, String shardId, ShardAllocation planned, String clusterId) { + private void orchestrateIndexShard(String indexName, String shardId, ShardAllocation planned, String clusterId, + Map pendingGoalStateUpdates) { // Process IngestSUs (Primary nodes) separately List ingestSUs = planned.getIngestSUs(); if (!ingestSUs.isEmpty()) { log.debug("Processing {} IngestSUs (primary nodes) for shard {}/{}", ingestSUs.size(), indexName, shardId); - orchestrateNodeGroup(indexName, shardId, ingestSUs, NodeRole.PRIMARY, planned, clusterId); + orchestrateNodeGroup(indexName, shardId, ingestSUs, NodeRole.PRIMARY, planned, clusterId, pendingGoalStateUpdates); } // Process SearchSUs (Replica nodes) separately List searchSUs = planned.getSearchSUs(); if (!searchSUs.isEmpty()) { log.debug("Processing {} SearchSUs (replica nodes) for shard {}/{}", searchSUs.size(), indexName, shardId); - orchestrateNodeGroup(indexName, shardId, searchSUs, NodeRole.REPLICA, planned, clusterId); + orchestrateNodeGroup(indexName, shardId, searchSUs, NodeRole.REPLICA, planned, clusterId, pendingGoalStateUpdates); } if (ingestSUs.isEmpty() && searchSUs.isEmpty()) { @@ -109,7 +118,9 @@ private void orchestrateIndexShard(String indexName, String shardId, ShardAlloca } } - private void orchestrateNodeGroup(String indexName, String shardId, List nodeGroup, NodeRole role, ShardAllocation planned, String clusterId) { + private void orchestrateNodeGroup(String indexName, String shardId, List nodeGroup, NodeRole role, + ShardAllocation planned, String clusterId, + Map pendingGoalStateUpdates) { String indexShard = indexName + "/" + shardId; // Check current progress for this node group @@ -135,11 +146,11 @@ private void orchestrateNodeGroup(String indexName, String shardId, List for (String nodeId : nextBatch) { try { - updateNodeGoalState(nodeId, indexName, shardId, planned, clusterId); + queueNodeGoalStateUpdate(nodeId, indexName, shardId, planned, clusterId, pendingGoalStateUpdates); successfulUpdates++; - log.debug("Started update for {} node {} with shard {}/{}", role, nodeId, indexName, shardId); + log.debug("Queued update for {} node {} with shard {}/{}", role, nodeId, indexName, shardId); } catch (Exception e) { - log.error("Failed to start update for {} node {}: {}", role, nodeId, e.getMessage(), e); + log.error("Failed to queue update for {} node {}: {}", role, nodeId, e.getMessage(), e); } } } else { @@ -221,8 +232,10 @@ private boolean hasActualStateConverged(String nodeId, String indexShard, String } List shards = nodeRouting.get(indexName); + // Only consider converged if shard is STARTED (not INITIALIZING/RELOCATING) return shards.stream() - .anyMatch(shard -> String.valueOf(shard.getShardId()).equals(shardId)); + .anyMatch(shard -> String.valueOf(shard.getShardId()).equals(shardId) + && ShardState.STARTED.equals(shard.getState())); } catch (Exception e) { log.error("Failed to check actual state for node {}: {}", nodeId, e.getMessage(), e); return false; @@ -236,40 +249,103 @@ private List getNodesToUpdate(List allNodes, IndexShardProgress .collect(java.util.stream.Collectors.toList()); } - private void updateNodeGoalState(String nodeId, String indexName, String shardId, ShardAllocation planned, String clusterId) throws Exception { + /** + * Queue a goal state update for a node. Instead of writing to etcd immediately, + * accumulates the change in pendingGoalStateUpdates map to be flushed later. + * This batches all shard assignments per node into a single etcd write. + */ + private void queueNodeGoalStateUpdate(String nodeId, String indexName, String shardId, ShardAllocation planned, + String clusterId, Map pendingGoalStateUpdates) throws Exception { try { - // Get current goal state for the node - SearchUnitGoalState currentGoalState = metadataStore.getSearchUnitGoalState(clusterId, nodeId); - - // If null, create new one - if (currentGoalState == null) { - currentGoalState = new SearchUnitGoalState(); - } - // Determine role based on whether this node is in IngestSUs or SearchSUs String role = planned.getIngestSUs().contains(nodeId) ? NodeRole.PRIMARY.getValue() : NodeRole.REPLICA.getValue(); - // Update goal state with new shard allocation - SearchUnitGoalState newGoalState = updateGoalStateForIndexShard(currentGoalState, indexName, shardId, role); + // Get or create pending goal state for this node + SearchUnitGoalState pendingState = pendingGoalStateUpdates.computeIfAbsent(nodeId, k -> { + try { + // Start with current state from etcd (or empty if none exists) + SearchUnitGoalState current = metadataStore.getSearchUnitGoalState(clusterId, nodeId); + if (current != null) { + // Create a copy to avoid modifying the cached version + SearchUnitGoalState copy = new SearchUnitGoalState(); + copy.setLocalShards(new HashMap<>(current.getLocalShards())); + // Deep copy the inner maps + for (Map.Entry> entry : current.getLocalShards().entrySet()) { + copy.getLocalShards().put(entry.getKey(), new HashMap<>(entry.getValue())); + } + return copy; + } + return new SearchUnitGoalState(); + } catch (Exception e) { + log.warn("Failed to get current goal state for node {}, starting fresh: {}", nodeId, e.getMessage()); + return new SearchUnitGoalState(); + } + }); - // Set new goal state in etcd - metadataStore.setSearchUnitGoalState(clusterId, nodeId, newGoalState); + // Check if the shard with this role already exists + if (pendingState.hasShardWithRole(indexName, shardId, role)) { + log.debug("Goal state for node {} already has shard {}/{} with role {}, skipping", + nodeId, indexName, shardId, role); + return; + } + + // Add shard to pending state (no etcd write here) + pendingState.getLocalShards().computeIfAbsent(indexName, k -> new HashMap<>()).put(shardId, role); + log.debug("Queued goal state update for node {} with shard {}/{} role {}", nodeId, indexName, shardId, role); } catch (Exception e) { - log.error("Failed to update goal state for node {}: {}", nodeId, e.getMessage(), e); + log.error("Failed to queue goal state update for node {}: {}", nodeId, e.getMessage(), e); throw e; } } - private SearchUnitGoalState updateGoalStateForIndexShard(SearchUnitGoalState current, String indexName, String shardId, String role) { - if (current == null) { - current = new SearchUnitGoalState(); + /** + * Flush all pending goal state updates to etcd. Each node gets exactly one write + * containing all accumulated shard assignments. Compares with current state to + * skip writes if unchanged. + * + * @param clusterId the cluster ID + * @param pendingGoalStateUpdates map of nodeId -> accumulated goal state + */ + private void flushPendingGoalStates(String clusterId, Map pendingGoalStateUpdates) { + if (pendingGoalStateUpdates.isEmpty()) { + log.debug("No pending goal state updates to flush"); + return; } - // Add new shard allocation using the localShards structure - current.getLocalShards().computeIfAbsent(indexName, k -> new HashMap<>()).put(shardId, role); + log.info("Flushing {} pending goal state updates to etcd", pendingGoalStateUpdates.size()); + int successCount = 0; + int skipCount = 0; + + for (Map.Entry entry : pendingGoalStateUpdates.entrySet()) { + String nodeId = entry.getKey(); + SearchUnitGoalState newGoalState = entry.getValue(); + + try { + // Read current state from etcd to compare - skip write if unchanged + SearchUnitGoalState currentState = metadataStore.getSearchUnitGoalState(clusterId, nodeId); + + if (currentState != null && currentState.equals(newGoalState)) { + log.debug("Goal state unchanged for node {}, skipping write", nodeId); + skipCount++; + continue; + } + + // Write the accumulated goal state in one etcd PUT + metadataStore.setSearchUnitGoalState(clusterId, nodeId, newGoalState); + successCount++; + + int shardCount = newGoalState.getLocalShards().values().stream() + .mapToInt(Map::size) + .sum(); + log.info("Flushed goal state for node {} ({} total shards)", nodeId, shardCount); + + } catch (Exception e) { + log.error("Failed to flush goal state for node {}: {}", nodeId, e.getMessage(), e); + } + } - return current; + log.info("Goal state flush complete: {} written, {} skipped (unchanged)", successCount, skipCount); } /** diff --git a/src/test/java/io/clustercontroller/models/SearchUnitGoalStateTest.java b/src/test/java/io/clustercontroller/models/SearchUnitGoalStateTest.java new file mode 100644 index 00000000..85a5d5bf --- /dev/null +++ b/src/test/java/io/clustercontroller/models/SearchUnitGoalStateTest.java @@ -0,0 +1,191 @@ +package io.clustercontroller.models; + +import org.junit.jupiter.api.Test; + +import java.util.HashMap; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +class SearchUnitGoalStateTest { + + // ========== equals() tests ========== + + @Test + void testEqualsWithSameLocalShards() { + SearchUnitGoalState state1 = createGoalState("index1", "0", "PRIMARY"); + SearchUnitGoalState state2 = createGoalState("index1", "0", "PRIMARY"); + + assertThat(state1).isEqualTo(state2); + assertThat(state1.hashCode()).isEqualTo(state2.hashCode()); + } + + @Test + void testEqualsIgnoresLastUpdatedAndVersion() { + SearchUnitGoalState state1 = createGoalState("index1", "0", "PRIMARY"); + state1.setLastUpdated("2024-01-01T00:00:00Z"); + state1.setVersion(1); + + SearchUnitGoalState state2 = createGoalState("index1", "0", "PRIMARY"); + state2.setLastUpdated("2024-12-31T23:59:59Z"); + state2.setVersion(999); + + // Should be equal because localShards are the same + assertThat(state1).isEqualTo(state2); + } + + @Test + void testNotEqualsWithDifferentShardId() { + SearchUnitGoalState state1 = createGoalState("index1", "0", "PRIMARY"); + SearchUnitGoalState state2 = createGoalState("index1", "1", "PRIMARY"); + + assertThat(state1).isNotEqualTo(state2); + } + + @Test + void testNotEqualsWithDifferentRole() { + SearchUnitGoalState state1 = createGoalState("index1", "0", "PRIMARY"); + SearchUnitGoalState state2 = createGoalState("index1", "0", "SEARCH_REPLICA"); + + assertThat(state1).isNotEqualTo(state2); + } + + @Test + void testNotEqualsWithDifferentIndex() { + SearchUnitGoalState state1 = createGoalState("index1", "0", "PRIMARY"); + SearchUnitGoalState state2 = createGoalState("index2", "0", "PRIMARY"); + + assertThat(state1).isNotEqualTo(state2); + } + + @Test + void testEqualsWithMultipleShards() { + SearchUnitGoalState state1 = new SearchUnitGoalState(); + Map> localShards1 = new HashMap<>(); + localShards1.put("index1", new HashMap<>()); + localShards1.get("index1").put("0", "PRIMARY"); + localShards1.get("index1").put("1", "SEARCH_REPLICA"); + localShards1.put("index2", new HashMap<>()); + localShards1.get("index2").put("0", "PRIMARY"); + state1.setLocalShards(localShards1); + + SearchUnitGoalState state2 = new SearchUnitGoalState(); + Map> localShards2 = new HashMap<>(); + localShards2.put("index1", new HashMap<>()); + localShards2.get("index1").put("0", "PRIMARY"); + localShards2.get("index1").put("1", "SEARCH_REPLICA"); + localShards2.put("index2", new HashMap<>()); + localShards2.get("index2").put("0", "PRIMARY"); + state2.setLocalShards(localShards2); + + assertThat(state1).isEqualTo(state2); + } + + @Test + void testEqualsWithBothNullLocalShards() { + SearchUnitGoalState state1 = new SearchUnitGoalState(); + state1.setLocalShards(null); + + SearchUnitGoalState state2 = new SearchUnitGoalState(); + state2.setLocalShards(null); + + assertThat(state1).isEqualTo(state2); + } + + @Test + void testNotEqualsWithOneNullLocalShards() { + SearchUnitGoalState state1 = createGoalState("index1", "0", "PRIMARY"); + + SearchUnitGoalState state2 = new SearchUnitGoalState(); + state2.setLocalShards(null); + + assertThat(state1).isNotEqualTo(state2); + } + + @Test + void testEqualsWithEmptyLocalShards() { + SearchUnitGoalState state1 = new SearchUnitGoalState(); + state1.setLocalShards(new HashMap<>()); + + SearchUnitGoalState state2 = new SearchUnitGoalState(); + state2.setLocalShards(new HashMap<>()); + + assertThat(state1).isEqualTo(state2); + } + + // ========== hasShardWithRole() tests ========== + + @Test + void testHasShardWithRoleReturnsTrue() { + SearchUnitGoalState state = createGoalState("index1", "0", "PRIMARY"); + + assertThat(state.hasShardWithRole("index1", "0", "PRIMARY")).isTrue(); + } + + @Test + void testHasShardWithRoleReturnsFalseForDifferentRole() { + SearchUnitGoalState state = createGoalState("index1", "0", "PRIMARY"); + + // Same index and shard, but different role + assertThat(state.hasShardWithRole("index1", "0", "SEARCH_REPLICA")).isFalse(); + } + + @Test + void testHasShardWithRoleReturnsFalseForDifferentShard() { + SearchUnitGoalState state = createGoalState("index1", "0", "PRIMARY"); + + assertThat(state.hasShardWithRole("index1", "1", "PRIMARY")).isFalse(); + } + + @Test + void testHasShardWithRoleReturnsFalseForDifferentIndex() { + SearchUnitGoalState state = createGoalState("index1", "0", "PRIMARY"); + + assertThat(state.hasShardWithRole("index2", "0", "PRIMARY")).isFalse(); + } + + @Test + void testHasShardWithRoleReturnsFalseForNullLocalShards() { + SearchUnitGoalState state = new SearchUnitGoalState(); + state.setLocalShards(null); + + assertThat(state.hasShardWithRole("index1", "0", "PRIMARY")).isFalse(); + } + + @Test + void testHasShardWithRoleReturnsFalseForEmptyLocalShards() { + SearchUnitGoalState state = new SearchUnitGoalState(); + state.setLocalShards(new HashMap<>()); + + assertThat(state.hasShardWithRole("index1", "0", "PRIMARY")).isFalse(); + } + + @Test + void testHasShardWithRoleWithMultipleShards() { + SearchUnitGoalState state = new SearchUnitGoalState(); + Map> localShards = new HashMap<>(); + localShards.put("index1", new HashMap<>()); + localShards.get("index1").put("0", "PRIMARY"); + localShards.get("index1").put("1", "SEARCH_REPLICA"); + localShards.put("index2", new HashMap<>()); + localShards.get("index2").put("0", "PRIMARY"); + state.setLocalShards(localShards); + + assertThat(state.hasShardWithRole("index1", "0", "PRIMARY")).isTrue(); + assertThat(state.hasShardWithRole("index1", "1", "SEARCH_REPLICA")).isTrue(); + assertThat(state.hasShardWithRole("index2", "0", "PRIMARY")).isTrue(); + assertThat(state.hasShardWithRole("index2", "0", "SEARCH_REPLICA")).isFalse(); + assertThat(state.hasShardWithRole("index3", "0", "PRIMARY")).isFalse(); + } + + // ========== Helper methods ========== + + private SearchUnitGoalState createGoalState(String indexName, String shardId, String role) { + SearchUnitGoalState state = new SearchUnitGoalState(); + Map> localShards = new HashMap<>(); + localShards.put(indexName, new HashMap<>()); + localShards.get(indexName).put(shardId, role); + state.setLocalShards(localShards); + return state; + } +} diff --git a/src/test/java/io/clustercontroller/orchestration/ImmediateOrchestrationStrategyTest.java b/src/test/java/io/clustercontroller/orchestration/ImmediateOrchestrationStrategyTest.java index ec7b2d08..40f24ae9 100644 --- a/src/test/java/io/clustercontroller/orchestration/ImmediateOrchestrationStrategyTest.java +++ b/src/test/java/io/clustercontroller/orchestration/ImmediateOrchestrationStrategyTest.java @@ -46,12 +46,15 @@ void testOrchestrateWithSingleIndexSingleShard() throws Exception { planned.setIngestSUs(Arrays.asList("node1")); planned.setSearchSUs(Arrays.asList("node2", "node3")); - SearchUnitGoalState existingGoalState = new SearchUnitGoalState(); - existingGoalState.setLocalShards(new HashMap<>()); - + // Return a fresh empty goal state for each node (simulating separate states per node) when(metadataStore.getAllIndexConfigs(clusterId)).thenReturn(Arrays.asList(indexConfig)); when(metadataStore.getPlannedAllocation(clusterId, indexName, "0")).thenReturn(planned); - when(metadataStore.getSearchUnitGoalState(eq(clusterId), anyString())).thenReturn(existingGoalState); + when(metadataStore.getSearchUnitGoalState(eq(clusterId), anyString())) + .thenAnswer(invocation -> { + SearchUnitGoalState state = new SearchUnitGoalState(); + state.setLocalShards(new HashMap<>()); + return state; + }); // When strategy.orchestrate(clusterId); @@ -86,14 +89,22 @@ void testOrchestrateWithMultipleIndexesAndShards() throws Exception { planned2_0.setIngestSUs(Arrays.asList("node1")); planned2_0.setSearchSUs(Arrays.asList("node2", "node3")); - SearchUnitGoalState existingGoalState = new SearchUnitGoalState(); - existingGoalState.setLocalShards(new HashMap<>()); + // Track goal states per node to simulate real behavior where each node has its own state + Map nodeGoalStates = new HashMap<>(); when(metadataStore.getAllIndexConfigs(clusterId)).thenReturn(Arrays.asList(index1, index2)); when(metadataStore.getPlannedAllocation(clusterId, "index1", "0")).thenReturn(planned1_0); when(metadataStore.getPlannedAllocation(clusterId, "index1", "1")).thenReturn(planned1_1); when(metadataStore.getPlannedAllocation(clusterId, "index2", "0")).thenReturn(planned2_0); - when(metadataStore.getSearchUnitGoalState(eq(clusterId), anyString())).thenReturn(existingGoalState); + when(metadataStore.getSearchUnitGoalState(eq(clusterId), anyString())) + .thenAnswer(invocation -> { + String nodeId = invocation.getArgument(1); + return nodeGoalStates.computeIfAbsent(nodeId, k -> { + SearchUnitGoalState state = new SearchUnitGoalState(); + state.setLocalShards(new HashMap<>()); + return state; + }); + }); // When strategy.orchestrate(clusterId); @@ -107,6 +118,7 @@ void testOrchestrateWithMultipleIndexesAndShards() throws Exception { verify(metadataStore).getPlannedAllocation(clusterId, "index2", "0"); // Verify goal states are set for all nodes (7 total calls due to duplicate nodes across shards) + // Each node/index/shard combination is unique, so all 7 updates should happen verify(metadataStore, times(7)).setSearchUnitGoalState(eq(clusterId), anyString(), any(SearchUnitGoalState.class)); } @@ -382,6 +394,115 @@ void testOrchestrateWithNullIngestAndSearchSUs() throws Exception { verify(metadataStore, never()).setSearchUnitGoalState(anyString(), anyString(), any(SearchUnitGoalState.class)); } + @Test + void testOrchestrateUpdatesWhenRoleChanges() throws Exception { + // Given: A node has shard 0 as REPLICA, but planned allocation says it should be PRIMARY + String clusterId = "test-cluster"; + String indexName = "test-index"; + + Index indexConfig = createIndex(indexName, 1); + + ShardAllocation planned = new ShardAllocation(); + planned.setIngestSUs(Arrays.asList("node1")); // node1 should be PRIMARY + planned.setSearchSUs(Arrays.asList()); + + // Existing goal state has node1 as REPLICA for this shard + SearchUnitGoalState existingGoalState = new SearchUnitGoalState(); + Map> localShards = new HashMap<>(); + localShards.put(indexName, new HashMap<>()); + localShards.get(indexName).put("0", "SEARCH_REPLICA"); // Wrong role! + existingGoalState.setLocalShards(localShards); + + when(metadataStore.getAllIndexConfigs(clusterId)).thenReturn(Arrays.asList(indexConfig)); + when(metadataStore.getPlannedAllocation(clusterId, indexName, "0")).thenReturn(planned); + when(metadataStore.getSearchUnitGoalState(clusterId, "node1")).thenReturn(existingGoalState); + + // When + strategy.orchestrate(clusterId); + + // Then - Should update because the role changed from SEARCH_REPLICA to PRIMARY + verify(metadataStore).setSearchUnitGoalState(eq(clusterId), eq("node1"), any(SearchUnitGoalState.class)); + } + + @Test + void testOrchestrateUpdatesWhenNewShardAdded() throws Exception { + // Given: A node has shard 0 but now needs shard 1 as well + String clusterId = "test-cluster"; + String indexName = "test-index"; + + Index indexConfig = createIndex(indexName, 2); + + ShardAllocation planned0 = new ShardAllocation(); + planned0.setIngestSUs(Arrays.asList("node1")); + planned0.setSearchSUs(Arrays.asList()); + + ShardAllocation planned1 = new ShardAllocation(); + planned1.setIngestSUs(Arrays.asList("node1")); // Same node needs shard 1 now + planned1.setSearchSUs(Arrays.asList()); + + // Track goal state - starts with only shard 0 + SearchUnitGoalState existingGoalState = new SearchUnitGoalState(); + Map> localShards = new HashMap<>(); + localShards.put(indexName, new HashMap<>()); + localShards.get(indexName).put("0", "PRIMARY"); + existingGoalState.setLocalShards(localShards); + + when(metadataStore.getAllIndexConfigs(clusterId)).thenReturn(Arrays.asList(indexConfig)); + when(metadataStore.getPlannedAllocation(clusterId, indexName, "0")).thenReturn(planned0); + when(metadataStore.getPlannedAllocation(clusterId, indexName, "1")).thenReturn(planned1); + when(metadataStore.getSearchUnitGoalState(clusterId, "node1")).thenReturn(existingGoalState); + + // When + strategy.orchestrate(clusterId); + + // Then + // Shard 0 should be skipped (already has PRIMARY role) + // Shard 1 should be updated (new shard) + verify(metadataStore, times(1)).setSearchUnitGoalState(eq(clusterId), eq("node1"), any(SearchUnitGoalState.class)); + } + + @Test + void testOrchestrateSkipsUnchangedGoalStates() throws Exception { + // Given + String clusterId = "test-cluster"; + String indexName = "test-index"; + int numberOfShards = 1; + + Index indexConfig = createIndex(indexName, numberOfShards); + + ShardAllocation planned = new ShardAllocation(); + planned.setIngestSUs(Arrays.asList("node1")); + planned.setSearchSUs(Arrays.asList("node2")); + + // Create existing goal states that already have the shard with correct role + SearchUnitGoalState existingPrimaryGoalState = new SearchUnitGoalState(); + Map> primaryLocalShards = new HashMap<>(); + primaryLocalShards.put(indexName, new HashMap<>()); + primaryLocalShards.get(indexName).put("0", "PRIMARY"); + existingPrimaryGoalState.setLocalShards(primaryLocalShards); + + SearchUnitGoalState existingReplicaGoalState = new SearchUnitGoalState(); + Map> replicaLocalShards = new HashMap<>(); + replicaLocalShards.put(indexName, new HashMap<>()); + replicaLocalShards.get(indexName).put("0", "SEARCH_REPLICA"); + existingReplicaGoalState.setLocalShards(replicaLocalShards); + + when(metadataStore.getAllIndexConfigs(clusterId)).thenReturn(Arrays.asList(indexConfig)); + when(metadataStore.getPlannedAllocation(clusterId, indexName, "0")).thenReturn(planned); + when(metadataStore.getSearchUnitGoalState(clusterId, "node1")).thenReturn(existingPrimaryGoalState); + when(metadataStore.getSearchUnitGoalState(clusterId, "node2")).thenReturn(existingReplicaGoalState); + + // When + strategy.orchestrate(clusterId); + + // Then + verify(metadataStore).getAllIndexConfigs(clusterId); + verify(metadataStore).getPlannedAllocation(clusterId, indexName, "0"); + + // Verify goal states are NOT updated since they already have the correct shard+role + verify(metadataStore, never()).setSearchUnitGoalState(eq(clusterId), anyString(), any(SearchUnitGoalState.class)); + } + // Helper method to create Index with initialized settings private Index createIndex(String indexName, int numberOfShards) { Index index = new Index(); diff --git a/src/test/java/io/clustercontroller/orchestration/RollingUpdateOrchestrationStrategyTest.java b/src/test/java/io/clustercontroller/orchestration/RollingUpdateOrchestrationStrategyTest.java index b832b46c..dfe85019 100644 --- a/src/test/java/io/clustercontroller/orchestration/RollingUpdateOrchestrationStrategyTest.java +++ b/src/test/java/io/clustercontroller/orchestration/RollingUpdateOrchestrationStrategyTest.java @@ -18,9 +18,12 @@ import java.util.List; import java.util.Map; +import static org.junit.jupiter.api.Assertions.*; import static org.mockito.ArgumentMatchers.*; import static org.mockito.Mockito.*; +import org.mockito.ArgumentCaptor; + @ExtendWith(MockitoExtension.class) class RollingUpdateOrchestrationStrategyTest { @@ -644,7 +647,7 @@ void testOrchestrateThrowsExceptionWhenGetSearchUnitGoalStateFails() throws Exce // Then - verify the method was called and exception was handled verify(metadataStore).getAllIndexConfigs(clusterId); verify(metadataStore).getPlannedAllocation(clusterId, indexName, "0"); - verify(metadataStore, times(2)).getSearchUnitGoalState(clusterId, "node1"); // Called in hasGoalStateUpdated and updateNodeGoalState + verify(metadataStore, times(3)).getSearchUnitGoalState(clusterId, "node1"); // Called in hasGoalStateUpdated, queueNodeGoalStateUpdate, and flushPendingGoalStates // The orchestration should complete without throwing (resilient behavior) } @@ -1041,6 +1044,11 @@ private SearchUnitGoalState createGoalStateWithShard(String indexName, String sh } private SearchUnitActualState createActualStateWithShard(String indexName, String shardId, String role) { + return createActualStateWithShardAndState(indexName, shardId, role, io.clustercontroller.enums.ShardState.STARTED); + } + + private SearchUnitActualState createActualStateWithShardAndState(String indexName, String shardId, String role, + io.clustercontroller.enums.ShardState state) { SearchUnitActualState actualState = new SearchUnitActualState(); Map> nodeRouting = new HashMap<>(); SearchUnitActualState.ShardRoutingInfo shardInfo = new SearchUnitActualState.ShardRoutingInfo(); @@ -1048,12 +1056,120 @@ private SearchUnitActualState createActualStateWithShard(String indexName, Strin shardInfo.setShardId(Integer.parseInt(shardId)); // Convert role to lowercase format matching worker output ("primary", "search_replica", "replica") shardInfo.setRole("PRIMARY".equals(role) ? "primary" : "replica"); - shardInfo.setState(io.clustercontroller.enums.ShardState.STARTED); + shardInfo.setState(state); nodeRouting.put(indexName, List.of(shardInfo)); actualState.setNodeRouting(nodeRouting); return actualState; } + // ========== BATCHING TESTS ========== + + @Test + void testBatchingMultipleShardsIntoSingleWrite() throws Exception { + // Given: A node needs to be updated for multiple shards + // This test verifies that all shard assignments are batched into ONE etcd write + String clusterId = "test-cluster"; + String indexName = "test-index"; + String nodeId = "node1"; + int numberOfShards = 5; // 5 shards all assigned to same node + + Index indexConfig = createIndex(indexName, numberOfShards); + when(metadataStore.getAllIndexConfigs(clusterId)).thenReturn(List.of(indexConfig)); + when(metadataStore.getAllNodesWithGoalStates(clusterId)).thenReturn(List.of()); + + // All 5 shards are assigned to the same primary node + for (int i = 0; i < numberOfShards; i++) { + ShardAllocation planned = new ShardAllocation(); + planned.setIngestSUs(List.of(nodeId)); + planned.setSearchSUs(List.of()); + when(metadataStore.getPlannedAllocation(clusterId, indexName, String.valueOf(i))) + .thenReturn(planned); + } + + // Node has no current goal state (new node) + when(metadataStore.getSearchUnitGoalState(clusterId, nodeId)).thenReturn(null); + + // When + strategy.orchestrate(clusterId); + + // Then: Should write ONCE with all 5 shards (not 5 separate writes) + ArgumentCaptor goalStateCaptor = ArgumentCaptor.forClass(SearchUnitGoalState.class); + verify(metadataStore, times(1)).setSearchUnitGoalState(eq(clusterId), eq(nodeId), goalStateCaptor.capture()); + + // Verify the single write contains all 5 shards + SearchUnitGoalState writtenGoalState = goalStateCaptor.getValue(); + assertNotNull(writtenGoalState.getLocalShards().get(indexName)); + assertEquals(5, writtenGoalState.getLocalShards().get(indexName).size(), + "All 5 shards should be in single goal state write"); + } + + @Test + void testBatchingSkipsWriteWhenGoalStateUnchanged() throws Exception { + // Given: A node already has all its shards in goal state + String clusterId = "test-cluster"; + String indexName = "test-index"; + String nodeId = "node1"; + + Index indexConfig = createIndex(indexName, 1); + when(metadataStore.getAllIndexConfigs(clusterId)).thenReturn(List.of(indexConfig)); + when(metadataStore.getAllNodesWithGoalStates(clusterId)).thenReturn(List.of()); + + ShardAllocation planned = new ShardAllocation(); + planned.setIngestSUs(List.of(nodeId)); + planned.setSearchSUs(List.of()); + when(metadataStore.getPlannedAllocation(clusterId, indexName, "0")).thenReturn(planned); + + // Node already has this shard in its goal state + SearchUnitGoalState existingGoalState = createGoalStateWithShard(indexName, "0", "PRIMARY"); + when(metadataStore.getSearchUnitGoalState(clusterId, nodeId)).thenReturn(existingGoalState); + + // When + strategy.orchestrate(clusterId); + + // Then: Should NOT write to etcd (goal state unchanged) + verify(metadataStore, never()).setSearchUnitGoalState(eq(clusterId), eq(nodeId), any()); + } + + @Test + void testBatchingMultipleIndicesForSameNode() throws Exception { + // Given: A node is assigned shards from multiple indices + String clusterId = "test-cluster"; + String nodeId = "node1"; + String index1 = "index-1"; + String index2 = "index-2"; + + Index indexConfig1 = createIndex(index1, 2); + Index indexConfig2 = createIndex(index2, 2); + when(metadataStore.getAllIndexConfigs(clusterId)).thenReturn(List.of(indexConfig1, indexConfig2)); + when(metadataStore.getAllNodesWithGoalStates(clusterId)).thenReturn(List.of()); + + // Node is primary for both indices + for (String indexName : List.of(index1, index2)) { + for (int i = 0; i < 2; i++) { + ShardAllocation planned = new ShardAllocation(); + planned.setIngestSUs(List.of(nodeId)); + planned.setSearchSUs(List.of()); + when(metadataStore.getPlannedAllocation(clusterId, indexName, String.valueOf(i))) + .thenReturn(planned); + } + } + + when(metadataStore.getSearchUnitGoalState(clusterId, nodeId)).thenReturn(null); + + // When + strategy.orchestrate(clusterId); + + // Then: Should write ONCE with shards from BOTH indices + ArgumentCaptor goalStateCaptor = ArgumentCaptor.forClass(SearchUnitGoalState.class); + verify(metadataStore, times(1)).setSearchUnitGoalState(eq(clusterId), eq(nodeId), goalStateCaptor.capture()); + + SearchUnitGoalState writtenGoalState = goalStateCaptor.getValue(); + assertNotNull(writtenGoalState.getLocalShards().get(index1), "Should have index-1 shards"); + assertNotNull(writtenGoalState.getLocalShards().get(index2), "Should have index-2 shards"); + assertEquals(2, writtenGoalState.getLocalShards().get(index1).size()); + assertEquals(2, writtenGoalState.getLocalShards().get(index2).size()); + } + // Helper method to create Index with initialized settings private Index createIndex(String indexName, int numberOfShards) { Index index = new Index(); @@ -1065,4 +1181,167 @@ private Index createIndex(String indexName, int numberOfShards) { return index; } + + // ========== SHARD STATE CONVERGENCE TESTS ========== + + @Test + void testInitializingShardsNotConsideredConverged() throws Exception { + // Given: Nodes have goal states but shards are still INITIALIZING (not STARTED) + // This tests that the rolling update waits for STARTED state before proceeding + String clusterId = "test-cluster"; + String indexName = "test-index"; + String shardId = "0"; + + Index indexConfig = createIndex(indexName, 1); + + // Create planned allocation with 10 replica nodes + ShardAllocation planned = new ShardAllocation(); + planned.setIngestSUs(List.of("primary1")); + planned.setSearchSUs(List.of("node1", "node2", "node3", "node4", "node5", + "node6", "node7", "node8", "node9", "node10")); + + // node1 and node2 have goal states (already updated) + SearchUnitGoalState goalState1 = createGoalStateWithShard(indexName, shardId, "SEARCH_REPLICA"); + SearchUnitGoalState goalState2 = createGoalStateWithShard(indexName, shardId, "SEARCH_REPLICA"); + + // But their actual states show INITIALIZING (not converged yet) + SearchUnitActualState initializingState1 = createActualStateWithShardAndState( + indexName, shardId, "SEARCH_REPLICA", io.clustercontroller.enums.ShardState.INITIALIZING); + SearchUnitActualState initializingState2 = createActualStateWithShardAndState( + indexName, shardId, "SEARCH_REPLICA", io.clustercontroller.enums.ShardState.INITIALIZING); + + when(metadataStore.getAllIndexConfigs(clusterId)).thenReturn(List.of(indexConfig)); + when(metadataStore.getPlannedAllocation(clusterId, indexName, shardId)).thenReturn(planned); + when(metadataStore.getAllNodesWithGoalStates(clusterId)).thenReturn(List.of()); + + // Primary is converged (already has goal state and is STARTED) + when(metadataStore.getSearchUnitGoalState(clusterId, "primary1")) + .thenReturn(createGoalStateWithShard(indexName, shardId, "PRIMARY")); + when(metadataStore.getSearchUnitActualState(clusterId, "primary1")) + .thenReturn(createActualStateWithShard(indexName, shardId, "PRIMARY")); + + // node1, node2 have goal states but are INITIALIZING + when(metadataStore.getSearchUnitGoalState(clusterId, "node1")).thenReturn(goalState1); + when(metadataStore.getSearchUnitGoalState(clusterId, "node2")).thenReturn(goalState2); + when(metadataStore.getSearchUnitActualState(clusterId, "node1")).thenReturn(initializingState1); + when(metadataStore.getSearchUnitActualState(clusterId, "node2")).thenReturn(initializingState2); + + // Rest have no goal states yet + for (int i = 3; i <= 10; i++) { + when(metadataStore.getSearchUnitGoalState(clusterId, "node" + i)).thenReturn(null); + } + + // When + strategy.orchestrate(clusterId); + + // Then: With 10 replicas and 20% limit, max 2 can be in transit + // node1 and node2 are already in transit (INITIALIZING), so NO new replica nodes should be updated + // Primary already has matching goal state, so it's skipped (unchanged) + // Result: 0 writes because all in-scope nodes either already have goal states or are blocked by transit limit + verify(metadataStore, never()).setSearchUnitGoalState(eq(clusterId), anyString(), any(SearchUnitGoalState.class)); + } + + @Test + void testStartedShardsAllowNextBatch() throws Exception { + // Given: Some nodes have STARTED shards (converged), allowing next batch to proceed + String clusterId = "test-cluster"; + String indexName = "test-index"; + String shardId = "0"; + + Index indexConfig = createIndex(indexName, 1); + + // Create planned allocation with 10 replica nodes + ShardAllocation planned = new ShardAllocation(); + planned.setIngestSUs(List.of("primary1")); + planned.setSearchSUs(List.of("node1", "node2", "node3", "node4", "node5", + "node6", "node7", "node8", "node9", "node10")); + + // node1 and node2 have goal states AND are STARTED (converged) + SearchUnitGoalState goalState1 = createGoalStateWithShard(indexName, shardId, "SEARCH_REPLICA"); + SearchUnitGoalState goalState2 = createGoalStateWithShard(indexName, shardId, "SEARCH_REPLICA"); + SearchUnitActualState startedState1 = createActualStateWithShard(indexName, shardId, "SEARCH_REPLICA"); + SearchUnitActualState startedState2 = createActualStateWithShard(indexName, shardId, "SEARCH_REPLICA"); + + when(metadataStore.getAllIndexConfigs(clusterId)).thenReturn(List.of(indexConfig)); + when(metadataStore.getPlannedAllocation(clusterId, indexName, shardId)).thenReturn(planned); + when(metadataStore.getAllNodesWithGoalStates(clusterId)).thenReturn(List.of()); + + // Primary is converged + when(metadataStore.getSearchUnitGoalState(clusterId, "primary1")) + .thenReturn(createGoalStateWithShard(indexName, shardId, "PRIMARY")); + when(metadataStore.getSearchUnitActualState(clusterId, "primary1")) + .thenReturn(createActualStateWithShard(indexName, shardId, "PRIMARY")); + + // node1, node2 are STARTED (converged - not in transit) + when(metadataStore.getSearchUnitGoalState(clusterId, "node1")).thenReturn(goalState1); + when(metadataStore.getSearchUnitGoalState(clusterId, "node2")).thenReturn(goalState2); + when(metadataStore.getSearchUnitActualState(clusterId, "node1")).thenReturn(startedState1); + when(metadataStore.getSearchUnitActualState(clusterId, "node2")).thenReturn(startedState2); + + // Rest have no goal states yet + for (int i = 3; i <= 10; i++) { + when(metadataStore.getSearchUnitGoalState(clusterId, "node" + i)).thenReturn(null); + } + + // When + strategy.orchestrate(clusterId); + + // Then: node1 and node2 are converged (0 in transit) + // With 20% limit of 10 nodes = 2 new nodes can be updated + // So 2 new replica nodes should be updated (node3 and node4) + // Total writes: 2 replicas (batched writes) + verify(metadataStore, times(2)).setSearchUnitGoalState(eq(clusterId), anyString(), any(SearchUnitGoalState.class)); + } + + @Test + void testMixedInitializingAndStartedShards() throws Exception { + // Given: Some nodes STARTED, some INITIALIZING - only INITIALIZING count as transit + String clusterId = "test-cluster"; + String indexName = "test-index"; + String shardId = "0"; + + Index indexConfig = createIndex(indexName, 1); + + // 10 replica nodes + ShardAllocation planned = new ShardAllocation(); + planned.setIngestSUs(List.of("primary1")); + planned.setSearchSUs(List.of("node1", "node2", "node3", "node4", "node5", + "node6", "node7", "node8", "node9", "node10")); + + when(metadataStore.getAllIndexConfigs(clusterId)).thenReturn(List.of(indexConfig)); + when(metadataStore.getPlannedAllocation(clusterId, indexName, shardId)).thenReturn(planned); + when(metadataStore.getAllNodesWithGoalStates(clusterId)).thenReturn(List.of()); + + // Primary is converged + when(metadataStore.getSearchUnitGoalState(clusterId, "primary1")) + .thenReturn(createGoalStateWithShard(indexName, shardId, "PRIMARY")); + when(metadataStore.getSearchUnitActualState(clusterId, "primary1")) + .thenReturn(createActualStateWithShard(indexName, shardId, "PRIMARY")); + + // node1: STARTED (converged) + when(metadataStore.getSearchUnitGoalState(clusterId, "node1")) + .thenReturn(createGoalStateWithShard(indexName, shardId, "SEARCH_REPLICA")); + when(metadataStore.getSearchUnitActualState(clusterId, "node1")) + .thenReturn(createActualStateWithShard(indexName, shardId, "SEARCH_REPLICA")); + + // node2: INITIALIZING (in transit - counts toward limit) + when(metadataStore.getSearchUnitGoalState(clusterId, "node2")) + .thenReturn(createGoalStateWithShard(indexName, shardId, "SEARCH_REPLICA")); + when(metadataStore.getSearchUnitActualState(clusterId, "node2")) + .thenReturn(createActualStateWithShardAndState(indexName, shardId, "SEARCH_REPLICA", + io.clustercontroller.enums.ShardState.INITIALIZING)); + + // Rest have no goal states yet + for (int i = 3; i <= 10; i++) { + when(metadataStore.getSearchUnitGoalState(clusterId, "node" + i)).thenReturn(null); + } + + // When + strategy.orchestrate(clusterId); + + // Then: 1 node is in transit (node2 INITIALIZING) + // With 20% limit of 10 = 2 max in transit, so 1 more can be added + // So 1 new node should be updated + verify(metadataStore, times(1)).setSearchUnitGoalState(eq(clusterId), anyString(), any(SearchUnitGoalState.class)); + } } diff --git a/src/test/java/io/clustercontroller/store/LeaderElectionTest.java b/src/test/java/io/clustercontroller/store/LeaderElectionTest.java index fa542b73..ef8c3f1c 100644 --- a/src/test/java/io/clustercontroller/store/LeaderElectionTest.java +++ b/src/test/java/io/clustercontroller/store/LeaderElectionTest.java @@ -143,12 +143,6 @@ void testLeaseKeepAliveError() throws Exception { when(leaseGrantResponse.getID()).thenReturn(12345L); when(leaseClient.grant(anyLong())).thenReturn(CompletableFuture.completedFuture(leaseGrantResponse)); - // Setup failed election campaign due to keep alive error - CompletableFuture failedCampaign = new CompletableFuture<>(); - failedCampaign.completeExceptionally(new RuntimeException("Campaign failed due to keep alive error")); - when(electionClient.campaign(any(ByteSequence.class), anyLong(), any(ByteSequence.class))) - .thenReturn(failedCampaign); - // Setup keep alive mock to simulate error callback doAnswer(invocation -> { StreamObserver observer = invocation.getArgument(1);