Examples with PartitionState com.linkedin.d2.balancer.strategies.relative.PartitionState used on opensource projects

Example 11 with PartitionState

use of com.linkedin.d2.balancer.strategies.relative.PartitionState in project rest.li by linkedin.

the class StateUpdaterTest method testClusterGenerationIdChange.

// Known to be flaky in CI
@Test(retryAnalyzer = ThreeRetries.class)
public void testClusterGenerationIdChange() throws InterruptedException {
    PartitionState state = new PartitionStateTestDataBuilder().setClusterGenerationId(DEFAULT_CLUSTER_GENERATION_ID).build();
    List<TrackerClient> trackerClients = TrackerClientMockHelper.mockTrackerClients(2, Arrays.asList(20, 20), Arrays.asList(10, 10), Arrays.asList(200L, 500L), Arrays.asList(100L, 200L), Arrays.asList(0, 0));
    ConcurrentMap<Integer, PartitionState> partitionLoadBalancerStateMap = new ConcurrentHashMap<>();
    partitionLoadBalancerStateMap.put(DEFAULT_PARTITION_ID, state);
    ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor();
    D2RelativeStrategyProperties relativeStrategyProperties = RelativeLoadBalancerStrategyFactory.putDefaultValues(new D2RelativeStrategyProperties());
    _stateUpdater = new StateUpdater(relativeStrategyProperties, _quarantineManager, executorService, partitionLoadBalancerStateMap, Collections.emptyList(), SERVICE_NAME);
    // update will be scheduled twice, once from interval update, once from cluster change
    CountDownLatch countDownLatch = new CountDownLatch(2);
    Mockito.doAnswer(new ExecutionCountDown<>(countDownLatch)).when(_quarantineManager).updateQuarantineState(any(), any(), anyLong());
    // Cluster generation id changed from 0 to 1
    _stateUpdater.updateState(new HashSet<>(trackerClients), DEFAULT_PARTITION_ID, 1, false);
    if (!countDownLatch.await(5, TimeUnit.SECONDS)) {
        fail("cluster update failed to finish within 5 seconds");
    }
    assertEquals(_stateUpdater.getPointsMap(DEFAULT_PARTITION_ID).size(), 2);
    executorService.shutdown();
}

Example 12 with PartitionState

use of com.linkedin.d2.balancer.strategies.relative.PartitionState in project rest.li by linkedin.

the class StateUpdaterTest method testForceUpdate.

// Known to be flaky in CI
@Test(retryAnalyzer = ThreeRetries.class)
public void testForceUpdate() throws InterruptedException {
    PartitionState state = new PartitionStateTestDataBuilder().setClusterGenerationId(DEFAULT_CLUSTER_GENERATION_ID).build();
    List<TrackerClient> trackerClients = TrackerClientMockHelper.mockTrackerClients(2, Arrays.asList(20, 20), Arrays.asList(10, 10), Arrays.asList(200L, 500L), Arrays.asList(100L, 200L), Arrays.asList(0, 0));
    ConcurrentMap<Integer, PartitionState> partitionLoadBalancerStateMap = new ConcurrentHashMap<>();
    partitionLoadBalancerStateMap.put(DEFAULT_PARTITION_ID, state);
    ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor();
    D2RelativeStrategyProperties relativeStrategyProperties = RelativeLoadBalancerStrategyFactory.putDefaultValues(new D2RelativeStrategyProperties());
    _stateUpdater = new StateUpdater(relativeStrategyProperties, _quarantineManager, executorService, partitionLoadBalancerStateMap, Collections.emptyList(), SERVICE_NAME);
    // update will be scheduled three times, once from interval update, once from cluster change, once from force update
    CountDownLatch countDownLatch = new CountDownLatch(3);
    Mockito.doAnswer(new ExecutionCountDown<>(countDownLatch)).when(_quarantineManager).updateQuarantineState(any(), any(), anyLong());
    // Cluster generation id changed from 0 to 1
    _stateUpdater.updateState(new HashSet<>(trackerClients), DEFAULT_PARTITION_ID, 1, false);
    _stateUpdater.updateState(new HashSet<>(trackerClients), DEFAULT_PARTITION_ID, 1, true);
    if (!countDownLatch.await(5, TimeUnit.SECONDS)) {
        fail("cluster update failed to finish within 5 seconds");
    }
    assertEquals(_stateUpdater.getPointsMap(DEFAULT_PARTITION_ID).size(), 2);
    executorService.shutdown();
}

Example 13 with PartitionState

use of com.linkedin.d2.balancer.strategies.relative.PartitionState in project rest.li by linkedin.

the class QuarantineManagerTest method testEnrollOneQuarantine.

