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

Example 26 with OffsetCommitCallback

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

the class WorkerSinkTaskTest method testPreCommit.

@Test
public void testPreCommit() throws Exception {
    createTask(initialState);
    expectInitializeTask();
    expectTaskGetTopic(true);
    // iter 1
    expectPollInitialAssignment();
    // iter 2
    expectConsumerPoll(2);
    expectConversionAndTransformation(2);
    sinkTask.put(EasyMock.anyObject());
    EasyMock.expectLastCall();
    final Map<TopicPartition, OffsetAndMetadata> workerStartingOffsets = new HashMap<>();
    workerStartingOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET));
    workerStartingOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
    final Map<TopicPartition, OffsetAndMetadata> workerCurrentOffsets = new HashMap<>();
    workerCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 2));
    workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
    final Map<TopicPartition, OffsetAndMetadata> taskOffsets = new HashMap<>();
    // act like FIRST_OFFSET+2 has not yet been flushed by the task
    taskOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
    // should be ignored because > current offset
    taskOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET + 1));
    // should be ignored because this partition is not assigned
    taskOffsets.put(new TopicPartition(TOPIC, 3), new OffsetAndMetadata(FIRST_OFFSET));
    final Map<TopicPartition, OffsetAndMetadata> committableOffsets = new HashMap<>();
    committableOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
    committableOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
    sinkTask.preCommit(workerCurrentOffsets);
    EasyMock.expectLastCall().andReturn(taskOffsets);
    // Expect extra invalid topic partition to be filtered, which causes the consumer assignment to be logged
    EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(2);
    final Capture<OffsetCommitCallback> callback = EasyMock.newCapture();
    consumer.commitAsync(EasyMock.eq(committableOffsets), EasyMock.capture(callback));
    EasyMock.expectLastCall().andAnswer(() -> {
        callback.getValue().onComplete(committableOffsets, null);
        return null;
    });
    expectConsumerPoll(0);
    sinkTask.put(EasyMock.anyObject());
    EasyMock.expectLastCall();
    PowerMock.replayAll();
    workerTask.initialize(TASK_CONFIG);
    workerTask.initializeAndStart();
    // iter 1 -- initial assignment
    workerTask.iteration();
    assertEquals(workerStartingOffsets, Whitebox.<Map<TopicPartition, OffsetAndMetadata>>getInternalState(workerTask, "currentOffsets"));
    // iter 2 -- deliver 2 records
    workerTask.iteration();
    assertEquals(workerCurrentOffsets, Whitebox.<Map<TopicPartition, OffsetAndMetadata>>getInternalState(workerTask, "currentOffsets"));
    assertEquals(workerStartingOffsets, Whitebox.<Map<TopicPartition, OffsetAndMetadata>>getInternalState(workerTask, "lastCommittedOffsets"));
    sinkTaskContext.getValue().requestCommit();
    // iter 3 -- commit
    workerTask.iteration();
    assertEquals(committableOffsets, Whitebox.<Map<TopicPartition, OffsetAndMetadata>>getInternalState(workerTask, "lastCommittedOffsets"));
    PowerMock.verifyAll();
}

Example 27 with OffsetCommitCallback

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

the class WorkerSinkTaskThreadedTest method expectOffsetCommit.

private Capture<OffsetCommitCallback> expectOffsetCommit(final long expectedMessages, final RuntimeException error, final Exception consumerCommitError, final long consumerCommitDelayMs, final boolean invokeCallback) {
    final long finalOffset = FIRST_OFFSET + expectedMessages;
    // All assigned partitions will have offsets committed, but we've only processed messages/updated offsets for one
    final Map<TopicPartition, OffsetAndMetadata> offsetsToCommit = new HashMap<>();
    offsetsToCommit.put(TOPIC_PARTITION, new OffsetAndMetadata(finalOffset));
    offsetsToCommit.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
    offsetsToCommit.put(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET));
    sinkTask.preCommit(offsetsToCommit);
    IExpectationSetters<Object> expectation = PowerMock.expectLastCall();
    if (error != null) {
        expectation.andThrow(error).once();
        return null;
    } else {
        expectation.andReturn(offsetsToCommit);
    }
    final Capture<OffsetCommitCallback> capturedCallback = EasyMock.newCapture();
    consumer.commitAsync(EasyMock.eq(offsetsToCommit), EasyMock.capture(capturedCallback));
    PowerMock.expectLastCall().andAnswer(() -> {
        time.sleep(consumerCommitDelayMs);
        if (invokeCallback)
            capturedCallback.getValue().onComplete(offsetsToCommit, consumerCommitError);
        return null;
    });
    return capturedCallback;
}

Example 28 with OffsetCommitCallback

use of org.apache.kafka.clients.consumer.OffsetCommitCallback in project incubator-gobblin by apache.

the class Kafka09ConsumerClient method commitOffsetsAsync.

/**
 * Commit offsets to Kafka asynchronously
 */
@Override
public void commitOffsetsAsync(Map<KafkaPartition, Long> partitionOffsets) {
    Map<TopicPartition, OffsetAndMetadata> offsets = partitionOffsets.entrySet().stream().collect(Collectors.toMap(e -> new TopicPartition(e.getKey().getTopicName(), e.getKey().getId()), e -> new OffsetAndMetadata(e.getValue())));
    consumer.commitAsync(offsets, new OffsetCommitCallback() {

        @Override
        public void onComplete(Map<TopicPartition, OffsetAndMetadata> offsets, Exception exception) {
            if (exception != null) {
                log.error("Exception while committing offsets " + offsets, exception);
                return;
            }
        }
    });
}