Skip to content

Commit 3d8dcce

Browse files
garyrussellgaryrussell
authored
GH-2359: Initial AOT Support for Native
Resolves #2359 This is sufficient for the basic Kafka smoke test. * Remove duplicates; register annotations properly; remove avro from factories. * Remove `@MessageMapping`, `@EnableKafka`. * Fix Spring proxies. * Add `KafkaAvroBeanRegistrationAotProcessor` to detect Avro types. * Use `registerTypeIfPresent()` for Kafka Streams. * Polish Avro detection. * Polish use of MergedAnnotations. * Kafka Avro AOT Processor polishing; no need to register for `@KafkaListener` only `GenericMessageListener`. * Fix `BatchMessageListener` (`List<ConsumerRecord<Foo, Bar>>`).
1 parent eb7648a commit 3d8dcce

File tree

4 files changed

+371
-0
lines changed

4 files changed

+371
-0
lines changed

‎spring-kafka/src/main/java/org/springframework/kafka/aot/KafkaAvroBeanRegistrationAotProcessor.java

+140
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
/*
2+
* Copyright 2022 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.kafka.aot;
18+
19+
import java.lang.reflect.ParameterizedType;
20+
import java.lang.reflect.Type;
21+
import java.util.HashSet;
22+
import java.util.List;
23+
import java.util.Set;
24+
25+
import org.apache.kafka.clients.consumer.ConsumerRecord;
26+
27+
import org.springframework.aot.hint.MemberCategory;
28+
import org.springframework.aot.hint.ReflectionHints;
29+
import org.springframework.beans.factory.aot.BeanRegistrationAotContribution;
30+
import org.springframework.beans.factory.aot.BeanRegistrationAotProcessor;
31+
import org.springframework.beans.factory.support.RegisteredBean;
32+
import org.springframework.core.ResolvableType;
33+
import org.springframework.core.annotation.MergedAnnotations;
34+
import org.springframework.kafka.listener.GenericMessageListener;
35+
import org.springframework.lang.Nullable;
36+
import org.springframework.util.ClassUtils;
37+
import org.springframework.util.ReflectionUtils;
38+
39+
/**
40+
* Detect and register Avro types for Apache Kafka listeners.
41+
*
42+
* @author Gary Russell
43+
* @since 3.0
44+
*
45+
*/
46+
public class KafkaAvroBeanRegistrationAotProcessor implements BeanRegistrationAotProcessor {
47+
48+
private static final String CONSUMER_RECORD_CLASS_NAME = ConsumerRecord.class.getName();
49+
50+
private static final String CONSUMER_RECORDS_CLASS_NAME = ConsumerRecord.class.getName();
51+
52+
private static final String AVRO_GENERATED_CLASS_NAME = "org.apache.avro.specific.AvroGenerated";
53+
54+
private static final boolean AVRO_PRESENT = ClassUtils.isPresent(AVRO_GENERATED_CLASS_NAME, null);
55+
56+
@Override
57+
@Nullable
58+
public BeanRegistrationAotContribution processAheadOfTime(RegisteredBean registeredBean) {
59+
if (!AVRO_PRESENT) {
60+
return null;
61+
}
62+
Class<?> beanType = registeredBean.getBeanClass();
63+
if (!isListener(beanType)) {
64+
return null;
65+
}
66+
Set<Class<?>> avroTypes = new HashSet<>();
67+
if (GenericMessageListener.class.isAssignableFrom(beanType)) {
68+
ReflectionUtils.doWithMethods(beanType, method -> {
69+
Type[] types = method.getGenericParameterTypes();
70+
if (types.length > 0) {
71+
ResolvableType resolvableType = ResolvableType.forType(types[0]);
72+
if (List.class.equals(resolvableType.getRawClass())) {
73+
resolvableType = resolvableType.getGeneric(0);
74+
}
75+
Class<?> keyType = resolvableType.resolveGeneric(0);
76+
Class<?> valueType = resolvableType.resolveGeneric(1);
77+
checkType(keyType, avroTypes);
78+
checkType(valueType, avroTypes);
79+
}
80+
}, method -> method.getName().equals("onMessage"));
81+
}
82+
if (avroTypes.size() > 0) {
83+
return (generationContext, beanRegistrationCode) -> {
84+
ReflectionHints reflectionHints = generationContext.getRuntimeHints().reflection();
85+
avroTypes.forEach(type -> reflectionHints.registerType(type,
86+
builder -> builder.withMembers(MemberCategory.INVOKE_DECLARED_CONSTRUCTORS,
87+
MemberCategory.INVOKE_PUBLIC_METHODS)));
88+
};
89+
}
90+
return null;
91+
}
92+
93+
private static boolean isListener(Class<?> beanType) {
94+
return GenericMessageListener.class.isAssignableFrom(beanType);
95+
}
96+
97+
private static void checkType(@Nullable Type paramType, Set<Class<?>> avroTypes) {
98+
if (paramType == null) {
99+
return;
100+
}
101+
boolean container = isContainer(paramType);
102+
if (!container && paramType instanceof Class) {
103+
MergedAnnotations mergedAnnotations = MergedAnnotations.from((Class<?>) paramType);
104+
if (mergedAnnotations.isPresent(AVRO_GENERATED_CLASS_NAME)) {
105+
avroTypes.add((Class<?>) paramType);
106+
}
107+
}
108+
else if (container) {
109+
if (paramType instanceof ParameterizedType) {
110+
Type[] generics = ((ParameterizedType) paramType).getActualTypeArguments();
111+
if (generics.length > 0) {
112+
checkAvro(generics[0], avroTypes);
113+
}
114+
if (generics.length == 2) {
115+
checkAvro(generics[1], avroTypes);
116+
}
117+
}
118+
}
119+
}
120+
121+
private static void checkAvro(@Nullable Type generic, Set<Class<?>> avroTypes) {
122+
if (generic instanceof Class) {
123+
MergedAnnotations methodAnnotations = MergedAnnotations.from((Class<?>) generic);
124+
if (methodAnnotations.isPresent(AVRO_GENERATED_CLASS_NAME)) {
125+
avroTypes.add((Class<?>) generic);
126+
}
127+
}
128+
}
129+
130+
private static boolean isContainer(Type paramType) {
131+
if (paramType instanceof ParameterizedType) {
132+
Type rawType = ((ParameterizedType) paramType).getRawType();
133+
return (rawType.equals(List.class))
134+
|| rawType.getTypeName().equals(CONSUMER_RECORD_CLASS_NAME)
135+
|| rawType.getTypeName().equals(CONSUMER_RECORDS_CLASS_NAME);
136+
}
137+
return false;
138+
}
139+
140+
}

