Examples with ConsumerRebalanceListener org.apache.kafka.clients.consumer.ConsumerRebalanceListener used on opensource projects

Example 1 with ConsumerRebalanceListener

use of org.apache.kafka.clients.consumer.ConsumerRebalanceListener in project kafka by apache.

the class StreamThreadTest method testMaybeCommit.

@Test
public void testMaybeCommit() throws Exception {
    File baseDir = Files.createTempDirectory("test").toFile();
    try {
        final long commitInterval = 1000L;
        Properties props = configProps();
        props.setProperty(StreamsConfig.STATE_DIR_CONFIG, baseDir.getCanonicalPath());
        props.setProperty(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, Long.toString(commitInterval));
        StreamsConfig config = new StreamsConfig(props);
        final MockTime mockTime = new MockTime();
        TopologyBuilder builder = new TopologyBuilder().setApplicationId("X");
        builder.addSource("source1", "topic1");
        MockClientSupplier mockClientSupplier = new MockClientSupplier();
        StreamThread thread = new StreamThread(builder, config, mockClientSupplier, applicationId, clientId, processId, new Metrics(), mockTime, new StreamsMetadataState(builder, StreamsMetadataState.UNKNOWN_HOST), 0) {

            @Override
            public void maybeCommit(long now) {
                super.maybeCommit(now);
            }

            @Override
            protected StreamTask createStreamTask(TaskId id, Collection<TopicPartition> partitionsForTask) {
                ProcessorTopology topology = builder.build(id.topicGroupId);
                return new TestStreamTask(id, applicationId, partitionsForTask, topology, consumer, producer, restoreConsumer, config, new MockStreamsMetrics(new Metrics()), stateDirectory);
            }
        };
        initPartitionGrouper(config, thread, mockClientSupplier);
        ConsumerRebalanceListener rebalanceListener = thread.rebalanceListener;
        List<TopicPartition> revokedPartitions;
        List<TopicPartition> assignedPartitions;
        //
        // Assign t1p1 and t1p2. This should create Task 1 & 2
        //
        revokedPartitions = Collections.emptyList();
        assignedPartitions = Arrays.asList(t1p1, t1p2);
        rebalanceListener.onPartitionsRevoked(revokedPartitions);
        rebalanceListener.onPartitionsAssigned(assignedPartitions);
        assertEquals(2, thread.tasks().size());
        // no task is committed before the commit interval
        mockTime.sleep(commitInterval - 10L);
        thread.maybeCommit(mockTime.milliseconds());
        for (StreamTask task : thread.tasks().values()) {
            assertFalse(((TestStreamTask) task).committed);
        }
        // all tasks are committed after the commit interval
        mockTime.sleep(11L);
        thread.maybeCommit(mockTime.milliseconds());
        for (StreamTask task : thread.tasks().values()) {
            assertTrue(((TestStreamTask) task).committed);
            ((TestStreamTask) task).committed = false;
        }
        // no task is committed before the commit interval, again
        mockTime.sleep(commitInterval - 10L);
        thread.maybeCommit(mockTime.milliseconds());
        for (StreamTask task : thread.tasks().values()) {
            assertFalse(((TestStreamTask) task).committed);
        }
        // all tasks are committed after the commit interval, again
        mockTime.sleep(11L);
        thread.maybeCommit(mockTime.milliseconds());
        for (StreamTask task : thread.tasks().values()) {
            assertTrue(((TestStreamTask) task).committed);
            ((TestStreamTask) task).committed = false;
        }
    } finally {
        Utils.delete(baseDir);
    }
}

Example 2 with ConsumerRebalanceListener

use of org.apache.kafka.clients.consumer.ConsumerRebalanceListener in project kafka by apache.

the class StreamThreadTest method testMaybeClean.

