Examples with UserTaskManager com.linkedin.kafka.cruisecontrol.servlet.UserTaskManager used on opensource projects

Example 1 with UserTaskManager

use of com.linkedin.kafka.cruisecontrol.servlet.UserTaskManager in project cruise-control by linkedin.

the class ExecutorTest method executeAndVerifyProposals.

private void executeAndVerifyProposals(KafkaZkClient kafkaZkClient, Collection<ExecutionProposal> proposalsToExecute, Collection<ExecutionProposal> proposalsToCheck, boolean completeWithError, Long replicationThrottle, boolean verifyProgress, boolean isTriggeredByUserRequest) throws OngoingExecutionException {
    KafkaCruiseControlConfig configs = new KafkaCruiseControlConfig(getExecutorProperties());
    UserTaskManager.UserTaskInfo mockUserTaskInfo = getMockUserTaskInfo();
    UserTaskManager mockUserTaskManager = isTriggeredByUserRequest ? getMockUserTaskManager(RANDOM_UUID, mockUserTaskInfo, Collections.singletonList(completeWithError)) : null;
    ExecutorNotifier mockExecutorNotifier = EasyMock.mock(ExecutorNotifier.class);
    LoadMonitor mockLoadMonitor = getMockLoadMonitor();
    Capture<String> captureMessage = Capture.newInstance(CaptureType.FIRST);
    AnomalyDetectorManager mockAnomalyDetectorManager = getMockAnomalyDetector(RANDOM_UUID);
    if (completeWithError) {
        mockExecutorNotifier.sendAlert(EasyMock.capture(captureMessage));
    } else {
        mockExecutorNotifier.sendNotification(EasyMock.capture(captureMessage));
    }
    if (isTriggeredByUserRequest) {
        EasyMock.replay(mockUserTaskInfo, mockUserTaskManager, mockExecutorNotifier, mockLoadMonitor, mockAnomalyDetectorManager);
    } else {
        EasyMock.replay(mockUserTaskInfo, mockExecutorNotifier, mockLoadMonitor, mockAnomalyDetectorManager);
    }
    Executor executor = new Executor(configs, new SystemTime(), new MetricRegistry(), null, mockExecutorNotifier, mockAnomalyDetectorManager);
    executor.setUserTaskManager(mockUserTaskManager);
    Map<TopicPartition, Integer> replicationFactors = new HashMap<>();
    for (ExecutionProposal proposal : proposalsToCheck) {
        TopicPartition tp = new TopicPartition(proposal.topic(), proposal.partitionId());
        replicationFactors.put(tp, proposal.oldReplicas().size());
    }
    executor.setGeneratingProposalsForExecution(RANDOM_UUID, ExecutorTest.class::getSimpleName, isTriggeredByUserRequest);
    executor.executeProposals(proposalsToExecute, Collections.emptySet(), null, mockLoadMonitor, null, null, null, null, null, null, replicationThrottle, isTriggeredByUserRequest, RANDOM_UUID, false, false);
    if (verifyProgress) {
        verifyOngoingPartitionReassignments(Collections.singleton(TP0));
    }
    waitUntilTrue(() -> (!executor.hasOngoingExecution() && executor.state().state() == ExecutorState.State.NO_TASK_IN_PROGRESS), "Proposal execution did not finish within the time limit", EXECUTION_DEADLINE_MS, EXECUTION_REGULAR_CHECK_MS);
    // Check notification is sent after execution has finished.
    String notification = captureMessage.getValue();
    assertTrue(notification.contains(RANDOM_UUID));
    for (ExecutionProposal proposal : proposalsToCheck) {
        TopicPartition tp = new TopicPartition(proposal.topic(), proposal.partitionId());
        int expectedReplicationFactor = replicationFactors.get(tp);
        assertEquals("Replication factor for partition " + tp + " should be " + expectedReplicationFactor, expectedReplicationFactor, kafkaZkClient.getReplicasForPartition(tp).size());
        if (proposal.hasReplicaAction()) {
            for (ReplicaPlacementInfo r : proposal.newReplicas()) {
                assertTrue("The partition should have moved for " + tp, kafkaZkClient.getReplicasForPartition(tp).contains(r.brokerId()));
            }
        }
        assertEquals("The leader should have moved for " + tp, proposal.newLeader().brokerId(), kafkaZkClient.getLeaderForPartition(tp).get());
    }
    if (isTriggeredByUserRequest) {
        EasyMock.verify(mockUserTaskInfo, mockUserTaskManager, mockExecutorNotifier, mockLoadMonitor, mockAnomalyDetectorManager);
    } else {
        EasyMock.verify(mockUserTaskInfo, mockExecutorNotifier, mockLoadMonitor, mockAnomalyDetectorManager);
    }
}

