diff --git a/apm-sniffer/apm-sdk-plugin/kafka-plugin/src/main/java/org/apache/skywalking/apm/plugin/kafka/CallbackAdapterInterceptor.java b/apm-sniffer/apm-sdk-plugin/kafka-plugin/src/main/java/org/apache/skywalking/apm/plugin/kafka/CallbackAdapterInterceptor.java new file mode 100644 index 000000000..d296c3fec --- /dev/null +++ b/apm-sniffer/apm-sdk-plugin/kafka-plugin/src/main/java/org/apache/skywalking/apm/plugin/kafka/CallbackAdapterInterceptor.java @@ -0,0 +1,65 @@ +/* + * 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.skywalking.apm.plugin.kafka; + +import org.apache.kafka.clients.producer.Callback; +import org.apache.kafka.clients.producer.RecordMetadata; +import org.apache.skywalking.apm.agent.core.context.ContextManager; +import org.apache.skywalking.apm.agent.core.context.ContextSnapshot; +import org.apache.skywalking.apm.agent.core.context.tag.Tags; +import org.apache.skywalking.apm.agent.core.context.trace.AbstractSpan; +import org.apache.skywalking.apm.network.trace.component.ComponentsDefine; + +/** + * implements Callback and EnhancedInstance, for kafka callback in lambda expression + */ +public class CallbackAdapterInterceptor implements Callback { + + /** + * user Callback object + */ + private CallbackCache callbackCache; + + public CallbackAdapterInterceptor(CallbackCache callbackCache) { + this.callbackCache = callbackCache; + } + + @Override + public void onCompletion(RecordMetadata metadata, Exception exception) { + ContextSnapshot snapshot = callbackCache.getSnapshot(); + AbstractSpan activeSpan = ContextManager.createLocalSpan("Kafka/Producer/Callback"); + activeSpan.setComponent(ComponentsDefine.KAFKA_PRODUCER); + if (metadata != null) { + Tags.MQ_TOPIC.set(activeSpan, metadata.topic()); + } + ContextManager.continued(snapshot); + + try { + callbackCache.getCallback().onCompletion(metadata, exception); + } catch (Throwable t) { + ContextManager.activeSpan().errorOccurred().log(t); + throw t; + } finally { + if (exception != null) { + ContextManager.activeSpan().errorOccurred().log(exception); + } + ContextManager.stopSpan(); + } + } +} \ No newline at end of file diff --git a/apm-sniffer/apm-sdk-plugin/kafka-plugin/src/main/java/org/apache/skywalking/apm/plugin/kafka/KafkaProducerInterceptor.java b/apm-sniffer/apm-sdk-plugin/kafka-plugin/src/main/java/org/apache/skywalking/apm/plugin/kafka/KafkaProducerInterceptor.java index 5f2260d2f..7a00a15b8 100644 --- a/apm-sniffer/apm-sdk-plugin/kafka-plugin/src/main/java/org/apache/skywalking/apm/plugin/kafka/KafkaProducerInterceptor.java +++ b/apm-sniffer/apm-sdk-plugin/kafka-plugin/src/main/java/org/apache/skywalking/apm/plugin/kafka/KafkaProducerInterceptor.java @@ -18,6 +18,7 @@ package org.apache.skywalking.apm.plugin.kafka; +import org.apache.kafka.clients.producer.Callback; import org.apache.kafka.clients.producer.ProducerRecord; import org.apache.skywalking.apm.agent.core.context.CarrierItem; import org.apache.skywalking.apm.agent.core.context.ContextCarrier; @@ -30,7 +31,6 @@ import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedI import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.InstanceMethodsAroundInterceptor; import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult; import org.apache.skywalking.apm.network.trace.component.ComponentsDefine; - import java.lang.reflect.Method; public class KafkaProducerInterceptor implements InstanceMethodsAroundInterceptor { @@ -40,14 +40,14 @@ public class KafkaProducerInterceptor implements InstanceMethodsAroundIntercepto @Override public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class[] argumentsTypes, - MethodInterceptResult result) throws Throwable { + MethodInterceptResult result) throws Throwable { ContextCarrier contextCarrier = new ContextCarrier(); ProducerRecord record = (ProducerRecord) allArguments[0]; String topicName = record.topic(); AbstractSpan activeSpan = ContextManager.createExitSpan(OPERATE_NAME_PREFIX + topicName + PRODUCER_OPERATE_NAME_SUFFIX, contextCarrier, (String) objInst - .getSkyWalkingDynamicField()); + .getSkyWalkingDynamicField()); Tags.MQ_BROKER.set(activeSpan, (String) objInst.getSkyWalkingDynamicField()); Tags.MQ_TOPIC.set(activeSpan, topicName); @@ -59,27 +59,42 @@ public class KafkaProducerInterceptor implements InstanceMethodsAroundIntercepto next = next.next(); record.headers().add(next.getHeadKey(), next.getHeadValue().getBytes()); } - EnhancedInstance callbackInstance = (EnhancedInstance) allArguments[1]; - if (null != callbackInstance) { - ContextSnapshot snapshot = ContextManager.capture(); - if (null != snapshot) { - CallbackCache cache = new CallbackCache(); - cache.setSnapshot(snapshot); - callbackInstance.setSkyWalkingDynamicField(cache); + + //when use lambda expression, not to generate inner class, + // and not to trigger kafka CallBack class define, so allArguments[1] can't to cast EnhancedInstance + Object shouldCallbackInstance = allArguments[1]; + if (null != shouldCallbackInstance) { + if (shouldCallbackInstance instanceof EnhancedInstance) { + EnhancedInstance callbackInstance = (EnhancedInstance) shouldCallbackInstance; + ContextSnapshot snapshot = ContextManager.capture(); + if (null != snapshot) { + CallbackCache cache = new CallbackCache(); + cache.setSnapshot(snapshot); + callbackInstance.setSkyWalkingDynamicField(cache); + } + } else if (shouldCallbackInstance instanceof Callback) { + Callback callback = (Callback) shouldCallbackInstance; + ContextSnapshot snapshot = ContextManager.capture(); + if (null != snapshot) { + CallbackCache cache = new CallbackCache(); + cache.setSnapshot(snapshot); + cache.setCallback(callback); + allArguments[1] = new CallbackAdapterInterceptor(cache); + } } } } @Override public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class[] argumentsTypes, - Object ret) throws Throwable { + Object ret) throws Throwable { ContextManager.stopSpan(); return ret; } @Override public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments, - Class[] argumentsTypes, Throwable t) { + Class[] argumentsTypes, Throwable t) { } } diff --git a/apm-sniffer/apm-sdk-plugin/kafka-plugin/src/main/java/org/apache/skywalking/apm/plugin/kafka/define/KafkaProducerInstrumentation.java b/apm-sniffer/apm-sdk-plugin/kafka-plugin/src/main/java/org/apache/skywalking/apm/plugin/kafka/define/KafkaProducerInstrumentation.java index a984df923..f8076e521 100644 --- a/apm-sniffer/apm-sdk-plugin/kafka-plugin/src/main/java/org/apache/skywalking/apm/plugin/kafka/define/KafkaProducerInstrumentation.java +++ b/apm-sniffer/apm-sdk-plugin/kafka-plugin/src/main/java/org/apache/skywalking/apm/plugin/kafka/define/KafkaProducerInstrumentation.java @@ -80,7 +80,7 @@ public class KafkaProducerInstrumentation extends AbstractKafkaInstrumentation { @Override public boolean isOverrideArgs() { - return false; + return true; } } }; diff --git a/apm-sniffer/apm-sdk-plugin/kafka-plugin/src/test/java/org/apache/skywalking/apm/plugin/kafka/KafkaProducerInterceptorTest.java b/apm-sniffer/apm-sdk-plugin/kafka-plugin/src/test/java/org/apache/skywalking/apm/plugin/kafka/KafkaProducerInterceptorTest.java index 753d74910..daf231195 100644 --- a/apm-sniffer/apm-sdk-plugin/kafka-plugin/src/test/java/org/apache/skywalking/apm/plugin/kafka/KafkaProducerInterceptorTest.java +++ b/apm-sniffer/apm-sdk-plugin/kafka-plugin/src/test/java/org/apache/skywalking/apm/plugin/kafka/KafkaProducerInterceptorTest.java @@ -19,6 +19,8 @@ package org.apache.skywalking.apm.plugin.kafka; import java.util.List; + +import org.apache.kafka.clients.producer.Callback; import org.apache.kafka.clients.producer.ProducerRecord; import org.apache.skywalking.apm.agent.core.context.trace.AbstractTracingSpan; import org.apache.skywalking.apm.agent.core.context.trace.SpanLayer; @@ -89,10 +91,14 @@ public class KafkaProducerInterceptorTest { @Before public void setUp() { producerInterceptor = new KafkaProducerInterceptor(); - - arguments = new Object[] { - messageInstance, - null + //when use lambda expression not to generate inner class,and not to trigger class define. + Callback callback = (metadata, exception) -> { + if (null != metadata) { + } + }; + arguments = new Object[]{ + messageInstance, + callback }; argumentType = new Class[] {ProducerRecord.class}; } @@ -112,6 +118,26 @@ public class KafkaProducerInterceptorTest { assertMessageSpan(spans.get(0)); } + @Test + public void testSendMessageAndCallBack() throws Throwable { + producerInterceptor.beforeMethod(kafkaProducerInstance, null, arguments, argumentType, null); + Object argument = arguments[1]; + if (null != argument) { + Callback callback = (Callback) argument; + callback.onCompletion(null, null); + } + producerInterceptor.afterMethod(kafkaProducerInstance, null, arguments, argumentType, null); + + List traceSegmentList = segmentStorage.getTraceSegments(); + assertThat(traceSegmentList.size(), is(1)); + + TraceSegment segment = traceSegmentList.get(0); + List spans = SegmentHelper.getSpans(segment); + assertThat(spans.size(), is(2)); + + assertCallbackSpan(spans.get(0)); + } + private void assertMessageSpan(AbstractTracingSpan span) { SpanAssert.assertTag(span, 0, "localhost:9092"); SpanAssert.assertTag(span, 1, "test"); @@ -119,4 +145,9 @@ public class KafkaProducerInterceptorTest { SpanAssert.assertLayer(span, SpanLayer.MQ); assertThat(span.getOperationName(), is("Kafka/test/Producer")); } + + private void assertCallbackSpan(AbstractTracingSpan span) { + SpanAssert.assertComponent(span, KAFKA_PRODUCER); + assertThat(span.getOperationName(), is("Kafka/Producer/Callback")); + } } \ No newline at end of file diff --git a/test/plugin/scenarios/kafka-scenario/config/expectedData.yaml b/test/plugin/scenarios/kafka-scenario/config/expectedData.yaml index 43995a17e..2e5da2a15 100644 --- a/test/plugin/scenarios/kafka-scenario/config/expectedData.yaml +++ b/test/plugin/scenarios/kafka-scenario/config/expectedData.yaml @@ -37,6 +37,26 @@ segmentItems: parentSpanId: 1, parentTraceSegmentId: not null, parentServiceInstance: not null, parentService: 'kafka-scenario', traceId: not null} skipAnalysis: 'false' + - segmentId: not null + spans: + - operationName: Kafka/Producer/Callback + operationId: 0 + parentSpanId: -1 + spanId: 0 + spanLayer: Unknown + startTime: nq 0 + endTime: nq 0 + componentId: 40 + isError: false + spanType: Local + peer: '' + tags: + - {key: mq.topic, value: test2} + refs: + - {parentEndpoint: /case/kafka-case, networkAddress: '', refType: CrossThread, + parentSpanId: 2, parentTraceSegmentId: not null, parentServiceInstance: not + null, parentService: 'kafka-scenario', traceId: not null} + skipAnalysis: 'false' - segmentId: not null spans: - operationName: Kafka/test/Producer @@ -54,6 +74,21 @@ segmentItems: - {key: mq.broker, value: 'kafka-server:9092'} - {key: mq.topic, value: test} skipAnalysis: 'false' + - operationName: Kafka/test2/Producer + operationId: 0 + parentSpanId: 0 + spanId: 2 + spanLayer: MQ + startTime: nq 0 + endTime: nq 0 + componentId: 40 + isError: false + spanType: Exit + peer: kafka-server:9092 + tags: + - {key: mq.broker, value: 'kafka-server:9092'} + - {key: mq.topic, value: test2} + skipAnalysis: 'false' - operationName: /case/kafka-case operationId: 0 parentSpanId: -1 diff --git a/test/plugin/scenarios/kafka-scenario/src/main/java/test/org/apache/skywalking/apm/testcase/kafka/controller/CaseController.java b/test/plugin/scenarios/kafka-scenario/src/main/java/test/org/apache/skywalking/apm/testcase/kafka/controller/CaseController.java index fa68d3a8b..e787c01bc 100644 --- a/test/plugin/scenarios/kafka-scenario/src/main/java/test/org/apache/skywalking/apm/testcase/kafka/controller/CaseController.java +++ b/test/plugin/scenarios/kafka-scenario/src/main/java/test/org/apache/skywalking/apm/testcase/kafka/controller/CaseController.java @@ -53,12 +53,14 @@ public class CaseController { private String bootstrapServers; private String topicName; + private String topicName2; private static volatile boolean KAFKA_STATUS = false; @PostConstruct private void setUp() { topicName = "test"; + topicName2 = "test2"; new CheckKafkaProducerThread(bootstrapServers).start(); } @@ -74,6 +76,13 @@ public class CaseController { logger.info("send success metadata={}", metadata); } }); + + ProducerRecord record2 = new ProducerRecord(topicName2, "testKey", Integer.toString(1)); + record2.headers().add("TEST", "TEST".getBytes()); + Callback callback2 = (metadata, exception) -> { + logger.info("send success metadata={}", metadata); + }; + producer.send(record2, callback2); }, bootstrapServers); Thread thread = new ConsumerThread(); thread.start(); @@ -134,14 +143,12 @@ public class CaseController { ProducerRecord record = new ProducerRecord("check", "checkKey", Integer .toString(1)); record.headers().add("CHECK", "CHECK".getBytes()); - producer.send(record, new Callback() { - @Override - public void onCompletion(RecordMetadata recordMetadata, Exception e) { - if (isNull(e)) { - KAFKA_STATUS = true; - } + Callback callback = (metadata, e) -> { + if (isNull(e)) { + KAFKA_STATUS = true; } - }); + }; + producer.send(record, callback); }, bootstrapServers); } catch (Exception e) { logger.error("check " + bootstrapServers + " " + e.getMessage(), e);