From 8bae6e8353b46ac7400336f20d3e0f0f93651c4d Mon Sep 17 00:00:00 2001 From: Hao Date: Mon, 22 Dec 2025 16:00:40 -0800 Subject: [PATCH] Add metrics in Actual States Allocator and Discovery --- .../ClusterControllerApplication.java | 8 +- .../allocation/ActualAllocationUpdater.java | 45 +++++++--- .../discovery/Discovery.java | 19 ++-- .../metrics/MetricsConstants.java | 2 + .../ActualAllocationUpdaterCleanupTest.java | 90 ++++++++++++++++--- .../discovery/DiscoveryTest.java | 81 ++++++++++++++--- 6 files changed, 196 insertions(+), 49 deletions(-) diff --git a/src/main/java/io/clustercontroller/ClusterControllerApplication.java b/src/main/java/io/clustercontroller/ClusterControllerApplication.java index cb40d973..89423d43 100644 --- a/src/main/java/io/clustercontroller/ClusterControllerApplication.java +++ b/src/main/java/io/clustercontroller/ClusterControllerApplication.java @@ -114,9 +114,9 @@ public ShardAllocator shardAllocator(MetadataStore metadataStore, MetricsProvide } @Bean - public ActualAllocationUpdater actualAllocationUpdater(MetadataStore metadataStore) { + public ActualAllocationUpdater actualAllocationUpdater(MetadataStore metadataStore, MetricsProvider metricsProvider) { log.info("Initializing ActualAllocationUpdater"); - return new ActualAllocationUpdater(metadataStore); + return new ActualAllocationUpdater(metadataStore, metricsProvider); } /** @@ -132,9 +132,9 @@ public GoalStateOrchestrator goalStateOrchestrator(MetadataStore metadataStore, * Discovery bean for node discovery */ @Bean - public Discovery discovery(MetadataStore metadataStore) { + public Discovery discovery(MetadataStore metadataStore, MetricsProvider metricsProvider) { log.info("Initializing Discovery for multi-cluster support"); - return new Discovery(metadataStore); + return new Discovery(metadataStore, metricsProvider); } /** diff --git a/src/main/java/io/clustercontroller/allocation/ActualAllocationUpdater.java b/src/main/java/io/clustercontroller/allocation/ActualAllocationUpdater.java index ca6fe28f..829e7d9f 100644 --- a/src/main/java/io/clustercontroller/allocation/ActualAllocationUpdater.java +++ b/src/main/java/io/clustercontroller/allocation/ActualAllocationUpdater.java @@ -1,22 +1,38 @@ package io.clustercontroller.allocation; import io.clustercontroller.enums.ShardState; -import io.clustercontroller.models.*; +import io.clustercontroller.metrics.MetricsProvider; +import io.clustercontroller.models.Alias; +import io.clustercontroller.models.CoordinatorGoalState; +import io.clustercontroller.models.Index; +import io.clustercontroller.models.IndexSettings; +import io.clustercontroller.models.SearchUnit; +import io.clustercontroller.models.SearchUnitActualState; +import io.clustercontroller.models.SearchUnitGoalState; +import io.clustercontroller.models.ShardAllocation; import io.clustercontroller.store.MetadataStore; import lombok.extern.slf4j.Slf4j; import io.clustercontroller.config.Constants; -import java.util.*; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static io.clustercontroller.metrics.MetricsConstants.UPDATE_ACTUAL_ALLOCATION_FAILURES_METRIC_NAME; +import static io.clustercontroller.metrics.MetricsUtils.buildMetricsTags; /** * ActualAllocationUpdater - Responsible for aggregating search unit actual states into index shard actual allocations - * + *

* This task: * 1. Reads actual states from /search-units//actual-state * 2. Aggregates this information by index and shard * 3. Updates actual allocations at /indices///actual-allocation * 4. Provides coordinator nodes with a goal-state view of remote_shards used for routing/monitoring - * + *

* Coordinator goal state specifics: * - For each index and shard, we compute a list of nodes with their role (primary/replica) flag: List>. * - A node is included only when it is converged for that index/shard. @@ -24,7 +40,7 @@ * (b) the node's actual state reports that shard in STARTED state. * - Primaries come from the current actual ingest set; replicas come from the current actual search set, * and both are filtered by convergence before inclusion. - * + *

* Example: * - Index "logs-2025" has 2 shards. * - Shard 0: su-a (primary, STARTED, requested in goal) and su-b (replica, STARTED, requested in goal) @@ -36,10 +52,12 @@ public class ActualAllocationUpdater { private final MetadataStore metadataStore; + private final MetricsProvider metricsProvider; private static final long ONE_MINUTE_MS = 60 * 1000; - public ActualAllocationUpdater(MetadataStore metadataStore) { + public ActualAllocationUpdater(MetadataStore metadataStore, MetricsProvider metricsProvider) { this.metadataStore = metadataStore; + this.metricsProvider = metricsProvider; } /** @@ -77,7 +95,7 @@ public void updateActualAllocations(String clusterId) throws Exception { * Collect actual allocations from all search units * Returns: Map>> */ - private Map>> collectActualAllocations(String clusterId, List searchUnits) throws Exception { + private Map>> collectActualAllocations(String clusterId, List searchUnits) { Map>> actualAllocations = new HashMap<>(); for (SearchUnit searchUnit : searchUnits) { @@ -152,7 +170,7 @@ private Map>> collectActualAllocations(String cl /** * Update actual allocation records based on collected data */ - private int updateActualAllocationRecords(String clusterId, Map>> actualAllocations) throws Exception { + private int updateActualAllocationRecords(String clusterId, Map>> actualAllocations) { int totalUpdates = 0; for (Map.Entry>> indexEntry : actualAllocations.entrySet()) { @@ -169,6 +187,10 @@ private int updateActualAllocationRecords(String clusterId, Map getAllIndicesWithActualAllocations(String clusterId) throws /** * Determines if a search unit is a coordinator node * Coordinator nodes belong to a separate cluster and are not part of shard allocation - * * Current logic: role="coordinator" OR name starts with "coordinator" * TODO: Use cluster_name information to properly identify coordinator nodes * instead of relying on naming conventions. Coordinator nodes should have @@ -368,11 +389,7 @@ private boolean isCoordinatorNode(SearchUnit unit) { // Check if name starts with "coordinator" prefix (for data nodes acting as coordinators) String name = unit.getName() != null ? unit.getName().toLowerCase() : ""; - if (name.startsWith("coordinator")) { - return true; - } - - return false; + return name.startsWith("coordinator"); } /** diff --git a/src/main/java/io/clustercontroller/discovery/Discovery.java b/src/main/java/io/clustercontroller/discovery/Discovery.java index 39dbb934..8d1ac233 100644 --- a/src/main/java/io/clustercontroller/discovery/Discovery.java +++ b/src/main/java/io/clustercontroller/discovery/Discovery.java @@ -2,6 +2,7 @@ import io.clustercontroller.config.Constants; import io.clustercontroller.enums.HealthState; +import io.clustercontroller.metrics.MetricsProvider; import io.clustercontroller.models.NodeAttributes; import io.clustercontroller.models.SearchUnit; import io.clustercontroller.models.SearchUnitActualState; @@ -13,10 +14,13 @@ import java.util.List; import java.util.Map; +import static io.clustercontroller.metrics.MetricsConstants.CLUSTER_ID_TAG; +import static io.clustercontroller.metrics.MetricsConstants.DISCOVERY_CLEANED_STALE_SEARCH_UNITS_COUNT_METRIC_NAME; + /** * Handles cluster topology discovery. * Internal component used by TaskManager. - * + *

* Note: This is a shared component across all clusters in multi-cluster mode. * Methods accept clusterName as a parameter instead of storing it as a field. */ @@ -24,13 +28,15 @@ public class Discovery { private final MetadataStore metadataStore; + private final MetricsProvider metricsProvider; - public Discovery(MetadataStore metadataStore) { + public Discovery(MetadataStore metadataStore, MetricsProvider metricsProvider) { this.metadataStore = metadataStore; + this.metricsProvider = metricsProvider; } - public void discoverSearchUnits(String clusterName) throws Exception { + public void discoverSearchUnits(String clusterName) { log.info("Discovery - Starting search unit discovery process for cluster: {}", clusterName); // Discover and update search units from Etcd actual-states @@ -75,7 +81,6 @@ private void processAllSearchUnits(String clusterName) { */ private void discoverSearchUnitsFromEtcd(String clusterName) { try { - log.info("Discovery - Discovering search units from Etcd..."); // fetch search units from actual-state paths List etcdSearchUnits = fetchSearchUnitsFromEtcd(clusterName); log.info("Discovery - Found {} search units from Etcd", etcdSearchUnits.size()); @@ -175,12 +180,12 @@ private SearchUnit convertActualStateToSearchUnit(SearchUnitActualState actualSt * Clean up search units with missing or stale actual state timestamp (older than configured timeout) */ private void cleanupStaleSearchUnits(String clusterName) { + int deletedCount = 0; try { log.info("Discovery - Starting cleanup of stale search units..."); // Get all existing search units from metadata store List allSearchUnits = metadataStore.getAllSearchUnits(clusterName); - int deletedCount = 0; for (SearchUnit searchUnit : allSearchUnits) { String unitName = searchUnit.getName(); @@ -221,6 +226,10 @@ private void cleanupStaleSearchUnits(String clusterName) { } catch (Exception e) { log.error("Discovery - Failed to cleanup stale search units: {}", e.getMessage(), e); + } finally { + metricsProvider + .counter(DISCOVERY_CLEANED_STALE_SEARCH_UNITS_COUNT_METRIC_NAME, Map.of(CLUSTER_ID_TAG, clusterName)) + .increment(deletedCount); } } diff --git a/src/main/java/io/clustercontroller/metrics/MetricsConstants.java b/src/main/java/io/clustercontroller/metrics/MetricsConstants.java index 243885c6..a394c4ba 100644 --- a/src/main/java/io/clustercontroller/metrics/MetricsConstants.java +++ b/src/main/java/io/clustercontroller/metrics/MetricsConstants.java @@ -8,6 +8,8 @@ public class MetricsConstants { public final static String ROLLING_UPDATE_TRANSIT_NODES_PERCENTAGE_METRIC_NAME = "rolling_update_transit_nodes_percentage"; public final static String PLANNED_INGEST_SU_ALLOCATION_METRIC_NAME = "planned_ingest_su_allocation"; public final static String PLANNED_SEARCH_SU_ALLOCATION_METRIC_NAME = "planned_search_su_allocation"; + public final static String UPDATE_ACTUAL_ALLOCATION_FAILURES_METRIC_NAME = "update_actual_allocation_failures_count"; + public final static String DISCOVERY_CLEANED_STALE_SEARCH_UNITS_COUNT_METRIC_NAME = "discovery_cleaned_stale_search_units_count"; public final static String CLUSTER_ID_TAG = "clusterId"; public final static String INDEX_NAME_TAG = "indexName"; public final static String SHARD_ID_TAG = "shardId"; diff --git a/src/test/java/io/clustercontroller/allocation/ActualAllocationUpdaterCleanupTest.java b/src/test/java/io/clustercontroller/allocation/ActualAllocationUpdaterCleanupTest.java index ef7b6c82..b18ba2fe 100644 --- a/src/test/java/io/clustercontroller/allocation/ActualAllocationUpdaterCleanupTest.java +++ b/src/test/java/io/clustercontroller/allocation/ActualAllocationUpdaterCleanupTest.java @@ -1,18 +1,19 @@ package io.clustercontroller.allocation; import io.clustercontroller.enums.ShardState; +import io.clustercontroller.metrics.MetricsProvider; import io.clustercontroller.models.Index; import io.clustercontroller.models.IndexSettings; import io.clustercontroller.models.SearchUnit; import io.clustercontroller.models.SearchUnitActualState; import io.clustercontroller.models.ShardAllocation; import io.clustercontroller.store.MetadataStore; +import io.micrometer.core.instrument.Counter; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.Mock; import org.mockito.MockitoAnnotations; -import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.List; @@ -29,12 +30,20 @@ class ActualAllocationUpdaterCleanupTest { @Mock private MetadataStore metadataStore; + @Mock + private MetricsProvider metricsProvider; + + @Mock + private Counter mockCounter; + private ActualAllocationUpdater updater; @BeforeEach void setUp() { MockitoAnnotations.openMocks(this); - updater = new ActualAllocationUpdater(metadataStore); + updater = new ActualAllocationUpdater(metadataStore, metricsProvider); + // Use lenient() for metrics since not all tests will trigger metric emissions + lenient().when(metricsProvider.counter(anyString(), anyMap())).thenReturn(mockCounter); } // ========== PHASE 1 TESTS: Existing Index - Selective Cleanup ========== @@ -49,22 +58,22 @@ void testPhase1_ExistingIndex_ClearsStaleShards() throws Exception { SearchUnit unit = createSearchUnit("node1"); SearchUnitActualState emptyState = new SearchUnitActualState(); emptyState.setNodeRouting(new HashMap<>()); - when(metadataStore.getAllSearchUnits(clusterId)).thenReturn(Arrays.asList(unit)); + when(metadataStore.getAllSearchUnits(clusterId)).thenReturn(List.of(unit)); when(metadataStore.getSearchUnitActualState(clusterId, "node1")).thenReturn(emptyState); // Index config exists Index indexConfig = createIndex(indexName, 2); - when(metadataStore.getAllIndexConfigs(clusterId)).thenReturn(Arrays.asList(indexConfig)); + when(metadataStore.getAllIndexConfigs(clusterId)).thenReturn(List.of(indexConfig)); // Mock that the index has actual-allocation entries in etcd when(metadataStore.getAllIndicesWithActualAllocations(clusterId)) .thenReturn(Collections.singleton(indexName)); // Stored actual allocations in etcd (stale) - ShardAllocation storedShard0 = createShardAllocation("0", Arrays.asList("node1"), Arrays.asList()); - ShardAllocation storedShard1 = createShardAllocation("1", Arrays.asList("node1"), Arrays.asList()); + ShardAllocation storedShard0 = createShardAllocation("0", List.of("node1"), List.of()); + ShardAllocation storedShard1 = createShardAllocation("1", List.of("node1"), List.of()); when(metadataStore.getAllActualAllocations(clusterId, indexName)) - .thenReturn(Arrays.asList(storedShard0, storedShard1)); + .thenReturn(List.of(storedShard0, storedShard1)); // When updater.updateActualAllocations(clusterId); @@ -88,9 +97,9 @@ void testPhase2_DeletedIndex_DeletesAllAllocations() throws Exception { // Node reports having shards from the deleted index SearchUnit unit = createSearchUnit("node1"); - SearchUnitActualState actualState = createActualStateWithShards(deletedIndex, Arrays.asList(0, 1, 2)); + SearchUnitActualState actualState = createActualStateWithShards(deletedIndex, List.of(0, 1, 2)); - when(metadataStore.getAllSearchUnits(clusterId)).thenReturn(Arrays.asList(unit)); + when(metadataStore.getAllSearchUnits(clusterId)).thenReturn(List.of(unit)); when(metadataStore.getSearchUnitActualState(clusterId, "node1")).thenReturn(actualState); // NO config exists for deleted-index (this makes it a deleted index) @@ -101,12 +110,12 @@ void testPhase2_DeletedIndex_DeletesAllAllocations() throws Exception { .thenReturn(Collections.singleton(deletedIndex)); // The deleted index has actual allocations stored in etcd (orphaned entries) - ShardAllocation shard0 = createShardAllocation("0", Arrays.asList("node1"), Arrays.asList()); - ShardAllocation shard1 = createShardAllocation("1", Arrays.asList("node1"), Arrays.asList()); - ShardAllocation shard2 = createShardAllocation("2", Arrays.asList("node1"), Arrays.asList()); + ShardAllocation shard0 = createShardAllocation("0", List.of("node1"), List.of()); + ShardAllocation shard1 = createShardAllocation("1", List.of("node1"), List.of()); + ShardAllocation shard2 = createShardAllocation("2", List.of("node1"), List.of()); when(metadataStore.getAllActualAllocations(clusterId, deletedIndex)) - .thenReturn(Arrays.asList(shard0, shard1, shard2)); + .thenReturn(List.of(shard0, shard1, shard2)); // When updater.updateActualAllocations(clusterId); @@ -123,6 +132,61 @@ void testPhase2_DeletedIndex_DeletesAllAllocations() throws Exception { verify(metadataStore).deleteActualAllocation(eq(clusterId), eq(deletedIndex), eq("2")); } + + @Test + void testMetricsCounter_IncrementedOnAllocationUpdateFailure() throws Exception { + String clusterId = "test-cluster"; + String indexName = "test-index"; + String shardId = "0"; + + SearchUnit unit = createSearchUnit("node1"); + SearchUnitActualState actualState = createActualStateWithShards(indexName, List.of(0)); + when(metadataStore.getAllSearchUnits(clusterId)).thenReturn(List.of(unit)); + when(metadataStore.getSearchUnitActualState(clusterId, "node1")).thenReturn(actualState); + Index indexConfig = createIndex(indexName, 1); + when(metadataStore.getAllIndexConfigs(clusterId)).thenReturn(List.of(indexConfig)); + when(metadataStore.getAllIndicesWithActualAllocations(clusterId)) + .thenReturn(Collections.singleton(indexName)); + when(metadataStore.getActualAllocation(clusterId, indexName, shardId)) + .thenThrow(new RuntimeException("Database connection failed")); + + updater.updateActualAllocations(clusterId); + + verify(metricsProvider).counter( + eq("update_actual_allocation_failures_count"), + argThat(tags -> + tags.get("clusterId").equals(clusterId) && + tags.get("indexName").equals(indexName) && + tags.get("shardId").equals(shardId) + ) + ); + verify(mockCounter).increment(); + } + + @Test + void testMetricsCounter_NotIncrementedOnSuccess() throws Exception { + String clusterId = "test-cluster"; + String indexName = "test-index"; + + SearchUnit unit = createSearchUnit("node1"); + unit.setRole("PRIMARY"); + SearchUnitActualState actualState = createActualStateWithShards(indexName, List.of(0)); + + when(metadataStore.getAllSearchUnits(clusterId)).thenReturn(List.of(unit)); + when(metadataStore.getSearchUnitActualState(clusterId, "node1")).thenReturn(actualState); + when(metadataStore.getSearchUnit(clusterId, "node1")) + .thenReturn(java.util.Optional.of(unit)); + Index indexConfig = createIndex(indexName, 1); + when(metadataStore.getAllIndexConfigs(clusterId)).thenReturn(List.of(indexConfig)); + when(metadataStore.getAllIndicesWithActualAllocations(clusterId)) + .thenReturn(Collections.singleton(indexName)); + when(metadataStore.getActualAllocation(clusterId, indexName, "0")).thenReturn(null); + doNothing().when(metadataStore).setActualAllocation(eq(clusterId), eq(indexName), eq("0"), any()); + + updater.updateActualAllocations(clusterId); + verify(metricsProvider, never()).counter(eq("update_actual_allocation_failures_count"), anyMap()); + verify(mockCounter, never()).increment(); + } // Helper methods private Index createIndex(String indexName, int numberOfShards) { Index index = new Index(); diff --git a/src/test/java/io/clustercontroller/discovery/DiscoveryTest.java b/src/test/java/io/clustercontroller/discovery/DiscoveryTest.java index 4b64cd7c..6713c854 100644 --- a/src/test/java/io/clustercontroller/discovery/DiscoveryTest.java +++ b/src/test/java/io/clustercontroller/discovery/DiscoveryTest.java @@ -1,21 +1,36 @@ package io.clustercontroller.discovery; import io.clustercontroller.enums.HealthState; -import io.clustercontroller.enums.NodeRole; +import io.clustercontroller.metrics.MetricsProvider; import io.clustercontroller.models.SearchUnit; import io.clustercontroller.models.SearchUnitActualState; import io.clustercontroller.store.MetadataStore; +import io.micrometer.core.instrument.Counter; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; -import java.util.*; -import static org.assertj.core.api.Assertions.*; -import static org.mockito.ArgumentMatchers.*; -import static org.mockito.Mockito.*; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import static org.assertj.core.api.AssertionsForClassTypes.assertThatCode; +import static org.assertj.core.api.AssertionsForInterfaceTypes.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyMap; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.argThat; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; /** * Comprehensive tests for Discovery flow. @@ -26,12 +41,19 @@ class DiscoveryTest { @Mock private MetadataStore metadataStore; + @Mock + private MetricsProvider metricsProvider; + + @Mock + private Counter mockCounter; + private Discovery discovery; private static final String TEST_CLUSTER = "test-cluster"; @BeforeEach void setUp() { - discovery = new Discovery(metadataStore); + discovery = new Discovery(metadataStore, metricsProvider); + lenient().when(metricsProvider.counter(anyString(), anyMap())).thenReturn(mockCounter); } // ================================================================= @@ -63,7 +85,7 @@ void testDiscoverSearchUnits_UpdatesExistingUnits() throws Exception { // Given Map actualStates = createMockActualStates(); SearchUnit existingUnit = createMockSearchUnit("coordinator-node-1", "COORDINATOR"); - List existingUnits = Arrays.asList(existingUnit); + List existingUnits = List.of(existingUnit); when(metadataStore.getAllSearchUnitActualStates(anyString())).thenReturn(actualStates); when(metadataStore.getAllSearchUnits(anyString())).thenReturn(existingUnits); @@ -190,7 +212,7 @@ void testFetchSearchUnitsFromEtcd_HandlesNullRoleAndShardId() throws Exception { // Then assertThat(result).hasSize(1); - SearchUnit unit = result.get(0); + SearchUnit unit = result.getFirst(); assertThat(unit.getName()).isEqualTo("test-node"); // Should handle null gracefully assertThat(unit.getRole()).isNull(); @@ -215,7 +237,7 @@ void testConvertActualStateToSearchUnit_CoordinatorNode() throws Exception { // Then assertThat(result).hasSize(1); - SearchUnit unit = result.get(0); + SearchUnit unit = result.getFirst(); assertThat(unit.getName()).isEqualTo("coordinator-node-1"); assertThat(unit.getRole()).isEqualTo("COORDINATOR"); assertThat(unit.getShardId()).isEqualTo("COORDINATOR"); @@ -240,7 +262,7 @@ void testConvertActualStateToSearchUnit_PrimaryNode() throws Exception { // Then assertThat(result).hasSize(1); - SearchUnit unit = result.get(0); + SearchUnit unit = result.getFirst(); assertThat(unit.getName()).isEqualTo("primary-node-1"); assertThat(unit.getRole()).isEqualTo("PRIMARY"); assertThat(unit.getShardId()).isEqualTo("shard-1"); @@ -263,7 +285,7 @@ void testConvertActualStateToSearchUnit_ReplicaNode() throws Exception { // Then assertThat(result).hasSize(1); - SearchUnit unit = result.get(0); + SearchUnit unit = result.getFirst(); assertThat(unit.getName()).isEqualTo("replica-node-1"); assertThat(unit.getRole()).isEqualTo("SEARCH_REPLICA"); assertThat(unit.getShardId()).isEqualTo("shard-2"); @@ -286,11 +308,44 @@ void testConvertActualStateToSearchUnit_UnhealthyNode() throws Exception { // Then assertThat(result).hasSize(1); - SearchUnit unit = result.get(0); + SearchUnit unit = result.getFirst(); assertThat(unit.getStateAdmin()).isEqualTo("DRAIN"); assertThat(unit.getStatePulled()).isEqualTo(HealthState.RED); } - + + @Test + void testCleanupStaleSearchUnits_MetricIncrementedWhenSearchUnitsDeleted() throws Exception { + // Given: 2 stale search units (missing actual state) + SearchUnit staleUnit1 = createMockSearchUnit("stale-node-1", "PRIMARY"); + SearchUnit staleUnit2 = createMockSearchUnit("stale-node-2", "SEARCH_REPLICA"); + SearchUnit healthyUnit = createMockSearchUnit("healthy-node-1", "PRIMARY"); + + when(metadataStore.getAllSearchUnitActualStates(anyString())).thenReturn(new HashMap<>()); + when(metadataStore.getAllSearchUnits(TEST_CLUSTER)) + .thenReturn(List.of(staleUnit1, staleUnit2, healthyUnit)); + + // staleUnit1 and staleUnit2 have no actual state (null) + when(metadataStore.getSearchUnitActualState(TEST_CLUSTER, "stale-node-1")).thenReturn(null); + when(metadataStore.getSearchUnitActualState(TEST_CLUSTER, "stale-node-2")).thenReturn(null); + + // healthyUnit has actual state with current timestamp + SearchUnitActualState healthyState = createHealthyActualState("healthy-node-1", "10.0.1.5", 9200, 9300); + healthyState.setTimestamp(System.currentTimeMillis()); // Current timestamp + when(metadataStore.getSearchUnitActualState(TEST_CLUSTER, "healthy-node-1")).thenReturn(healthyState); + + discovery.discoverSearchUnits(TEST_CLUSTER); + + verify(metadataStore).deleteSearchUnit(TEST_CLUSTER, "stale-node-1"); + verify(metadataStore).deleteSearchUnit(TEST_CLUSTER, "stale-node-2"); + verify(metadataStore, never()).deleteSearchUnit(TEST_CLUSTER, "healthy-node-1"); + verify(metricsProvider, times(1)).counter( + eq("discovery_cleaned_stale_search_units_count"), + argThat(tags -> + tags.get("clusterId").equals(TEST_CLUSTER) + ) + ); + verify(mockCounter).increment(2); + } // ================================================================= // HELPER METHODS