Examples with TimestampType org.apache.kafka.common.record.TimestampType used on opensource projects

Example 6 with TimestampType

use of org.apache.kafka.common.record.TimestampType in project kafka by apache.

the class Fetcher method parseRecord.

/**
     * Parse the record entry, deserializing the key / value fields if necessary
     */
private ConsumerRecord<K, V> parseRecord(TopicPartition partition, LogEntry logEntry) {
    Record record = logEntry.record();
    if (this.checkCrcs) {
        try {
            record.ensureValid();
        } catch (InvalidRecordException e) {
            throw new KafkaException("Record for partition " + partition + " at offset " + logEntry.offset() + " is invalid, cause: " + e.getMessage());
        }
    }
    try {
        long offset = logEntry.offset();
        long timestamp = record.timestamp();
        TimestampType timestampType = record.timestampType();
        ByteBuffer keyBytes = record.key();
        byte[] keyByteArray = keyBytes == null ? null : Utils.toArray(keyBytes);
        K key = keyBytes == null ? null : this.keyDeserializer.deserialize(partition.topic(), keyByteArray);
        ByteBuffer valueBytes = record.value();
        byte[] valueByteArray = valueBytes == null ? null : Utils.toArray(valueBytes);
        V value = valueBytes == null ? null : this.valueDeserializer.deserialize(partition.topic(), valueByteArray);
        return new ConsumerRecord<>(partition.topic(), partition.partition(), offset, timestamp, timestampType, record.checksum(), keyByteArray == null ? ConsumerRecord.NULL_SIZE : keyByteArray.length, valueByteArray == null ? ConsumerRecord.NULL_SIZE : valueByteArray.length, key, value);
    } catch (RuntimeException e) {
        throw new SerializationException("Error deserializing key/value for partition " + partition + " at offset " + logEntry.offset(), e);
    }
}

Example 7 with TimestampType

use of org.apache.kafka.common.record.TimestampType in project incubator-pulsar by apache.

the class PulsarKafkaConsumer method poll.

@SuppressWarnings("unchecked")
@Override
public ConsumerRecords<K, V> poll(long timeoutMillis) {
    try {
        QueueItem item = receivedMessages.poll(timeoutMillis, TimeUnit.MILLISECONDS);
        if (item == null) {
            return (ConsumerRecords<K, V>) ConsumerRecords.EMPTY;
        }
        Map<TopicPartition, List<ConsumerRecord<K, V>>> records = new HashMap<>();
        int numberOfRecords = 0;
        while (item != null && ++numberOfRecords < MAX_RECORDS_IN_SINGLE_POLL) {
            TopicName topicName = TopicName.get(item.consumer.getTopic());
            String topic = topicName.getPartitionedTopicName();
            int partition = topicName.isPartitioned() ? topicName.getPartitionIndex() : 0;
            Message<byte[]> msg = item.message;
            MessageIdImpl msgId = (MessageIdImpl) msg.getMessageId();
            long offset = MessageIdUtils.getOffset(msgId);
            TopicPartition tp = new TopicPartition(topic, partition);
            K key = getKey(topic, msg);
            V value = valueDeserializer.deserialize(topic, msg.getData());
            TimestampType timestampType = TimestampType.LOG_APPEND_TIME;
            long timestamp = msg.getPublishTime();
            if (msg.getEventTime() > 0) {
                // If we have Event time, use that in preference
                timestamp = msg.getEventTime();
                timestampType = TimestampType.CREATE_TIME;
            }
            ConsumerRecord<K, V> consumerRecord = new ConsumerRecord<>(topic, partition, offset, timestamp, timestampType, -1, msg.hasKey() ? msg.getKey().length() : 0, msg.getData().length, key, value);
            records.computeIfAbsent(tp, k -> new ArrayList<>()).add(consumerRecord);
            // Update last offset seen by application
            lastReceivedOffset.put(tp, offset);
            // Check if we have an item already available
            item = receivedMessages.poll(0, TimeUnit.MILLISECONDS);
        }
        if (isAutoCommit) {
            // Commit the offset of previously dequeued messages
            commitAsync();
        }
        return new ConsumerRecords<>(records);
    } catch (InterruptedException e) {
        throw new RuntimeException(e);
    }
}

Example 8 with TimestampType