‎spring-kafka/src/main/java/org/springframework/kafka/aot/KafkaRuntimeHintsRegistrar.java

+223
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,223 @@
1+
/*
2+
* Copyright 2022 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.kafka.aot;
18+
19+
import java.util.stream.Stream;
20+
import java.util.zip.CRC32C;
21+
22+
import org.apache.kafka.clients.admin.NewTopic;
23+
import org.apache.kafka.clients.consumer.Consumer;
24+
import org.apache.kafka.clients.consumer.CooperativeStickyAssignor;
25+
import org.apache.kafka.clients.consumer.RangeAssignor;
26+
import org.apache.kafka.clients.consumer.RoundRobinAssignor;
27+
import org.apache.kafka.clients.consumer.StickyAssignor;
28+
import org.apache.kafka.clients.producer.Producer;
29+
import org.apache.kafka.clients.producer.RoundRobinPartitioner;
30+
import org.apache.kafka.clients.producer.UniformStickyPartitioner;
31+
import org.apache.kafka.clients.producer.internals.DefaultPartitioner;
32+
import org.apache.kafka.common.message.CreateTopicsRequestData.CreatableTopic;
33+
import org.apache.kafka.common.protocol.Message;
34+
import org.apache.kafka.common.serialization.ByteArrayDeserializer;
35+
import org.apache.kafka.common.serialization.ByteArraySerializer;
36+
import org.apache.kafka.common.serialization.ByteBufferDeserializer;
37+
import org.apache.kafka.common.serialization.ByteBufferSerializer;
38+
import org.apache.kafka.common.serialization.BytesDeserializer;
39+
import org.apache.kafka.common.serialization.BytesSerializer;
40+
import org.apache.kafka.common.serialization.DoubleDeserializer;
41+
import org.apache.kafka.common.serialization.DoubleSerializer;
42+
import org.apache.kafka.common.serialization.FloatDeserializer;
43+
import org.apache.kafka.common.serialization.FloatSerializer;
44+
import org.apache.kafka.common.serialization.IntegerDeserializer;
45+
import org.apache.kafka.common.serialization.IntegerSerializer;
46+
import org.apache.kafka.common.serialization.ListDeserializer;
47+
import org.apache.kafka.common.serialization.ListSerializer;
48+
import org.apache.kafka.common.serialization.LongDeserializer;
49+
import org.apache.kafka.common.serialization.LongSerializer;
50+
import org.apache.kafka.common.serialization.Serdes;
51+
import org.apache.kafka.common.serialization.StringDeserializer;
52+
import org.apache.kafka.common.serialization.StringSerializer;
53+
import org.apache.kafka.common.utils.AppInfoParser.AppInfo;
54+
import org.apache.kafka.common.utils.ImplicitLinkedHashCollection;
55+
56+
import org.springframework.aop.framework.AopProxyUtils;
57+
import org.springframework.aot.hint.MemberCategory;
58+
import org.springframework.aot.hint.ReflectionHints;
59+
import org.springframework.aot.hint.RuntimeHints;
60+
import org.springframework.aot.hint.RuntimeHintsRegistrar;
61+
import org.springframework.aot.hint.support.RuntimeHintsUtils;
62+
import org.springframework.kafka.annotation.KafkaBootstrapConfiguration;
63+
import org.springframework.kafka.annotation.KafkaListener;
64+
import org.springframework.kafka.annotation.KafkaListenerAnnotationBeanPostProcessor;
65+
import org.springframework.kafka.annotation.KafkaListeners;
66+
import org.springframework.kafka.annotation.PartitionOffset;
67+
import org.springframework.kafka.annotation.TopicPartition;
68+
import org.springframework.kafka.config.AbstractKafkaListenerContainerFactory;
69+
import org.springframework.kafka.config.ConcurrentKafkaListenerContainerFactory;
70+
import org.springframework.kafka.config.KafkaListenerContainerFactory;
71+
import org.springframework.kafka.config.KafkaListenerEndpointRegistry;
72+
import org.springframework.kafka.core.ConsumerFactory;
73+
import org.springframework.kafka.core.DefaultKafkaConsumerFactory;
74+
import org.springframework.kafka.core.DefaultKafkaProducerFactory;
75+
import org.springframework.kafka.core.KafkaAdmin;
76+
import org.springframework.kafka.core.KafkaOperations;
77+
import org.springframework.kafka.core.KafkaResourceFactory;
78+
import org.springframework.kafka.core.KafkaTemplate;
79+
import org.springframework.kafka.core.ProducerFactory;
80+
import org.springframework.kafka.listener.ConsumerProperties;
81+
import org.springframework.kafka.listener.ContainerProperties;
82+
import org.springframework.kafka.support.LoggingProducerListener;
83+
import org.springframework.kafka.support.ProducerListener;
84+
import org.springframework.kafka.support.serializer.DelegatingByTopicDeserializer;
85+
import org.springframework.kafka.support.serializer.DelegatingByTypeSerializer;
86+
import org.springframework.kafka.support.serializer.DelegatingDeserializer;
87+
import org.springframework.kafka.support.serializer.DelegatingSerializer;
88+
import org.springframework.kafka.support.serializer.ErrorHandlingDeserializer;
89+
import org.springframework.kafka.support.serializer.JsonDeserializer;
90+
import org.springframework.kafka.support.serializer.JsonSerializer;
91+
import org.springframework.kafka.support.serializer.ParseStringDeserializer;
92+
import org.springframework.kafka.support.serializer.StringOrBytesSerializer;
93+
import org.springframework.kafka.support.serializer.ToStringSerializer;
94+
import org.springframework.lang.Nullable;
95+
96+
/**
97+
* {@link RuntimeHintsRegistrar} for Spring for Apache Kafka.
98+
*
99+
* @author Gary Russell
100+
* @since 3.0
101+
*
102+
*/
103+
public class KafkaRuntimeHintsRegistrar implements RuntimeHintsRegistrar {
104+
105+
@Override
106+
public void registerHints(RuntimeHints hints, @Nullable ClassLoader classLoader) {
107+
RuntimeHintsUtils.registerAnnotation(hints, KafkaListener.class);
108+
RuntimeHintsUtils.registerAnnotation(hints, KafkaListeners.class);
109+
RuntimeHintsUtils.registerAnnotation(hints, PartitionOffset.class);
110+
RuntimeHintsUtils.registerAnnotation(hints, TopicPartition.class);
111+
ReflectionHints reflectionHints = hints.reflection();
112+
Stream.of(
113+
ConsumerProperties.class,
114+
ContainerProperties.class,
115+
ProducerListener.class)
116+
.forEach(type -> reflectionHints.registerType(type,
117+
builder -> builder.withMembers(MemberCategory.INVOKE_DECLARED_METHODS)));
118+
119+
Stream.of(
120+
Message.class,
121+
ImplicitLinkedHashCollection.Element.class,
122+
NewTopic.class,
123+
AbstractKafkaListenerContainerFactory.class,
124+
ConcurrentKafkaListenerContainerFactory.class,
125+
KafkaListenerContainerFactory.class,
126+
KafkaListenerEndpointRegistry.class,
127+
DefaultKafkaConsumerFactory.class,
128+
DefaultKafkaProducerFactory.class,
129+
KafkaAdmin.class,
130+
KafkaOperations.class,
131+
KafkaResourceFactory.class,
132+
KafkaTemplate.class,
133+
ProducerFactory.class,
134+
KafkaOperations.class,
135+
ConsumerFactory.class,
136+
LoggingProducerListener.class,
137+
ImplicitLinkedHashCollection.Element.class,
138+
KafkaListenerAnnotationBeanPostProcessor.class)
139+
.forEach(type -> reflectionHints.registerType(type,
140+
builder -> builder.withMembers(MemberCategory.INVOKE_DECLARED_CONSTRUCTORS,
141+
MemberCategory.INVOKE_DECLARED_METHODS,
142+
MemberCategory.INTROSPECT_PUBLIC_METHODS)));
143+
144+
Stream.of(
145+
KafkaBootstrapConfiguration.class,
146+
CreatableTopic.class,
147+
KafkaListenerEndpointRegistry.class)
148+
.forEach(type -> reflectionHints.registerType(type,
149+
builder -> builder.withMembers(MemberCategory.INVOKE_DECLARED_CONSTRUCTORS)));
150+
151+
Stream.of(
152+
AppInfo.class,
153+
// standard assignors
154+
CooperativeStickyAssignor.class,
155+
RangeAssignor.class,
156+
RoundRobinAssignor.class,
157+
StickyAssignor.class,
158+
// standard partitioners
159+
DefaultPartitioner.class,
160+
RoundRobinPartitioner.class,
161+
UniformStickyPartitioner.class,
162+
// standard serialization
163+
ByteArrayDeserializer.class,
164+
ByteArraySerializer.class,
165+
ByteBufferDeserializer.class,
166+
ByteBufferSerializer.class,
167+
BytesDeserializer.class,
168+
BytesSerializer.class,
169+
DoubleSerializer.class,
170+
DoubleDeserializer.class,
171+
FloatSerializer.class,
172+
FloatDeserializer.class,
173+
IntegerSerializer.class,
174+
IntegerDeserializer.class,
175+
ListDeserializer.class,
176+
ListSerializer.class,
177+
LongSerializer.class,
178+
LongDeserializer.class,
179+
StringDeserializer.class,
180+
StringSerializer.class,
181+
// Spring serialization
182+
DelegatingByTopicDeserializer.class,
183+
DelegatingByTypeSerializer.class,
184+
DelegatingDeserializer.class,
185+
ErrorHandlingDeserializer.class,
186+
DelegatingSerializer.class,
187+
JsonDeserializer.class,
188+
JsonSerializer.class,
189+
ParseStringDeserializer.class,
190+
StringOrBytesSerializer.class,
191+
ToStringSerializer.class,
192+
Serdes.class,
193+
Serdes.ByteArraySerde.class,
194+
Serdes.BytesSerde.class,
195+
Serdes.ByteBufferSerde.class,
196+
Serdes.DoubleSerde.class,
197+
Serdes.FloatSerde.class,
198+
Serdes.IntegerSerde.class,
199+
Serdes.LongSerde.class,
200+
Serdes.ShortSerde.class,
201+
Serdes.StringSerde.class,
202+
Serdes.UUIDSerde.class,
203+
Serdes.VoidSerde.class,
204+
CRC32C.class)
205+
.forEach(type -> reflectionHints.registerType(type, builder ->
206+
builder.withMembers(MemberCategory.INVOKE_PUBLIC_CONSTRUCTORS)));
207+
208+
hints.proxies().registerJdkProxy(AopProxyUtils.completeJdkProxyInterfaces(Consumer.class));
209+
hints.proxies().registerJdkProxy(AopProxyUtils.completeJdkProxyInterfaces(Producer.class));
210+
211+
Stream.of(
212+
"org.apache.kafka.streams.processor.internals.StreamsPartitionAssignor",
213+
"org.apache.kafka.streams.errors.DefaultProductionExceptionHandler",
214+
"org.apache.kafka.streams.processor.FailOnInvalidTimestamp",
215+
"org.apache.kafka.streams.processor.internals.assignment.HighAvailabilityTaskAssignor",
216+
"org.apache.kafka.streams.processor.internals.assignment.StickyTaskAssignor",
217+
"org.apache.kafka.streams.processor.internals.assignment.FallbackPriorTaskAssignor",
218+
"org.apache.kafka.streams.errors.LogAndFailExceptionHandler")
219+
.forEach(type -> reflectionHints.registerTypeIfPresent(classLoader, type, builder ->
220+
builder.withMembers(MemberCategory.INVOKE_PUBLIC_CONSTRUCTORS)));
221+
}
222+
223+
}

‎spring-kafka/src/main/java/org/springframework/kafka/aot/package-info.java

+6
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
/**
2+
* Provides classes to support Spring AOT.
3+
*/
4+
@org.springframework.lang.NonNullApi
5+
@org.springframework.lang.NonNullFields
6+
package org.springframework.kafka.aot;

‎spring-kafka/src/main/resources/META-INF/spring/aot.factories

+2
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
org.springframework.aot.hint.RuntimeHintsRegistrar=org.springframework.kafka.aot.KafkaRuntimeHintsRegistrar
2+
org.springframework.beans.factory.aot.BeanRegistrationAotProcessor=org.springframework.kafka.aot.KafkaAvroBeanRegistrationAotProcessor

0 commit comments

Comments
 (0)