@Test
public void testMaybeClean() throws Exception {
    File baseDir = Files.createTempDirectory("test").toFile();
    try {
        final long cleanupDelay = 1000L;
        Properties props = configProps();
        props.setProperty(StreamsConfig.STATE_CLEANUP_DELAY_MS_CONFIG, Long.toString(cleanupDelay));
        props.setProperty(StreamsConfig.STATE_DIR_CONFIG, baseDir.getCanonicalPath());
        StreamsConfig config = new StreamsConfig(props);
        File applicationDir = new File(baseDir, applicationId);
        applicationDir.mkdir();
        File stateDir1 = new File(applicationDir, task1.toString());
        File stateDir2 = new File(applicationDir, task2.toString());
        File stateDir3 = new File(applicationDir, task3.toString());
        File extraDir = new File(applicationDir, "X");
        stateDir1.mkdir();
        stateDir2.mkdir();
        stateDir3.mkdir();
        extraDir.mkdir();
        final MockTime mockTime = new MockTime();
        TopologyBuilder builder = new TopologyBuilder().setApplicationId("X");
        builder.addSource("source1", "topic1");
        MockClientSupplier mockClientSupplier = new MockClientSupplier();
        StreamThread thread = new StreamThread(builder, config, mockClientSupplier, applicationId, clientId, processId, new Metrics(), mockTime, new StreamsMetadataState(builder, StreamsMetadataState.UNKNOWN_HOST), 0) {

            @Override
            public void maybeClean(long now) {
                super.maybeClean(now);
            }

            @Override
            protected StreamTask createStreamTask(TaskId id, Collection<TopicPartition> partitionsForTask) {
                ProcessorTopology topology = builder.build(id.topicGroupId);
                return new TestStreamTask(id, applicationId, partitionsForTask, topology, consumer, producer, restoreConsumer, config, new MockStreamsMetrics(new Metrics()), stateDirectory);
            }
        };
        initPartitionGrouper(config, thread, mockClientSupplier);
        ConsumerRebalanceListener rebalanceListener = thread.rebalanceListener;
        assertTrue(thread.tasks().isEmpty());
        mockTime.sleep(cleanupDelay);
        // all directories exist since an assignment didn't happen
        assertTrue(stateDir1.exists());
        assertTrue(stateDir2.exists());
        assertTrue(stateDir3.exists());
        assertTrue(extraDir.exists());
        List<TopicPartition> revokedPartitions;
        List<TopicPartition> assignedPartitions;
        Map<TaskId, StreamTask> prevTasks;
        //
        // Assign t1p1 and t1p2. This should create task1 & task2
        //
        final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
        activeTasks.put(task1, Collections.singleton(t1p1));
        activeTasks.put(task2, Collections.singleton(t1p2));
        thread.partitionAssignor(new MockStreamsPartitionAssignor(activeTasks));
        revokedPartitions = Collections.emptyList();
        assignedPartitions = Arrays.asList(t1p1, t1p2);
        prevTasks = new HashMap<>(thread.tasks());
        rebalanceListener.onPartitionsRevoked(revokedPartitions);
        rebalanceListener.onPartitionsAssigned(assignedPartitions);
        // there shouldn't be any previous task
        assertTrue(prevTasks.isEmpty());
        // task 1 & 2 are created
        assertEquals(2, thread.tasks().size());
        // all directories should still exit before the cleanup delay time
        mockTime.sleep(cleanupDelay - 10L);
        thread.maybeClean(mockTime.milliseconds());
        assertTrue(stateDir1.exists());
        assertTrue(stateDir2.exists());
        assertTrue(stateDir3.exists());
        assertTrue(extraDir.exists());
        // all state directories except for task task2 & task3 will be removed. the extra directory should still exists
        mockTime.sleep(11L);
        thread.maybeClean(mockTime.milliseconds());
        assertTrue(stateDir1.exists());
        assertTrue(stateDir2.exists());
        assertFalse(stateDir3.exists());
        assertTrue(extraDir.exists());
        //
        // Revoke t1p1 and t1p2. This should remove task1 & task2
        //
        activeTasks.clear();
        revokedPartitions = assignedPartitions;
        assignedPartitions = Collections.emptyList();
        prevTasks = new HashMap<>(thread.tasks());
        rebalanceListener.onPartitionsRevoked(revokedPartitions);
        rebalanceListener.onPartitionsAssigned(assignedPartitions);
        // previous tasks should be committed
        assertEquals(2, prevTasks.size());
        for (StreamTask task : prevTasks.values()) {
            assertTrue(((TestStreamTask) task).committed);
            ((TestStreamTask) task).committed = false;
        }
        // no task
        assertTrue(thread.tasks().isEmpty());
        // all state directories for task task1 & task2 still exist before the cleanup delay time
        mockTime.sleep(cleanupDelay - 10L);
        thread.maybeClean(mockTime.milliseconds());
        assertTrue(stateDir1.exists());
        assertTrue(stateDir2.exists());
        assertFalse(stateDir3.exists());
        assertTrue(extraDir.exists());
        // all state directories for task task1 & task2 are removed
        mockTime.sleep(11L);
        thread.maybeClean(mockTime.milliseconds());
        assertFalse(stateDir1.exists());
        assertFalse(stateDir2.exists());
        assertFalse(stateDir3.exists());
        assertTrue(extraDir.exists());
    } finally {
        Utils.delete(baseDir);
    }
}

