/
ConsumerRecords.java
129 lines (111 loc) · 4.59 KB
/
ConsumerRecords.java
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
/*
* 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.clients.consumer;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.utils.AbstractIterator;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
/**
* A container that holds the list {@link ConsumerRecord} per partition for a
* particular topic. There is one {@link ConsumerRecord} list for every topic
* partition returned by a {@link Consumer#poll(java.time.Duration)} operation.
*/
public class ConsumerRecords<K, V> implements Iterable<ConsumerRecord<K, V>> {
public static final ConsumerRecords<Object, Object> EMPTY = new ConsumerRecords<>(Collections.emptyMap());
private final Map<TopicPartition, List<ConsumerRecord<K, V>>> records;
public ConsumerRecords(Map<TopicPartition, List<ConsumerRecord<K, V>>> records) {
this.records = records;
}
/**
* Get just the records for the given partition
*
* @param partition The partition to get records for
*/
public List<ConsumerRecord<K, V>> records(TopicPartition partition) {
List<ConsumerRecord<K, V>> recs = this.records.get(partition);
if (recs == null)
return Collections.emptyList();
else
return Collections.unmodifiableList(recs);
}
/**
* Get just the records for the given topic
*/
public Iterable<ConsumerRecord<K, V>> records(String topic) {
if (topic == null)
throw new IllegalArgumentException("Topic must be non-null.");
List<List<ConsumerRecord<K, V>>> recs = new ArrayList<>();
for (Map.Entry<TopicPartition, List<ConsumerRecord<K, V>>> entry : records.entrySet()) {
if (entry.getKey().topic().equals(topic))
recs.add(entry.getValue());
}
return new ConcatenatedIterable<>(recs);
}
/**
* Get the partitions which have records contained in this record set.
* @return the set of partitions with data in this record set (may be empty if no data was returned)
*/
public Set<TopicPartition> partitions() {
return Collections.unmodifiableSet(records.keySet());
}
@Override
public Iterator<ConsumerRecord<K, V>> iterator() {
return new ConcatenatedIterable<>(records.values()).iterator();
}
/**
* The number of records for all topics
*/
public int count() {
int count = 0;
for (List<ConsumerRecord<K, V>> recs: this.records.values())
count += recs.size();
return count;
}
private static class ConcatenatedIterable<K, V> implements Iterable<ConsumerRecord<K, V>> {
private final Iterable<? extends Iterable<ConsumerRecord<K, V>>> iterables;
public ConcatenatedIterable(Iterable<? extends Iterable<ConsumerRecord<K, V>>> iterables) {
this.iterables = iterables;
}
@Override
public Iterator<ConsumerRecord<K, V>> iterator() {
return new AbstractIterator<ConsumerRecord<K, V>>() {
final Iterator<? extends Iterable<ConsumerRecord<K, V>>> iters = iterables.iterator();
Iterator<ConsumerRecord<K, V>> current;
protected ConsumerRecord<K, V> makeNext() {
while (current == null || !current.hasNext()) {
if (iters.hasNext())
current = iters.next().iterator();
else
return allDone();
}
return current.next();
}
};
}
}
public boolean isEmpty() {
return records.isEmpty();
}
@SuppressWarnings("unchecked")
public static <K, V> ConsumerRecords<K, V> empty() {
return (ConsumerRecords<K, V>) EMPTY;
}
}