use of org.apache.kafka.common.record.TimestampType in project apache-kafka-on-k8s by banzaicloud.

the class WorkerSinkTaskTest method testCommitWithOutOfOrderCallback.

// Verify that when commitAsync is called but the supplied callback is not called by the consumer before a
// rebalance occurs, the async callback does not reset the last committed offset from the rebalance.
// See KAFKA-5731 for more information.
@Test
public void testCommitWithOutOfOrderCallback() throws Exception {
    createTask(initialState);
    expectInitializeTask();
    // iter 1
    expectPollInitialAssignment();
    // iter 2
    expectConsumerPoll(1);
    expectConversionAndTransformation(4);
    sinkTask.put(EasyMock.<Collection<SinkRecord>>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 + 1));
    workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
    final List<TopicPartition> originalPartitions = asList(TOPIC_PARTITION, TOPIC_PARTITION2);
    final List<TopicPartition> rebalancedPartitions = asList(TOPIC_PARTITION, TOPIC_PARTITION2, TOPIC_PARTITION3);
    final Map<TopicPartition, OffsetAndMetadata> rebalanceOffsets = new HashMap<>();
    rebalanceOffsets.put(TOPIC_PARTITION, workerCurrentOffsets.get(TOPIC_PARTITION));
    rebalanceOffsets.put(TOPIC_PARTITION2, workerCurrentOffsets.get(TOPIC_PARTITION2));
    rebalanceOffsets.put(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET));
    final Map<TopicPartition, OffsetAndMetadata> postRebalanceCurrentOffsets = new HashMap<>();
    postRebalanceCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 3));
    postRebalanceCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
    postRebalanceCurrentOffsets.put(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET + 2));
    // iter 3 - note that we return the current offset to indicate they should be committed
    sinkTask.preCommit(workerCurrentOffsets);
    EasyMock.expectLastCall().andReturn(workerCurrentOffsets);
    // We need to delay the result of trying to commit offsets to Kafka via the consumer.commitAsync
    // method. We do this so that we can test that the callback is not called until after the rebalance
    // changes the lastCommittedOffsets. To fake this for tests we have the commitAsync build a function
    // that will call the callback with the appropriate parameters, and we'll run that function later.
    final AtomicReference<Runnable> asyncCallbackRunner = new AtomicReference<>();
    final AtomicBoolean asyncCallbackRan = new AtomicBoolean();
    consumer.commitAsync(EasyMock.eq(workerCurrentOffsets), EasyMock.<OffsetCommitCallback>anyObject());
    EasyMock.expectLastCall().andAnswer(new IAnswer<Void>() {

        @SuppressWarnings("unchecked")
        @Override
        public Void answer() throws Throwable {
            // Grab the arguments passed to the consumer.commitAsync method
            final Object[] args = EasyMock.getCurrentArguments();
            final Map<TopicPartition, OffsetAndMetadata> offsets = (Map<TopicPartition, OffsetAndMetadata>) args[0];
            final OffsetCommitCallback callback = (OffsetCommitCallback) args[1];
            asyncCallbackRunner.set(new Runnable() {

                @Override
                public void run() {
                    callback.onComplete(offsets, null);
                    asyncCallbackRan.set(true);
                }
            });
            return null;
        }
    });
    // Expect the next poll to discover and perform the rebalance, THEN complete the previous callback handler,
    // and then return one record for TP1 and one for TP3.
    final AtomicBoolean rebalanced = new AtomicBoolean();
    EasyMock.expect(consumer.poll(EasyMock.anyLong())).andAnswer(new IAnswer<ConsumerRecords<byte[], byte[]>>() {

        @Override
        public ConsumerRecords<byte[], byte[]> answer() throws Throwable {
            // Rebalance always begins with revoking current partitions ...
            rebalanceListener.getValue().onPartitionsRevoked(originalPartitions);
            // Respond to the rebalance
            Map<TopicPartition, Long> offsets = new HashMap<>();
            offsets.put(TOPIC_PARTITION, rebalanceOffsets.get(TOPIC_PARTITION).offset());
            offsets.put(TOPIC_PARTITION2, rebalanceOffsets.get(TOPIC_PARTITION2).offset());
            offsets.put(TOPIC_PARTITION3, rebalanceOffsets.get(TOPIC_PARTITION3).offset());
            sinkTaskContext.getValue().offset(offsets);
            rebalanceListener.getValue().onPartitionsAssigned(rebalancedPartitions);
            rebalanced.set(true);
            // Run the previous async commit handler
            asyncCallbackRunner.get().run();
            // And prep the two records to return
            long timestamp = RecordBatch.NO_TIMESTAMP;
            TimestampType timestampType = TimestampType.NO_TIMESTAMP_TYPE;
            List<ConsumerRecord<byte[], byte[]>> records = new ArrayList<>();
            records.add(new ConsumerRecord<>(TOPIC, PARTITION, FIRST_OFFSET + recordsReturnedTp1 + 1, timestamp, timestampType, 0L, 0, 0, RAW_KEY, RAW_VALUE));
            records.add(new ConsumerRecord<>(TOPIC, PARTITION3, FIRST_OFFSET + recordsReturnedTp3 + 1, timestamp, timestampType, 0L, 0, 0, RAW_KEY, RAW_VALUE));
            recordsReturnedTp1 += 1;
            recordsReturnedTp3 += 1;
            return new ConsumerRecords<>(Collections.singletonMap(new TopicPartition(TOPIC, PARTITION), records));
        }
    });
    // onPartitionsRevoked
    sinkTask.preCommit(workerCurrentOffsets);
    EasyMock.expectLastCall().andReturn(workerCurrentOffsets);
    sinkTask.put(EasyMock.<Collection<SinkRecord>>anyObject());
    EasyMock.expectLastCall();
    sinkTask.close(workerCurrentOffsets.keySet());
    EasyMock.expectLastCall();
    consumer.commitSync(workerCurrentOffsets);
    EasyMock.expectLastCall();
    // onPartitionsAssigned - step 1
    final long offsetTp1 = rebalanceOffsets.get(TOPIC_PARTITION).offset();
    final long offsetTp2 = rebalanceOffsets.get(TOPIC_PARTITION2).offset();
    final long offsetTp3 = rebalanceOffsets.get(TOPIC_PARTITION3).offset();
    EasyMock.expect(consumer.position(TOPIC_PARTITION)).andReturn(offsetTp1);
    EasyMock.expect(consumer.position(TOPIC_PARTITION2)).andReturn(offsetTp2);
    EasyMock.expect(consumer.position(TOPIC_PARTITION3)).andReturn(offsetTp3);
    // onPartitionsAssigned - step 2
    sinkTask.open(rebalancedPartitions);
    EasyMock.expectLastCall();
    // onPartitionsAssigned - step 3 rewind
    consumer.seek(TOPIC_PARTITION, offsetTp1);
    EasyMock.expectLastCall();
    consumer.seek(TOPIC_PARTITION2, offsetTp2);
    EasyMock.expectLastCall();
    consumer.seek(TOPIC_PARTITION3, offsetTp3);
    EasyMock.expectLastCall();
    // iter 4 - note that we return the current offset to indicate they should be committed
    sinkTask.preCommit(postRebalanceCurrentOffsets);
    EasyMock.expectLastCall().andReturn(postRebalanceCurrentOffsets);
    final Capture<OffsetCommitCallback> callback = EasyMock.newCapture();
    consumer.commitAsync(EasyMock.eq(postRebalanceCurrentOffsets), EasyMock.capture(callback));
    EasyMock.expectLastCall().andAnswer(new IAnswer<Void>() {

        @Override
        public Void answer() throws Throwable {
            callback.getValue().onComplete(postRebalanceCurrentOffsets, null);
            return null;
        }
    });
    // no actual consumer.commit() triggered
    expectConsumerPoll(1);
    sinkTask.put(EasyMock.<Collection<SinkRecord>>anyObject());
    EasyMock.expectLastCall();
    PowerMock.replayAll();
    workerTask.initialize(TASK_CONFIG);
    workerTask.initializeAndStart();
    // iter 1 -- initial assignment
    workerTask.iteration();
    assertEquals(workerStartingOffsets, Whitebox.getInternalState(workerTask, "currentOffsets"));
    assertEquals(workerStartingOffsets, Whitebox.getInternalState(workerTask, "lastCommittedOffsets"));
    time.sleep(WorkerConfig.OFFSET_COMMIT_TIMEOUT_MS_DEFAULT);
    // iter 2 -- deliver 2 records
    workerTask.iteration();
    sinkTaskContext.getValue().requestCommit();
    // iter 3 -- commit in progress
    workerTask.iteration();
    assertSinkMetricValue("partition-count", 3);
    assertSinkMetricValue("sink-record-read-total", 3.0);
    assertSinkMetricValue("sink-record-send-total", 3.0);
    assertSinkMetricValue("sink-record-active-count", 4.0);
    assertSinkMetricValue("sink-record-active-count-max", 4.0);
    assertSinkMetricValue("sink-record-active-count-avg", 0.71429);
    assertSinkMetricValue("offset-commit-seq-no", 2.0);
    assertSinkMetricValue("offset-commit-completion-total", 1.0);
    assertSinkMetricValue("offset-commit-skip-total", 1.0);
    assertTaskMetricValue("status", "running");
    assertTaskMetricValue("running-ratio", 1.0);
    assertTaskMetricValue("pause-ratio", 0.0);
    assertTaskMetricValue("batch-size-max", 2.0);
    assertTaskMetricValue("batch-size-avg", 1.0);
    assertTaskMetricValue("offset-commit-max-time-ms", 0.0);
    assertTaskMetricValue("offset-commit-avg-time-ms", 0.0);
    assertTaskMetricValue("offset-commit-failure-percentage", 0.0);
    assertTaskMetricValue("offset-commit-success-percentage", 1.0);
    assertTrue(asyncCallbackRan.get());
    assertTrue(rebalanced.get());
    // Check that the offsets were not reset by the out-of-order async commit callback
    assertEquals(postRebalanceCurrentOffsets, Whitebox.getInternalState(workerTask, "currentOffsets"));
    assertEquals(rebalanceOffsets, Whitebox.getInternalState(workerTask, "lastCommittedOffsets"));
    time.sleep(WorkerConfig.OFFSET_COMMIT_TIMEOUT_MS_DEFAULT);
    sinkTaskContext.getValue().requestCommit();
    // iter 4 -- commit in progress
    workerTask.iteration();
    // Check that the offsets were not reset by the out-of-order async commit callback
    assertEquals(postRebalanceCurrentOffsets, Whitebox.getInternalState(workerTask, "currentOffsets"));
    assertEquals(postRebalanceCurrentOffsets, Whitebox.getInternalState(workerTask, "lastCommittedOffsets"));
    assertSinkMetricValue("partition-count", 3);
    assertSinkMetricValue("sink-record-read-total", 4.0);
    assertSinkMetricValue("sink-record-send-total", 4.0);
    assertSinkMetricValue("sink-record-active-count", 0.0);
    assertSinkMetricValue("sink-record-active-count-max", 4.0);
    assertSinkMetricValue("sink-record-active-count-avg", 0.5555555);
    assertSinkMetricValue("offset-commit-seq-no", 3.0);
    assertSinkMetricValue("offset-commit-completion-total", 2.0);
    assertSinkMetricValue("offset-commit-skip-total", 1.0);
    assertTaskMetricValue("status", "running");
    assertTaskMetricValue("running-ratio", 1.0);
    assertTaskMetricValue("pause-ratio", 0.0);
    assertTaskMetricValue("batch-size-max", 2.0);
    assertTaskMetricValue("batch-size-avg", 1.0);
    assertTaskMetricValue("offset-commit-max-time-ms", 0.0);
    assertTaskMetricValue("offset-commit-avg-time-ms", 0.0);
    assertTaskMetricValue("offset-commit-failure-percentage", 0.0);
    assertTaskMetricValue("offset-commit-success-percentage", 1.0);
    PowerMock.verifyAll();
}