Example 3 with ConsumerRebalanceListener

use of org.apache.kafka.clients.consumer.ConsumerRebalanceListener in project kafka by apache.

the class StreamThreadTest method testPartitionAssignmentChange.

@SuppressWarnings("unchecked")
@Test
public void testPartitionAssignmentChange() throws Exception {
    StreamsConfig config = new StreamsConfig(configProps());
    StateListenerStub stateListener = new StateListenerStub();
    TopologyBuilder builder = new TopologyBuilder().setApplicationId("X");
    builder.addSource("source1", "topic1");
    builder.addSource("source2", "topic2");
    builder.addSource("source3", "topic3");
    builder.addProcessor("processor", new MockProcessorSupplier(), "source2", "source3");
    MockClientSupplier mockClientSupplier = new MockClientSupplier();
    StreamThread thread = new StreamThread(builder, config, mockClientSupplier, applicationId, clientId, processId, new Metrics(), Time.SYSTEM, new StreamsMetadataState(builder, StreamsMetadataState.UNKNOWN_HOST), 0) {

        @Override
        protected StreamTask createStreamTask(TaskId id, Collection<TopicPartition> partitionsForTask) {
            ProcessorTopology topology = builder.build(id.topicGroupId);
            return new TestStreamTask(id, applicationId, partitionsForTask, topology, consumer, producer, restoreConsumer, config, new MockStreamsMetrics(new Metrics()), stateDirectory);
        }
    };
    thread.setStateListener(stateListener);
    assertEquals(thread.state(), StreamThread.State.RUNNING);
    initPartitionGrouper(config, thread, mockClientSupplier);
    ConsumerRebalanceListener rebalanceListener = thread.rebalanceListener;
    assertTrue(thread.tasks().isEmpty());
    List<TopicPartition> revokedPartitions;
    List<TopicPartition> assignedPartitions;
    Set<TopicPartition> expectedGroup1;
    Set<TopicPartition> expectedGroup2;
    revokedPartitions = Collections.emptyList();
    assignedPartitions = Collections.singletonList(t1p1);
    expectedGroup1 = new HashSet<>(Arrays.asList(t1p1));
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    assertEquals(thread.state(), StreamThread.State.PARTITIONS_REVOKED);
    Assert.assertEquals(stateListener.numChanges, 1);
    Assert.assertEquals(stateListener.oldState, StreamThread.State.RUNNING);
    Assert.assertEquals(stateListener.newState, StreamThread.State.PARTITIONS_REVOKED);
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    assertEquals(thread.state(), StreamThread.State.RUNNING);
    Assert.assertEquals(stateListener.numChanges, 3);
    Assert.assertEquals(stateListener.oldState, StreamThread.State.ASSIGNING_PARTITIONS);
    Assert.assertEquals(stateListener.newState, StreamThread.State.RUNNING);
    assertTrue(thread.tasks().containsKey(task1));
    assertEquals(expectedGroup1, thread.tasks().get(task1).partitions());
    assertEquals(1, thread.tasks().size());
    revokedPartitions = assignedPartitions;
    assignedPartitions = Collections.singletonList(t1p2);
    expectedGroup2 = new HashSet<>(Arrays.asList(t1p2));
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    assertFalse(thread.tasks().containsKey(task1));
    assertEquals(0, thread.tasks().size());
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    assertTrue(thread.tasks().containsKey(task2));
    assertEquals(expectedGroup2, thread.tasks().get(task2).partitions());
    assertEquals(1, thread.tasks().size());
    revokedPartitions = assignedPartitions;
    assignedPartitions = Arrays.asList(t1p1, t1p2);
    expectedGroup1 = new HashSet<>(Collections.singleton(t1p1));
    expectedGroup2 = new HashSet<>(Collections.singleton(t1p2));
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    assertTrue(thread.tasks().containsKey(task1));
    assertTrue(thread.tasks().containsKey(task2));
    assertEquals(expectedGroup1, thread.tasks().get(task1).partitions());
    assertEquals(expectedGroup2, thread.tasks().get(task2).partitions());
    assertEquals(2, thread.tasks().size());
    revokedPartitions = assignedPartitions;
    assignedPartitions = Arrays.asList(t2p1, t2p2, t3p1, t3p2);
    expectedGroup1 = new HashSet<>(Arrays.asList(t2p1, t3p1));
    expectedGroup2 = new HashSet<>(Arrays.asList(t2p2, t3p2));
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    assertTrue(thread.tasks().containsKey(task4));
    assertTrue(thread.tasks().containsKey(task5));
    assertEquals(expectedGroup1, thread.tasks().get(task4).partitions());
    assertEquals(expectedGroup2, thread.tasks().get(task5).partitions());
    assertEquals(2, thread.tasks().size());
    revokedPartitions = assignedPartitions;
    assignedPartitions = Arrays.asList(t1p1, t2p1, t3p1);
    expectedGroup1 = new HashSet<>(Arrays.asList(t1p1));
    expectedGroup2 = new HashSet<>(Arrays.asList(t2p1, t3p1));
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    assertTrue(thread.tasks().containsKey(task1));
    assertTrue(thread.tasks().containsKey(task4));
    assertEquals(expectedGroup1, thread.tasks().get(task1).partitions());
    assertEquals(expectedGroup2, thread.tasks().get(task4).partitions());
    assertEquals(2, thread.tasks().size());
    revokedPartitions = assignedPartitions;
    assignedPartitions = Arrays.asList(t1p1, t2p1, t3p1);
    expectedGroup1 = new HashSet<>(Arrays.asList(t1p1));
    expectedGroup2 = new HashSet<>(Arrays.asList(t2p1, t3p1));
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    assertTrue(thread.tasks().containsKey(task1));
    assertTrue(thread.tasks().containsKey(task4));
    assertEquals(expectedGroup1, thread.tasks().get(task1).partitions());
    assertEquals(expectedGroup2, thread.tasks().get(task4).partitions());
    assertEquals(2, thread.tasks().size());
    revokedPartitions = assignedPartitions;
    assignedPartitions = Collections.emptyList();
    rebalanceListener.onPartitionsRevoked(revokedPartitions);
    rebalanceListener.onPartitionsAssigned(assignedPartitions);
    assertTrue(thread.tasks().isEmpty());
    thread.close();
    assertTrue((thread.state() == StreamThread.State.PENDING_SHUTDOWN) || (thread.state() == StreamThread.State.NOT_RUNNING));
}