Example 2 with UserTaskManager

use of com.linkedin.kafka.cruisecontrol.servlet.UserTaskManager in project cruise-control by linkedin.

the class ExecutorTest method testTimeoutAndExecutionStop.

@Test
public void testTimeoutAndExecutionStop() throws InterruptedException, OngoingExecutionException {
    createTopics(0);
    // The proposal tries to move the leader. We fake the replica list to be unchanged so there is no replica
    // movement, but only leader movement.
    ExecutionProposal proposal = new ExecutionProposal(TP1, 0, new ReplicaPlacementInfo(1), Arrays.asList(new ReplicaPlacementInfo(0), new ReplicaPlacementInfo(1)), Arrays.asList(new ReplicaPlacementInfo(0), new ReplicaPlacementInfo(1)));
    KafkaCruiseControlConfig configs = new KafkaCruiseControlConfig(getExecutorProperties());
    Time time = new MockTime();
    MetadataClient mockMetadataClient = EasyMock.mock(MetadataClient.class);
    // Fake the metadata to never change so the leader movement will timeout.
    Node node0 = new Node(0, "host0", 100);
    Node node1 = new Node(1, "host1", 100);
    Node[] replicas = new Node[2];
    replicas[0] = node0;
    replicas[1] = node1;
    PartitionInfo partitionInfo = new PartitionInfo(TP1.topic(), TP1.partition(), node1, replicas, replicas);
    Cluster cluster = new Cluster("id", Arrays.asList(node0, node1), Collections.singleton(partitionInfo), Collections.emptySet(), Collections.emptySet());
    MetadataClient.ClusterAndGeneration clusterAndGeneration = new MetadataClient.ClusterAndGeneration(cluster, 0);
    EasyMock.expect(mockMetadataClient.refreshMetadata()).andReturn(clusterAndGeneration).anyTimes();
    EasyMock.expect(mockMetadataClient.cluster()).andReturn(clusterAndGeneration.cluster()).anyTimes();
    LoadMonitor mockLoadMonitor = getMockLoadMonitor();
    AnomalyDetectorManager mockAnomalyDetectorManager = getMockAnomalyDetector(RANDOM_UUID);
    UserTaskManager.UserTaskInfo mockUserTaskInfo = getMockUserTaskInfo();
    // This tests runs two consecutive executions. First one completes w/o error, but the second one with error.
    UserTaskManager mockUserTaskManager = getMockUserTaskManager(RANDOM_UUID, mockUserTaskInfo, Arrays.asList(false, true));
    EasyMock.replay(mockMetadataClient, mockLoadMonitor, mockAnomalyDetectorManager, mockUserTaskInfo, mockUserTaskManager);
    Collection<ExecutionProposal> proposalsToExecute = Collections.singletonList(proposal);
    Executor executor = new Executor(configs, time, new MetricRegistry(), mockMetadataClient, null, mockAnomalyDetectorManager);
    executor.setUserTaskManager(mockUserTaskManager);
    executor.setGeneratingProposalsForExecution(RANDOM_UUID, ExecutorTest.class::getSimpleName, true);
    executor.executeProposals(proposalsToExecute, Collections.emptySet(), null, mockLoadMonitor, null, null, null, null, null, null, null, true, RANDOM_UUID, false, false);
    waitUntilTrue(() -> (executor.state().state() == ExecutorState.State.LEADER_MOVEMENT_TASK_IN_PROGRESS && !executor.inExecutionTasks().isEmpty()), "Leader movement task did not start within the time limit", EXECUTION_DEADLINE_MS, EXECUTION_SHORT_CHECK_MS);
    // Sleep over ExecutorConfig#DEFAULT_LEADER_MOVEMENT_TIMEOUT_MS with some margin for inter-thread synchronization.
    time.sleep(ExecutorConfig.DEFAULT_LEADER_MOVEMENT_TIMEOUT_MS + 1L);
    // The execution should finish.
    waitUntilTrue(() -> (!executor.hasOngoingExecution() && executor.state().state() == ExecutorState.State.NO_TASK_IN_PROGRESS), "Proposal execution did not finish within the time limit", EXECUTION_DEADLINE_MS, EXECUTION_REGULAR_CHECK_MS);
    // The proposal tries to move replicas.
    proposal = new ExecutionProposal(TP1, 0, new ReplicaPlacementInfo(1), Arrays.asList(new ReplicaPlacementInfo(0), new ReplicaPlacementInfo(1)), Arrays.asList(new ReplicaPlacementInfo(1), new ReplicaPlacementInfo(0)));
    Collection<ExecutionProposal> newProposalsToExecute = Collections.singletonList(proposal);
    // Expect exception in case of UUID mismatch between UNKNOWN_UUID and RANDOM_UUID.
    executor.setGeneratingProposalsForExecution(UNKNOWN_UUID, ExecutorTest.class::getSimpleName, true);
    assertThrows(IllegalStateException.class, () -> executor.executeProposals(newProposalsToExecute, Collections.emptySet(), null, mockLoadMonitor, null, null, null, null, null, null, null, true, RANDOM_UUID, false, false));
    executor.failGeneratingProposalsForExecution(UNKNOWN_UUID);
    // Now successfully start the execution..
    executor.setGeneratingProposalsForExecution(RANDOM_UUID, ExecutorTest.class::getSimpleName, true);
    executor.executeProposals(newProposalsToExecute, Collections.emptySet(), null, mockLoadMonitor, null, null, null, null, null, null, null, true, RANDOM_UUID, false, false);
    waitUntilTrue(() -> (executor.state().state() == ExecutorState.State.INTER_BROKER_REPLICA_MOVEMENT_TASK_IN_PROGRESS), "Inter-broker replica movement task did not start within the time limit", EXECUTION_DEADLINE_MS, EXECUTION_SHORT_CHECK_MS);
    // Stop execution.
    executor.userTriggeredStopExecution(false);
    // The execution should finish.
    waitUntilTrue(() -> (!executor.hasOngoingExecution() && executor.state().state() == ExecutorState.State.NO_TASK_IN_PROGRESS), "Proposal execution did not finish within the time limit", EXECUTION_DEADLINE_MS, EXECUTION_REGULAR_CHECK_MS);
    EasyMock.verify(mockMetadataClient, mockLoadMonitor, mockAnomalyDetectorManager, mockUserTaskInfo, mockUserTaskManager);
}

Example 3 with UserTaskManager

use of com.linkedin.kafka.cruisecontrol.servlet.UserTaskManager in project cruise-control by linkedin.

the class ExecutorTest method getMockUserTaskManager.

private static UserTaskManager getMockUserTaskManager(String uuid, UserTaskManager.UserTaskInfo userTaskInfo, List<Boolean> completeWithError) {
    UserTaskManager mockUserTaskManager = EasyMock.mock(UserTaskManager.class);
    // Handle the case that the execution started, but did not finish.
    if (completeWithError.isEmpty()) {
        EasyMock.expect(mockUserTaskManager.markTaskExecutionBegan(uuid)).andReturn(userTaskInfo).once();
    } else {
        // Return as many times as the executions to ensure that the same test can run multiple executions.
        for (boolean completeStatus : completeWithError) {
            EasyMock.expect(mockUserTaskManager.markTaskExecutionBegan(uuid)).andReturn(userTaskInfo).once();
            mockUserTaskManager.markTaskExecutionFinished(uuid, completeStatus);
        }
    }
    return mockUserTaskManager;
}