@Test
public void testEnrollOneQuarantine() {
    LoadBalancerQuarantine quarantine = Mockito.mock(LoadBalancerQuarantine.class);
    List<TrackerClient> trackerClients = TrackerClientMockHelper.mockTrackerClients(3);
    Map<TrackerClient, LoadBalancerQuarantine> existingQuarantineMap = new HashMap<>();
    setup(0.5, true, true);
    _quarantineManager.tryEnableQuarantine();
    PartitionState state = new PartitionStateTestDataBuilder().setTrackerClientStateMap(trackerClients, Arrays.asList(StateUpdater.MIN_HEALTH_SCORE, StateUpdater.MIN_HEALTH_SCORE, QuarantineManager.INITIAL_RECOVERY_HEALTH_SCORE), Arrays.asList(TrackerClientState.HealthState.UNHEALTHY, TrackerClientState.HealthState.NEUTRAL, TrackerClientState.HealthState.UNHEALTHY), Arrays.asList(20, 20, 20)).build();
    _quarantineManager.updateQuarantineState(state, state, 10000);
    assertEquals(state.getQuarantineMap().size(), 1, "No exception expected");
    assertTrue(state.getQuarantineMap().containsKey(trackerClients.get(0)));
}

Example 14 with PartitionState

use of com.linkedin.d2.balancer.strategies.relative.PartitionState in project rest.li by linkedin.

the class QuarantineManagerTest method testQuarantineCheck.

@Test(dataProvider = "trueFalse")
public void testQuarantineCheck(boolean quarantineCheckResult) {
    setup(0.5, false, false);
    LoadBalancerQuarantine quarantine = Mockito.mock(LoadBalancerQuarantine.class);
    List<TrackerClient> trackerClients = TrackerClientMockHelper.mockTrackerClients(3);
    Map<TrackerClient, LoadBalancerQuarantine> existingQuarantineMap = new HashMap<>();
    existingQuarantineMap.put(trackerClients.get(0), quarantine);
    Mockito.when(quarantine.checkUpdateQuarantineState()).thenReturn(quarantineCheckResult);
    PartitionState state = new PartitionStateTestDataBuilder().setTrackerClientStateMap(trackerClients, Arrays.asList(StateUpdater.MIN_HEALTH_SCORE, 0.6, 0.6), Arrays.asList(TrackerClientState.HealthState.NEUTRAL, TrackerClientState.HealthState.UNHEALTHY, TrackerClientState.HealthState.UNHEALTHY), Arrays.asList(20, 20, 20)).setQuarantineMap(existingQuarantineMap).build();
    _quarantineManager.tryEnableQuarantine();
    _quarantineManager.updateQuarantineState(state, state, DEFAULT_AVG_CLUSTER_LATENCY);
    if (quarantineCheckResult) {
        assertTrue(state.getQuarantineMap().isEmpty());
        assertEquals(state.getRecoveryTrackerClients().size(), 1, "The quarantine should be over and the host should be put into recovery");
        assertEquals(state.getTrackerClientStateMap().get(trackerClients.get(0)).getHealthScore(), QuarantineManager.INITIAL_RECOVERY_HEALTH_SCORE);
    } else {
        assertEquals(state.getQuarantineMap().size(), 1, "Quarantine health check failed, the host should be kept in quarantine state");
        assertTrue(state.getRecoveryTrackerClients().isEmpty(), "No client should be in recovery state");
        assertTrue(state.getQuarantineMap().containsKey(trackerClients.get(0)));
    }
}

Example 15 with PartitionState

use of com.linkedin.d2.balancer.strategies.relative.PartitionState in project rest.li by linkedin.

the class RelativeLoadBalancerStrategyJmxTest method mockRelativeLoadBalancerStrategyJmx.

private RelativeLoadBalancerStrategyJmx mockRelativeLoadBalancerStrategyJmx(List<TrackerClient> trackerClients) {
    Map<TrackerClient, TrackerClientState> trackerClientsMap = new HashMap<>();
    for (TrackerClient trackerClient : trackerClients) {
        trackerClientsMap.put(trackerClient, new TrackerClientState(1, 1));
    }
    RelativeLoadBalancerStrategy strategy = Mockito.mock(RelativeLoadBalancerStrategy.class);
    PartitionState state = Mockito.mock(PartitionState.class);
    Mockito.when(state.getTrackerClientStateMap()).thenReturn(trackerClientsMap);
    Mockito.when(strategy.getFirstValidPartitionId()).thenReturn(DefaultPartitionAccessor.DEFAULT_PARTITION_ID);
    Mockito.when(strategy.getPartitionState(anyInt())).thenReturn(state);
    return new RelativeLoadBalancerStrategyJmx(strategy);
}