Skip to content

Commit 7412f58

Browse files
sjvanrossumsjvanrossum
committed
Replace end offset consumer with admin client
1 parent 423a3c3 commit 7412f58

14 files changed

Lines changed: 308 additions & 203 deletions

File tree

‎buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -629,7 +629,7 @@ class BeamModulePlugin implements Plugin<Project> {
629629
def jaxb_api_version = "2.3.3"
630630
def jsr305_version = "3.0.2"
631631
def everit_json_version = "1.14.2"
632-
def kafka_version = "2.4.1"
632+
def kafka_version = "2.8.2"
633633
def log4j2_version = "2.20.0"
634634
def nemo_version = "0.1"
635635
// [bomupgrader] determined by: io.grpc:grpc-netty, consistent with: google_cloud_platform_libraries_bom

‎sdks/java/io/kafka/build.gradle‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -36,10 +36,6 @@ ext {
3636
}
3737

3838
def kafkaVersions = [
39-
'201': "2.0.1",
40-
'231': "2.3.1",
41-
'241': "2.4.1",
42-
'251': "2.5.1",
4339
'282': "2.8.2",
4440
'312': "3.1.2",
4541
'390': "3.9.0",

‎sdks/java/io/kafka/kafka-201/build.gradle‎

Lines changed: 0 additions & 24 deletions
This file was deleted.

‎sdks/java/io/kafka/kafka-231/build.gradle‎

Lines changed: 0 additions & 24 deletions
This file was deleted.

‎sdks/java/io/kafka/kafka-241/build.gradle‎

Lines changed: 0 additions & 24 deletions
This file was deleted.

‎sdks/java/io/kafka/kafka-251/build.gradle‎

Lines changed: 0 additions & 24 deletions
This file was deleted.

‎sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,7 @@
114114
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
115115
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
116116
import org.apache.kafka.clients.CommonClientConfigs;
117+
import org.apache.kafka.clients.admin.Admin;
117118
import org.apache.kafka.clients.consumer.Consumer;
118119
import org.apache.kafka.clients.consumer.ConsumerConfig;
119120
import org.apache.kafka.clients.consumer.KafkaConsumer;
@@ -604,6 +605,7 @@ public static <K, V> Read<K, V> read() {
604605
return new AutoValue_KafkaIO_Read.Builder<K, V>()
605606
.setTopics(new ArrayList<>())
606607
.setTopicPartitions(new ArrayList<>())
608+
.setAdminFactoryFn(KafkaIOUtils.KAFKA_ADMIN_FACTORY_FN)
607609
.setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
608610
.setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
609611
.setMaxNumRecords(Long.MAX_VALUE)
@@ -695,6 +697,9 @@ public abstract static class Read<K, V>
695697
@Pure
696698
public abstract @Nullable Coder<V> getValueCoder();
697699

700+
@Pure
701+
public abstract SerializableFunction<Map<String, Object>, Admin> getAdminFactoryFn();
702+
698703
@Pure
699704
public abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
700705
getConsumerFactoryFn();
@@ -778,6 +783,9 @@ abstract static class Builder<K, V> {
778783

779784
abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
780785

786+
abstract Builder<K, V> setAdminFactoryFn(
787+
SerializableFunction<Map<String, Object>, Admin> adminFactoryFn);
788+
781789
abstract Builder<K, V> setConsumerFactoryFn(
782790
SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
783791

@@ -861,6 +869,7 @@ static <K, V> void setupExternalBuilder(
861869

862870
// Set required defaults
863871
builder.setTopicPartitions(Collections.emptyList());
872+
builder.setAdminFactoryFn(KafkaIOUtils.KAFKA_ADMIN_FACTORY_FN);
864873
builder.setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN);
865874
if (config.maxReadTime != null) {
866875
builder.setMaxReadTime(Duration.standardSeconds(config.maxReadTime));
@@ -1306,6 +1315,15 @@ public Read<K, V> withValueDeserializerProviderAndCoder(
13061315
.build();
13071316
}
13081317

1318+
/**
1319+
* A factory to create Kafka {@link Admin} from offset consumer configuration. This is useful
1320+
* for supporting another version of Kafka admin. Default is {@link Admin#create(Map)}.
1321+
*/
1322+
public Read<K, V> withAdminFactoryFn(
1323+
SerializableFunction<Map<String, Object>, Admin> adminFactoryFn) {
1324+
return toBuilder().setAdminFactoryFn(adminFactoryFn).build();
1325+
}
1326+
13091327
/**
13101328
* A factory to create Kafka {@link Consumer} from consumer configuration. This is useful for
13111329
* supporting another version of Kafka consumer. Default is {@link KafkaConsumer}.
@@ -1963,6 +1981,7 @@ public PCollection<KafkaRecord<K, V>> expand(PBegin input) {
19631981
ReadSourceDescriptors.<K, V>read()
19641982
.withConsumerConfigOverrides(kafkaRead.getConsumerConfig())
19651983
.withOffsetConsumerConfigOverrides(kafkaRead.getOffsetConsumerConfig())
1984+
.withAdminFactoryFn(kafkaRead.getAdminFactoryFn())
19661985
.withConsumerFactoryFn(kafkaRead.getConsumerFactoryFn())
19671986
.withKeyDeserializerProviderAndCoder(
19681987
kafkaRead.getKeyDeserializerProvider(), keyCoder)
@@ -2479,6 +2498,9 @@ public abstract static class ReadSourceDescriptors<K, V>
24792498
@Pure
24802499
abstract @Nullable Coder<V> getValueCoder();
24812500

2501+
@Pure
2502+
abstract SerializableFunction<Map<String, Object>, Admin> getAdminFactoryFn();
2503+
24822504
@Pure
24832505
abstract SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
24842506
getConsumerFactoryFn();
@@ -2529,6 +2551,9 @@ abstract static class Builder<K, V> {
25292551
abstract ReadSourceDescriptors.Builder<K, V> setOffsetConsumerConfig(
25302552
@Nullable Map<String, Object> offsetConsumerConfig);
25312553

2554+
abstract ReadSourceDescriptors.Builder<K, V> setAdminFactoryFn(
2555+
SerializableFunction<Map<String, Object>, Admin> adminFactoryFn);
2556+
25322557
abstract ReadSourceDescriptors.Builder<K, V> setConsumerFactoryFn(
25332558
SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn);
25342559

@@ -2583,6 +2608,7 @@ abstract ReadSourceDescriptors.Builder<K, V> setBadRecordErrorHandler(
25832608

25842609
public static <K, V> ReadSourceDescriptors<K, V> read() {
25852610
return new AutoValue_KafkaIO_ReadSourceDescriptors.Builder<K, V>()
2611+
.setAdminFactoryFn(KafkaIOUtils.KAFKA_ADMIN_FACTORY_FN)
25862612
.setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN)
25872613
.setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES)
25882614
.setCommitOffsetEnabled(false)
@@ -2683,6 +2709,15 @@ public ReadSourceDescriptors<K, V> withValueDeserializerProviderAndCoder(
26832709
.build();
26842710
}
26852711

2712+
/**
2713+
* A factory to create Kafka {@link Admin} from offset consumer configuration. This is useful
2714+
* for supporting another version of Kafka admin. Default is {@link Admin#create(Map)}.
2715+
*/
2716+
public ReadSourceDescriptors<K, V> withAdminFactoryFn(
2717+
SerializableFunction<Map<String, Object>, Admin> adminFactoryFn) {
2718+
return toBuilder().setAdminFactoryFn(adminFactoryFn).build();
2719+
}
2720+
26862721
/**
26872722
* A factory to create Kafka {@link Consumer} from consumer configuration. This is useful for
26882723
* supporting another version of Kafka consumer. Default is {@link KafkaConsumer}.

‎sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIOReadImplementationCompatibility.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@ enum KafkaIOReadProperties {
8383
TOPIC_PATTERN,
8484
KEY_CODER,
8585
VALUE_CODER,
86+
ADMIN_FACTORY_FN,
8687
CONSUMER_FACTORY_FN,
8788
WATERMARK_FN(LEGACY),
8889
MAX_NUM_RECORDS(LEGACY) {

‎sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIOUtils.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import org.apache.beam.sdk.transforms.SerializableFunction;
2929
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
3030
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.Longs;
31+
import org.apache.kafka.clients.admin.Admin;
3132
import org.apache.kafka.clients.consumer.Consumer;
3233
import org.apache.kafka.clients.consumer.ConsumerConfig;
3334
import org.apache.kafka.clients.consumer.KafkaConsumer;
@@ -76,6 +77,10 @@ public final class KafkaIOUtils {
7677
// lets allow these, applications can have better resume point for restarts.
7778
);
7879

80+
// Default Kafka Admin supplier.
81+
static final SerializableFunction<Map<String, Object>, Admin> KAFKA_ADMIN_FACTORY_FN =
82+
Admin::create;
83+
7984
// default Kafka 0.9 Consumer supplier.
8085
static final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>
8186
KAFKA_CONSUMER_FACTORY_FN = KafkaConsumer::new;

‎sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaDoFn.java‎

Lines changed: 35 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -67,10 +67,14 @@
6767
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
6868
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
6969
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.Closeables;
70+
import org.apache.kafka.clients.admin.Admin;
71+
import org.apache.kafka.clients.admin.ListOffsetsResult;
72+
import org.apache.kafka.clients.admin.OffsetSpec;
7073
import org.apache.kafka.clients.consumer.Consumer;
7174
import org.apache.kafka.clients.consumer.ConsumerConfig;
7275
import org.apache.kafka.clients.consumer.ConsumerRecord;
7376
import org.apache.kafka.clients.consumer.ConsumerRecords;
77+
import org.apache.kafka.common.KafkaFuture;
7478
import org.apache.kafka.common.PartitionInfo;
7579
import org.apache.kafka.common.TopicPartition;
7680
import org.apache.kafka.common.config.ConfigDef;
@@ -202,6 +206,8 @@ private static class Bounded<K, V> extends ReadFromKafkaDoFn<K, V> {
202206
private ReadFromKafkaDoFn(
203207
ReadSourceDescriptors<K, V> transform,
204208
TupleTag<KV<KafkaSourceDescriptor, KafkaRecord<K, V>>> recordTag) {
209+
final SerializableFunction<Map<String, Object>, Admin> adminFactoryFn =
210+
transform.getAdminFactoryFn();
205211
final SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> consumerFactoryFn =
206212
transform.getConsumerFactoryFn();
207213
this.consumerConfig = transform.getConsumerConfig();
@@ -250,15 +256,14 @@ public MovingAvg load(KafkaSourceDescriptor kafkaSourceDescriptor)
250256
public KafkaLatestOffsetEstimator load(
251257
final KafkaSourceDescriptor sourceDescriptor) {
252258
LOG.info(
253-
"Creating Kafka consumer for offset estimation for {}",
259+
"Creating Kafka admin for offset estimation for {}",
254260
sourceDescriptor);
255261
final Map<String, Object> config =
256262
KafkaIOUtils.overrideBootstrapServersConfig(
257263
consumerConfig, sourceDescriptor);
258-
final Consumer<byte[], byte[]> consumer =
259-
consumerFactoryFn.apply(config);
264+
final Admin admin = adminFactoryFn.apply(config);
260265
return new KafkaLatestOffsetEstimator(
261-
consumer, sourceDescriptor.getTopicPartition());
266+
admin, sourceDescriptor.getTopicPartition());
262267
}
263268
}));
264269
this.pollConsumerCacheSupplier =
@@ -355,37 +360,43 @@ public Consumer<byte[], byte[]> load(
355360
*/
356361
private static class KafkaLatestOffsetEstimator
357362
implements GrowableOffsetRangeTracker.RangeEndEstimator, Closeable {
358-
private final Consumer<byte[], byte[]> offsetConsumer;
359-
private final Supplier<Long> offsetSupplier;
360-
361-
KafkaLatestOffsetEstimator(
362-
final Consumer<byte[], byte[]> offsetConsumer, final TopicPartition topicPartition) {
363-
this.offsetConsumer = offsetConsumer;
364-
this.offsetSupplier =
363+
private static final ListOffsetsResult.ListOffsetsResultInfo DEFAULT_RESULT =
364+
new ListOffsetsResult.ListOffsetsResultInfo(
365+
Long.MIN_VALUE, Long.MIN_VALUE, Optional.empty());
366+
367+
private final Admin admin;
368+
private final Supplier<KafkaFuture<ListOffsetsResult.ListOffsetsResultInfo>>
369+
latestOffsetFutureSupplier;
370+
private ListOffsetsResult.ListOffsetsResultInfo latestOffsetResult;
371+
372+
KafkaLatestOffsetEstimator(final Admin admin, final TopicPartition topicPartition) {
373+
this.admin = admin;
374+
this.latestOffsetFutureSupplier =
365375
new ExpiringMemoizingSerializableSupplier<>(
366-
() -> {
367-
try {
368-
return offsetConsumer
369-
.endOffsets(Collections.singleton(topicPartition))
370-
.getOrDefault(topicPartition, Long.MIN_VALUE);
371-
} catch (Throwable t) {
372-
LOG.error("Failed to get end offset for {}", topicPartition, t);
373-
return Long.MIN_VALUE;
374-
}
375-
},
376+
() ->
377+
admin
378+
.listOffsets(Collections.singletonMap(topicPartition, OffsetSpec.latest()))
379+
.partitionResult(topicPartition),
376380
Duration.ofSeconds(1),
377-
Long.MIN_VALUE,
381+
KafkaFuture.completedFuture(DEFAULT_RESULT),
378382
Duration.ZERO);
383+
this.latestOffsetResult = DEFAULT_RESULT;
379384
}
380385

381386
@Override
382387
public long estimate() {
383-
return offsetSupplier.get();
388+
try {
389+
latestOffsetResult = latestOffsetFutureSupplier.get().getNow(latestOffsetResult);
390+
} catch (Throwable t) {
391+
LOG.error("Failed to get latest offset", t);
392+
}
393+
394+
return latestOffsetResult.offset();
384395
}
385396

386397
@Override
387398
public void close() {
388-
offsetConsumer.close();
399+
admin.close();
389400
}
390401
}
391402

0 commit comments

Comments
 (0)