/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.kafka.streams.processor.internals;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.common.TopicPartition;
import java.util.Collections;
import java.util.Comparator;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.Set;
/**
* A PartitionGroup is composed from a set of partitions. It also maintains the timestamp of this
* group, hence the associated task as the min timestamp across all partitions in the group.
*/
public class PartitionGroup {
private final Map<TopicPartition, RecordQueue> partitionQueues;
private final PriorityQueue<RecordQueue> queuesByTime;
public static class RecordInfo {
RecordQueue queue;
public ProcessorNode node() {
return queue.source();
}
public TopicPartition partition() {
return queue.partition();
}
RecordQueue queue() {
return queue;
}
}
// since task is thread-safe, we do not need to synchronize on local variables
private int totalBuffered;
PartitionGroup(final Map<TopicPartition, RecordQueue> partitionQueues) {
queuesByTime = new PriorityQueue<>(partitionQueues.size(), new Comparator<RecordQueue>() {
@Override
public int compare(final RecordQueue queue1, final RecordQueue queue2) {
final long time1 = queue1.timestamp();
final long time2 = queue2.timestamp();
if (time1 < time2) {
return -1;
}
if (time1 > time2) {
return 1;
}
return 0;
}
});
this.partitionQueues = partitionQueues;
totalBuffered = 0;
}
/**
* Get the next record and queue
*
* @return StampedRecord
*/
StampedRecord nextRecord(final RecordInfo info) {
StampedRecord record = null;
final RecordQueue queue = queuesByTime.poll();
if (queue != null) {
// get the first record from this queue.
record = queue.poll();
if (!queue.isEmpty()) {
queuesByTime.offer(queue);
}
}
info.queue = queue;
if (record != null) {
--totalBuffered;
}
return record;
}
/**
* Adds raw records to this partition group
*
* @param partition the partition
* @param rawRecords the raw records
* @return the queue size for the partition
*/
int addRawRecords(final TopicPartition partition, final Iterable<ConsumerRecord<byte[], byte[]>> rawRecords) {
final RecordQueue recordQueue = partitionQueues.get(partition);
final int oldSize = recordQueue.size();
final int newSize = recordQueue.addRawRecords(rawRecords);
// add this record queue to be considered for processing in the future if it was empty before
if (oldSize == 0 && newSize > 0) {
queuesByTime.offer(recordQueue);
}
totalBuffered += newSize - oldSize;
return newSize;
}
public Set<TopicPartition> partitions() {
return Collections.unmodifiableSet(partitionQueues.keySet());
}
/**
* Return the timestamp of this partition group as the smallest
* partition timestamp among all its partitions
*/
public long timestamp() {
// we should always return the smallest timestamp of all partitions
// to avoid group partition time goes backward
long timestamp = Long.MAX_VALUE;
for (final RecordQueue queue : partitionQueues.values()) {
if (timestamp > queue.timestamp()) {
timestamp = queue.timestamp();
}
}
return timestamp;
}
/**
* @throws IllegalStateException if the record's partition does not belong to this partition group
*/
int numBuffered(final TopicPartition partition) {
final RecordQueue recordQueue = partitionQueues.get(partition);
if (recordQueue == null) {
throw new IllegalStateException("Record's partition does not belong to this partition-group.");
}
return recordQueue.size();
}
int numBuffered() {
return totalBuffered;
}
public void close() {
queuesByTime.clear();
partitionQueues.clear();
}
public void clear() {
queuesByTime.clear();
for (final RecordQueue queue : partitionQueues.values()) {
queue.clear();
}
}
}