Example 9 with TimestampType

use of org.apache.kafka.common.record.TimestampType in project apache-kafka-on-k8s by banzaicloud.

the class Fetcher method parseRecord.

/**
 * Parse the record entry, deserializing the key / value fields if necessary
 */
private ConsumerRecord<K, V> parseRecord(TopicPartition partition, RecordBatch batch, Record record) {
    try {
        long offset = record.offset();
        long timestamp = record.timestamp();
        TimestampType timestampType = batch.timestampType();
        Headers headers = new RecordHeaders(record.headers());
        ByteBuffer keyBytes = record.key();
        byte[] keyByteArray = keyBytes == null ? null : Utils.toArray(keyBytes);
        K key = keyBytes == null ? null : this.keyDeserializer.deserialize(partition.topic(), headers, keyByteArray);
        ByteBuffer valueBytes = record.value();
        byte[] valueByteArray = valueBytes == null ? null : Utils.toArray(valueBytes);
        V value = valueBytes == null ? null : this.valueDeserializer.deserialize(partition.topic(), headers, valueByteArray);
        return new ConsumerRecord<>(partition.topic(), partition.partition(), offset, timestamp, timestampType, record.checksumOrNull(), keyByteArray == null ? ConsumerRecord.NULL_SIZE : keyByteArray.length, valueByteArray == null ? ConsumerRecord.NULL_SIZE : valueByteArray.length, key, value, headers);
    } catch (RuntimeException e) {
        throw new SerializationException("Error deserializing key/value for partition " + partition + " at offset " + record.offset() + ". If needed, please seek past the record to continue consumption.", e);
    }
}