Example 4 with ConsumerRebalanceListener

use of org.apache.kafka.clients.consumer.ConsumerRebalanceListener in project kafka by apache.

the class ConsumerCoordinator method onJoinPrepare.

@Override
protected void onJoinPrepare(int generation, String memberId) {
    // commit offsets prior to rebalance if auto-commit enabled
    maybeAutoCommitOffsetsSync(rebalanceTimeoutMs);
    // execute the user's callback before rebalance
    ConsumerRebalanceListener listener = subscriptions.listener();
    log.info("Revoking previously assigned partitions {} for group {}", subscriptions.assignedPartitions(), groupId);
    try {
        Set<TopicPartition> revoked = new HashSet<>(subscriptions.assignedPartitions());
        listener.onPartitionsRevoked(revoked);
    } catch (WakeupException | InterruptException e) {
        throw e;
    } catch (Exception e) {
        log.error("User provided listener {} for group {} failed on partition revocation", listener.getClass().getName(), groupId, e);
    }
    isLeader = false;
    subscriptions.resetGroupSubscription();
}

Example 5 with ConsumerRebalanceListener

use of org.apache.kafka.clients.consumer.ConsumerRebalanceListener in project kafka by apache.

the class ConsumerCoordinator method onJoinComplete.

@Override
protected void onJoinComplete(int generation, String memberId, String assignmentStrategy, ByteBuffer assignmentBuffer) {
    // only the leader is responsible for monitoring for metadata changes (i.e. partition changes)
    if (!isLeader)
        assignmentSnapshot = null;
    PartitionAssignor assignor = lookupAssignor(assignmentStrategy);
    if (assignor == null)
        throw new IllegalStateException("Coordinator selected invalid assignment protocol: " + assignmentStrategy);
    Assignment assignment = ConsumerProtocol.deserializeAssignment(assignmentBuffer);
    // set the flag to refresh last committed offsets
    subscriptions.needRefreshCommits();
    // update partition assignment
    subscriptions.assignFromSubscribed(assignment.partitions());
    // check if the assignment contains some topics that were not in the original
    // subscription, if yes we will obey what leader has decided and add these topics
    // into the subscriptions as long as they still match the subscribed pattern
    //
    // TODO this part of the logic should be removed once we allow regex on leader assign
    Set<String> addedTopics = new HashSet<>();
    for (TopicPartition tp : subscriptions.assignedPartitions()) {
        if (!joinedSubscription.contains(tp.topic()))
            addedTopics.add(tp.topic());
    }
    if (!addedTopics.isEmpty()) {
        Set<String> newSubscription = new HashSet<>(subscriptions.subscription());
        Set<String> newJoinedSubscription = new HashSet<>(joinedSubscription);
        newSubscription.addAll(addedTopics);
        newJoinedSubscription.addAll(addedTopics);
        this.subscriptions.subscribeFromPattern(newSubscription);
        this.joinedSubscription = newJoinedSubscription;
    }
    // update the metadata and enforce a refresh to make sure the fetcher can start
    // fetching data in the next iteration
    this.metadata.setTopics(subscriptions.groupSubscription());
    client.ensureFreshMetadata();
    // give the assignor a chance to update internal state based on the received assignment
    assignor.onAssignment(assignment);
    // reschedule the auto commit starting from now
    this.nextAutoCommitDeadline = time.milliseconds() + autoCommitIntervalMs;
    // execute the user's callback after rebalance
    ConsumerRebalanceListener listener = subscriptions.listener();
    log.info("Setting newly assigned partitions {} for group {}", subscriptions.assignedPartitions(), groupId);
    try {
        Set<TopicPartition> assigned = new HashSet<>(subscriptions.assignedPartitions());
        listener.onPartitionsAssigned(assigned);
    } catch (WakeupException | InterruptException e) {
        throw e;
    } catch (Exception e) {
        log.error("User provided listener {} for group {} failed on partition assignment", listener.getClass().getName(), groupId, e);
    }
}