Example 10 with TimestampType

use of org.apache.kafka.common.record.TimestampType in project kafka by apache.

the class WorkerSinkTaskTest method testCommitWithOutOfOrderCallback.

// Verify that when commitAsync is called but the supplied callback is not called by the consumer before a
// rebalance occurs, the async callback does not reset the last committed offset from the rebalance.
// See KAFKA-5731 for more information.
@Test
public void testCommitWithOutOfOrderCallback() throws Exception {
    createTask(initialState);
    expectInitializeTask();
    expectTaskGetTopic(true);
    // iter 1
    expectPollInitialAssignment();
    // iter 2
    expectConsumerPoll(1);
    expectConversionAndTransformation(4);
    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 + 1));
    workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
    final List<TopicPartition> originalPartitions = new ArrayList<>(INITIAL_ASSIGNMENT);
    final List<TopicPartition> rebalancedPartitions = asList(TOPIC_PARTITION, TOPIC_PARTITION2, TOPIC_PARTITION3);
    final Map<TopicPartition, OffsetAndMetadata> rebalanceOffsets = new HashMap<>();
    rebalanceOffsets.put(TOPIC_PARTITION, workerCurrentOffsets.get(TOPIC_PARTITION));
    rebalanceOffsets.put(TOPIC_PARTITION2, workerCurrentOffsets.get(TOPIC_PARTITION2));
    rebalanceOffsets.put(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET));
    final Map<TopicPartition, OffsetAndMetadata> postRebalanceCurrentOffsets = new HashMap<>();
    postRebalanceCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 3));
    postRebalanceCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
    postRebalanceCurrentOffsets.put(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET + 2));
    EasyMock.expect(consumer.assignment()).andReturn(new HashSet<>(originalPartitions)).times(2);
    // iter 3 - note that we return the current offset to indicate they should be committed
    sinkTask.preCommit(workerCurrentOffsets);
    EasyMock.expectLastCall().andReturn(workerCurrentOffsets);
    // We need to delay the result of trying to commit offsets to Kafka via the consumer.commitAsync
    // method. We do this so that we can test that the callback is not called until after the rebalance
    // changes the lastCommittedOffsets. To fake this for tests we have the commitAsync build a function
    // that will call the callback with the appropriate parameters, and we'll run that function later.
    final AtomicReference<Runnable> asyncCallbackRunner = new AtomicReference<>();
    final AtomicBoolean asyncCallbackRan = new AtomicBoolean();
    consumer.commitAsync(EasyMock.eq(workerCurrentOffsets), EasyMock.anyObject());
    EasyMock.expectLastCall().andAnswer(() -> {
        // Grab the arguments passed to the consumer.commitAsync method
        final Object[] args = EasyMock.getCurrentArguments();
        @SuppressWarnings("unchecked") final Map<TopicPartition, OffsetAndMetadata> offsets = (Map<TopicPartition, OffsetAndMetadata>) args[0];
        final OffsetCommitCallback callback = (OffsetCommitCallback) args[1];
        asyncCallbackRunner.set(() -> {
            callback.onComplete(offsets, null);
            asyncCallbackRan.set(true);
        });
        return null;
    });
    // Expect the next poll to discover and perform the rebalance, THEN complete the previous callback handler,
    // and then return one record for TP1 and one for TP3.
    final AtomicBoolean rebalanced = new AtomicBoolean();
    EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(() -> {
        // Rebalance always begins with revoking current partitions ...
        rebalanceListener.getValue().onPartitionsRevoked(originalPartitions);
        // Respond to the rebalance
        Map<TopicPartition, Long> offsets = new HashMap<>();
        offsets.put(TOPIC_PARTITION, rebalanceOffsets.get(TOPIC_PARTITION).offset());
        offsets.put(TOPIC_PARTITION2, rebalanceOffsets.get(TOPIC_PARTITION2).offset());
        offsets.put(TOPIC_PARTITION3, rebalanceOffsets.get(TOPIC_PARTITION3).offset());
        sinkTaskContext.getValue().offset(offsets);
        rebalanceListener.getValue().onPartitionsAssigned(rebalancedPartitions);
        rebalanced.set(true);
        // Run the previous async commit handler
        asyncCallbackRunner.get().run();
        // And prep the two records to return
        long timestamp = RecordBatch.NO_TIMESTAMP;
        TimestampType timestampType = TimestampType.NO_TIMESTAMP_TYPE;
        List<ConsumerRecord<byte[], byte[]>> records = new ArrayList<>();
        records.add(new ConsumerRecord<>(TOPIC, PARTITION, FIRST_OFFSET + recordsReturnedTp1 + 1, timestamp, timestampType, 0, 0, RAW_KEY, RAW_VALUE, new RecordHeaders(), Optional.empty()));
        records.add(new ConsumerRecord<>(TOPIC, PARTITION3, FIRST_OFFSET + recordsReturnedTp3 + 1, timestamp, timestampType, 0, 0, RAW_KEY, RAW_VALUE, new RecordHeaders(), Optional.empty()));
        recordsReturnedTp1 += 1;
        recordsReturnedTp3 += 1;
        return new ConsumerRecords<>(Collections.singletonMap(new TopicPartition(TOPIC, PARTITION), records));
    });
    // onPartitionsRevoked
    sinkTask.preCommit(workerCurrentOffsets);
    EasyMock.expectLastCall().andReturn(workerCurrentOffsets);
    sinkTask.put(EasyMock.anyObject());
    EasyMock.expectLastCall();
    sinkTask.close(new ArrayList<>(workerCurrentOffsets.keySet()));
    EasyMock.expectLastCall();
    consumer.commitSync(workerCurrentOffsets);
    EasyMock.expectLastCall();
    // onPartitionsAssigned - step 1
    final long offsetTp1 = rebalanceOffsets.get(TOPIC_PARTITION).offset();
    final long offsetTp2 = rebalanceOffsets.get(TOPIC_PARTITION2).offset();
    final long offsetTp3 = rebalanceOffsets.get(TOPIC_PARTITION3).offset();
    EasyMock.expect(consumer.position(TOPIC_PARTITION)).andReturn(offsetTp1);
    EasyMock.expect(consumer.position(TOPIC_PARTITION2)).andReturn(offsetTp2);
    EasyMock.expect(consumer.position(TOPIC_PARTITION3)).andReturn(offsetTp3);
    EasyMock.expect(consumer.assignment()).andReturn(new HashSet<>(rebalancedPartitions)).times(6);
    // onPartitionsAssigned - step 2
    sinkTask.open(EasyMock.eq(rebalancedPartitions));
    EasyMock.expectLastCall();
    // onPartitionsAssigned - step 3 rewind
    consumer.seek(TOPIC_PARTITION, offsetTp1);
    EasyMock.expectLastCall();
    consumer.seek(TOPIC_PARTITION2, offsetTp2);
    EasyMock.expectLastCall();
    consumer.seek(TOPIC_PARTITION3, offsetTp3);
    EasyMock.expectLastCall();
    // iter 4 - note that we return the current offset to indicate they should be committed
    sinkTask.preCommit(postRebalanceCurrentOffsets);
    EasyMock.expectLastCall().andReturn(postRebalanceCurrentOffsets);
    final Capture<OffsetCommitCallback> callback = EasyMock.newCapture();
    consumer.commitAsync(EasyMock.eq(postRebalanceCurrentOffsets), EasyMock.capture(callback));
    EasyMock.expectLastCall().andAnswer(() -> {
        callback.getValue().onComplete(postRebalanceCurrentOffsets, null);
        return null;
    });
    // no actual consumer.commit() triggered
    expectConsumerPoll(1);
    sinkTask.put(EasyMock.anyObject());
    EasyMock.expectLastCall();
    PowerMock.replayAll();
    workerTask.initialize(TASK_CONFIG);
    workerTask.initializeAndStart();
    // iter 1 -- initial assignment
    workerTask.iteration();
    assertEquals(workerStartingOffsets, Whitebox.getInternalState(workerTask, "currentOffsets"));
    assertEquals(workerStartingOffsets, Whitebox.getInternalState(workerTask, "lastCommittedOffsets"));
    time.sleep(WorkerConfig.OFFSET_COMMIT_TIMEOUT_MS_DEFAULT);
    // iter 2 -- deliver 2 records
    workerTask.iteration();
    sinkTaskContext.getValue().requestCommit();
    // iter 3 -- commit in progress
    workerTask.iteration();
    assertSinkMetricValue("partition-count", 3);
    assertSinkMetricValue("sink-record-read-total", 3.0);
    assertSinkMetricValue("sink-record-send-total", 3.0);
    assertSinkMetricValue("sink-record-active-count", 4.0);
    assertSinkMetricValue("sink-record-active-count-max", 4.0);
    assertSinkMetricValue("sink-record-active-count-avg", 0.71429);
    assertSinkMetricValue("offset-commit-seq-no", 2.0);
    assertSinkMetricValue("offset-commit-completion-total", 1.0);
    assertSinkMetricValue("offset-commit-skip-total", 1.0);
    assertTaskMetricValue("status", "running");
    assertTaskMetricValue("running-ratio", 1.0);
    assertTaskMetricValue("pause-ratio", 0.0);
    assertTaskMetricValue("batch-size-max", 2.0);
    assertTaskMetricValue("batch-size-avg", 1.0);
    assertTaskMetricValue("offset-commit-max-time-ms", 0.0);
    assertTaskMetricValue("offset-commit-avg-time-ms", 0.0);
    assertTaskMetricValue("offset-commit-failure-percentage", 0.0);
    assertTaskMetricValue("offset-commit-success-percentage", 1.0);
    assertTrue(asyncCallbackRan.get());
    assertTrue(rebalanced.get());
    // Check that the offsets were not reset by the out-of-order async commit callback
    assertEquals(postRebalanceCurrentOffsets, Whitebox.getInternalState(workerTask, "currentOffsets"));
    assertEquals(rebalanceOffsets, Whitebox.getInternalState(workerTask, "lastCommittedOffsets"));
    time.sleep(WorkerConfig.OFFSET_COMMIT_TIMEOUT_MS_DEFAULT);
    sinkTaskContext.getValue().requestCommit();
    // iter 4 -- commit in progress
    workerTask.iteration();
    // Check that the offsets were not reset by the out-of-order async commit callback
    assertEquals(postRebalanceCurrentOffsets, Whitebox.getInternalState(workerTask, "currentOffsets"));
    assertEquals(postRebalanceCurrentOffsets, Whitebox.getInternalState(workerTask, "lastCommittedOffsets"));
    assertSinkMetricValue("partition-count", 3);
    assertSinkMetricValue("sink-record-read-total", 4.0);
    assertSinkMetricValue("sink-record-send-total", 4.0);
    assertSinkMetricValue("sink-record-active-count", 0.0);
    assertSinkMetricValue("sink-record-active-count-max", 4.0);
    assertSinkMetricValue("sink-record-active-count-avg", 0.5555555);
    assertSinkMetricValue("offset-commit-seq-no", 3.0);
    assertSinkMetricValue("offset-commit-completion-total", 2.0);
    assertSinkMetricValue("offset-commit-skip-total", 1.0);
    assertTaskMetricValue("status", "running");
    assertTaskMetricValue("running-ratio", 1.0);
    assertTaskMetricValue("pause-ratio", 0.0);
    assertTaskMetricValue("batch-size-max", 2.0);
    assertTaskMetricValue("batch-size-avg", 1.0);
    assertTaskMetricValue("offset-commit-max-time-ms", 0.0);
    assertTaskMetricValue("offset-commit-avg-time-ms", 0.0);
    assertTaskMetricValue("offset-commit-failure-percentage", 0.0);
    assertTaskMetricValue("offset-commit-success-percentage", 1.0);
    PowerMock.verifyAll();
}