Support @KafkaPollAndInvoke annotation to replace @Trace when just using the Kafka scenario agent plugin (#5304)

apm-application-toolkit/apm-toolkit-kafka/pom.xml

@@ -0,0 +1,33 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+~ 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.
+~
+-->
+
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <parent>
+        <artifactId>apm-application-toolkit</artifactId>
+        <groupId>org.apache.skywalking</groupId>
+        <version>8.2.0-SNAPSHOT</version>
+    </parent>
+    <modelVersion>4.0.0</modelVersion>
+
+    <artifactId>apm-toolkit-kafka</artifactId>
+    <packaging>jar</packaging>
+
+    <url>http://maven.apache.org</url>
+</project>

apm-application-toolkit/apm-toolkit-kafka/src/main/java/org/apache/skywalking/apm/toolkit/kafka/KafkaPollAndInvoke.java

@@ -0,0 +1,30 @@
+/*
+ * 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.toolkit.kafka;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@Target(ElementType.METHOD)
+@Retention(RetentionPolicy.RUNTIME)
+public @interface KafkaPollAndInvoke {
+
+}

apm-application-toolkit/pom.xml

@@ -38,5 +38,6 @@
         <module>apm-toolkit-trace</module>
         <module>apm-toolkit-meter</module>
         <module>apm-toolkit-micrometer-registry</module>
+        <module>apm-toolkit-kafka</module>
     </modules>
 </project>

apm-sniffer/apm-sdk-plugin/kafka-commons/src/main/java/org/apache/skywalking/apm/plugin/kafka/define/Constants.java

@@ -21,7 +21,7 @@ package org.apache.skywalking.apm.plugin.kafka.define;
 
 public class Constants {
 
-    public static final String SPRING_KAFKA_FLAG = "SW_SPRING_KAFKA_FLAG";
+    public static final String KAFKA_FLAG = "SW_KAFKA_FLAG";
 
-    public static final String SPRING_KAFKA_POLL_AND_INVOKE_OPERATION_NAME = "/spring-kafka/pollAndInvoke";
+    public static final String KAFKA_POLL_AND_INVOKE_OPERATION_NAME = "/pollAndInvoke";
 }

apm-sniffer/apm-sdk-plugin/kafka-commons/src/main/java/org/apache/skywalking/apm/plugin/kafka/define/InterceptorMethod.java

@@ -0,0 +1,48 @@
+/*
+ * 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.define;
+
+import org.apache.skywalking.apm.agent.core.context.ContextManager;
+
+public class InterceptorMethod {
+
+    public static void beforeMethod(String operationName) {
+        ContextManager.getRuntimeContext().put(Constants.KAFKA_FLAG, new KafkaContext(operationName));
+    }
+
+    public static Object afterMethod(Object ret) {
+        KafkaContext context = (KafkaContext) ContextManager.getRuntimeContext().get(Constants.KAFKA_FLAG);
+        if (context == null) {
+            return ret;
+        }
+        if (context.isNeedStop()) {
+            ContextManager.stopSpan();
+        } else {
+            ContextManager.getRuntimeContext().remove(Constants.KAFKA_FLAG);
+        }
+        return ret;
+    }
+
+    public static void handleMethodException(Throwable t) {
+        KafkaContext context = (KafkaContext) ContextManager.getRuntimeContext().get(Constants.KAFKA_FLAG);
+        if (context != null && context.isNeedStop()) {
+            ContextManager.activeSpan().errorOccurred().log(t);
+        }
+    }
+}

apm-sniffer/apm-sdk-plugin/kafka-commons/src/main/java/org/apache/skywalking/apm/plugin/kafka/define/KafkaContext.java

@@ -19,13 +19,15 @@
 
 package org.apache.skywalking.apm.plugin.kafka.define;
 
-public class SpringKafkaContext {
+public class KafkaContext {
 
-    public SpringKafkaContext() {
-        needStop = false;
+    public KafkaContext(String operationName) {
+        this.operationName = operationName;
     }
 
-    private boolean needStop;
+    private boolean needStop = false;
+
+    private String operationName;
 
     public boolean isNeedStop() {
         return needStop;
@@ -34,4 +36,12 @@ public class SpringKafkaContext {
     public void setNeedStop(boolean needStop) {
         this.needStop = needStop;
     }
+
+    public String getOperationName() {
+        return operationName;
+    }
+
+    public void setOperationName(String operationName) {
+        this.operationName = operationName;
+    }
 }

apm-sniffer/apm-sdk-plugin/kafka-plugin/src/main/java/org/apache/skywalking/apm/plugin/kafka/KafkaConsumerInterceptor.java

@@ -32,7 +32,7 @@ import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.InstanceM
 import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult;
 import org.apache.skywalking.apm.network.trace.component.ComponentsDefine;
 import org.apache.skywalking.apm.plugin.kafka.define.Constants;
-import org.apache.skywalking.apm.plugin.kafka.define.SpringKafkaContext;
+import org.apache.skywalking.apm.plugin.kafka.define.KafkaContext;
 
 import java.lang.reflect.Method;
 import java.util.Iterator;
@@ -66,9 +66,10 @@ public class KafkaConsumerInterceptor implements InstanceMethodsAroundIntercepto
         //
         if (records.size() > 0) {
             ConsumerEnhanceRequiredInfo requiredInfo = (ConsumerEnhanceRequiredInfo) objInst.getSkyWalkingDynamicField();
-            if (ContextManager.getRuntimeContext().get(Constants.SPRING_KAFKA_FLAG) != null) {
-                ContextManager.createEntrySpan(Constants.SPRING_KAFKA_POLL_AND_INVOKE_OPERATION_NAME, null);
-                ((SpringKafkaContext) ContextManager.getRuntimeContext().get(Constants.SPRING_KAFKA_FLAG)).setNeedStop(true);
+            KafkaContext context = (KafkaContext) ContextManager.getRuntimeContext().get(Constants.KAFKA_FLAG);
+            if (context != null) {
+                ContextManager.createEntrySpan(context.getOperationName(), null);
+                context.setNeedStop(true);
             }
             String operationName = OPERATE_NAME_PREFIX + requiredInfo.getTopics() + CONSUMER_OPERATE_NAME + requiredInfo.getGroupId();
             AbstractSpan activeSpan = ContextManager.createEntrySpan(operationName, null).start(requiredInfo.getStartTime());

apm-sniffer/apm-sdk-plugin/spring-plugins/spring-kafka-2.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/kafka/PollAndInvokeMethodInterceptor.java

@@ -18,40 +18,33 @@
 
 package org.apache.skywalking.apm.plugin.spring.kafka;
 
-import org.apache.skywalking.apm.agent.core.context.ContextManager;
 import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance;
 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.plugin.kafka.define.Constants;
-import org.apache.skywalking.apm.plugin.kafka.define.SpringKafkaContext;
+import org.apache.skywalking.apm.plugin.kafka.define.InterceptorMethod;
 
 import java.lang.reflect.Method;
 
 public class PollAndInvokeMethodInterceptor implements InstanceMethodsAroundInterceptor {
 
+    private static final String OPERATION_NAME = "/spring-kafka" + Constants.KAFKA_POLL_AND_INVOKE_OPERATION_NAME;
+
     @Override
     public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
                              MethodInterceptResult result) throws Throwable {
-        ContextManager.getRuntimeContext().put(Constants.SPRING_KAFKA_FLAG, new SpringKafkaContext());
+        InterceptorMethod.beforeMethod(OPERATION_NAME);
     }
 
     @Override
     public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
                               Object ret) throws Throwable {
-        SpringKafkaContext context = (SpringKafkaContext) ContextManager.getRuntimeContext().get(Constants.SPRING_KAFKA_FLAG);
-        if (context == null) {
-            return ret;
-        }
-        if (context.isNeedStop()) {
-            ContextManager.stopSpan();
-        } else {
-            ContextManager.getRuntimeContext().remove(Constants.SPRING_KAFKA_FLAG);
-        }
-        return ret;
+        return InterceptorMethod.afterMethod(ret);
     }
 
     @Override
     public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments,
                                       Class<?>[] argumentsTypes, Throwable t) {
+        InterceptorMethod.handleMethodException(t);
     }
 }

apm-sniffer/apm-toolkit-activation/apm-toolkit-kafka-activation/pom.xml

@@ -0,0 +1,44 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+  ~ 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.
+  ~
+  -->
+
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+<parent>
+    <artifactId>apm-toolkit-activation</artifactId>
+    <groupId>org.apache.skywalking</groupId>
+    <version>8.2.0-SNAPSHOT</version>
+</parent>
+<modelVersion>4.0.0</modelVersion>
+
+<artifactId>apm-toolkit-kafka-activation</artifactId>
+
+<dependencies>
+    <dependency>
+        <groupId>org.apache.skywalking</groupId>
+        <artifactId>apm-toolkit-trace</artifactId>
+        <version>${project.version}</version>
+        <scope>provided</scope>
+    </dependency>
+    <dependency>
+        <groupId>org.apache.skywalking</groupId>
+        <artifactId>apm-kafka-commons</artifactId>
+        <version>${project.version}</version>
+        <scope>provided</scope>
+    </dependency>
+</dependencies>
+</project>

apm-sniffer/apm-toolkit-activation/apm-toolkit-kafka-activation/src/main/java/org/apache/skywalking/apm/toolkit/activation/kafka/KafkaOnMessageAnnotationMethodInterceptor.java

@@ -0,0 +1,50 @@
+/*
+ * 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.toolkit.activation.kafka;
+
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance;
+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.plugin.kafka.define.Constants;
+import org.apache.skywalking.apm.plugin.kafka.define.InterceptorMethod;
+
+import java.lang.reflect.Method;
+
+public class KafkaOnMessageAnnotationMethodInterceptor implements InstanceMethodsAroundInterceptor {
+
+    private static final String OPERATION_NAME = "/kafka-toolkit" + Constants.KAFKA_POLL_AND_INVOKE_OPERATION_NAME;
+
+    @Override
+    public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
+                             MethodInterceptResult result) throws Throwable {
+        InterceptorMethod.beforeMethod(OPERATION_NAME);
+    }
+
+    @Override
+    public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
+                              Object ret) throws Throwable {
+        return InterceptorMethod.afterMethod(ret);
+    }
+
+    @Override
+    public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments,
+                                      Class<?>[] argumentsTypes, Throwable t) {
+        InterceptorMethod.handleMethodException(t);
+    }
+}

apm-sniffer/apm-toolkit-activation/apm-toolkit-kafka-activation/src/main/java/org/apache/skywalking/apm/toolkit/activation/kafka/define/KafkaOnMessageAnnotationInstrumentation.java

@@ -0,0 +1,69 @@
+/*
+ * 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.toolkit.activation.kafka.define;
+
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.DeclaredInstanceMethodsInterceptPoint;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.InstanceMethodsInterceptPoint;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.ClassInstanceMethodsEnhancePluginDefine;
+import org.apache.skywalking.apm.agent.core.plugin.match.ClassMatch;
+import org.apache.skywalking.apm.agent.core.plugin.match.MethodAnnotationMatch;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.ConstructorInterceptPoint;
+
+import static net.bytebuddy.matcher.ElementMatchers.isAnnotatedWith;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+public class KafkaOnMessageAnnotationInstrumentation extends ClassInstanceMethodsEnhancePluginDefine {
+
+    public static final String TRACE_ANNOTATION_METHOD_INTERCEPTOR = "org.apache.skywalking.apm.toolkit.activation.kafka.KafkaOnMessageAnnotationMethodInterceptor";
+    public static final String TRACE_ANNOTATION = "org.apache.skywalking.apm.toolkit.kafka.KafkaPollAndInvoke";
+
+    @Override
+    public ConstructorInterceptPoint[] getConstructorsInterceptPoints() {
+        return new ConstructorInterceptPoint[0];
+    }
+
+    @Override
+    public InstanceMethodsInterceptPoint[] getInstanceMethodsInterceptPoints() {
+        return new InstanceMethodsInterceptPoint[] {
+            new DeclaredInstanceMethodsInterceptPoint() {
+                @Override
+                public ElementMatcher<MethodDescription> getMethodsMatcher() {
+                    return isAnnotatedWith(named(TRACE_ANNOTATION));
+                }
+
+                @Override
+                public String getMethodsInterceptor() {
+                    return TRACE_ANNOTATION_METHOD_INTERCEPTOR;
+                }
+
+                @Override
+                public boolean isOverrideArgs() {
+                    return false;
+                }
+            }
+        };
+    }
+
+    @Override
+    protected ClassMatch enhanceClass() {
+        return MethodAnnotationMatch.byMethodAnnotationMatch(TRACE_ANNOTATION);
+    }
+}

apm-sniffer/apm-toolkit-activation/apm-toolkit-kafka-activation/src/main/resources/skywalking-plugin.def

@@ -0,0 +1,17 @@
+# 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.
+
+toolkit-kafka=org.apache.skywalking.apm.toolkit.activation.kafka.define.KafkaOnMessageAnnotationInstrumentation

apm-sniffer/apm-toolkit-activation/pom.xml

@@ -26,6 +26,7 @@
     <modelVersion>4.0.0</modelVersion>
     <packaging>pom</packaging>
     <modules>
+        <module>apm-toolkit-kafka-activation</module>
         <module>apm-toolkit-log4j-1.x-activation</module>
         <module>apm-toolkit-log4j-2.x-activation</module>
         <module>apm-toolkit-logback-1.x-activation</module>

docs/en/FAQ/kafka-plugin.md

@@ -5,4 +5,4 @@ The trace doesn't continue in kafka consumer side.
 The kafka client is pulling message from server, the plugin also just traces the pull action. As that, you need to do the manual instrument before the pull action, and include the further data process.
 
 ### Resolve
-Use Application Toolkit libraries to do manual instrumentation. such as `@Trace` annotation or OpenTracing API, Or if you're using `spring-kafka` 2.2.x or above, you can track the Consumer side without any code change.
+Use Application Toolkit libraries to do manual instrumentation. such as `@KafkaPollAndInvoke` annotation at `apm-toolkit-kafka` or OpenTracing API, Or if you're using `spring-kafka` 2.2.x or above, you can track the Consumer side without any code change.

test/plugin/scenarios/kafka-scenario/config/expectedData.yaml

@@ -127,6 +127,21 @@ segmentItems:
       skipAnalysis: 'false'
   - segmentId: not null
     spans:
+      - operationName: /kafka-scenario/case/kafka-thread2-ping
+        operationId: 0
+        parentSpanId: 0
+        spanId: 1
+        spanLayer: Http
+        startTime: nq 0
+        endTime: nq 0
+        componentId: 12
+        isError: false
+        spanType: Exit
+        peer: localhost:8080
+        skipAnalysis: false
+        tags:
+        - {key: http.method, value: GET}
+        - {key: url, value: 'http://localhost:8080/kafka-scenario/case/kafka-thread2-ping'}
       - operationName: Kafka/test./Consumer/testGroup2
         operationId: 0
         parentSpanId: -1
@@ -146,3 +161,25 @@ segmentItems:
              parentSpanId: 2, parentTraceSegmentId: not null, parentServiceInstance: not
                                                                 null, parentService: 'kafka-scenario', traceId: not null}
         skipAnalysis: 'false'
+  - segmentId: not null
+    spans:
+      - operationName: /case/kafka-thread2-ping
+        operationId: 0
+        parentSpanId: -1
+        spanId: 0
+        spanLayer: Http
+        startTime: not null
+        endTime: not null
+        componentId: 14
+        isError: false
+        spanType: Entry
+        peer: ''
+        skipAnalysis: false
+        tags:
+          - {key: url, value: 'http://localhost:8080/kafka-scenario/case/kafka-thread2-ping'}
+          - {key: http.method, value: GET}
+        refs:
+          - {parentEndpoint: not null, networkAddress: 'localhost:8080',
+             refType: CrossProcess, parentSpanId: 1, parentTraceSegmentId: not null,
+             parentServiceInstance: not null, parentService: kafka-scenario,
+             traceId: not null}

test/plugin/scenarios/kafka-scenario/pom.xml

@@ -34,6 +34,7 @@
         <log4j.version>2.6.2</log4j.version>
         <spring.version>4.3.8.RELEASE</spring.version>
         <spring-boot-version>1.5.2.RELEASE</spring-boot-version>
+        <okhttp-version>3.0.0</okhttp-version>
     </properties>
 
     <name>skywalking-kafka-scenario</name>
@@ -80,6 +81,11 @@
             <artifactId>spring-boot-starter-web</artifactId>
             <version>${spring-boot-version}</version>
         </dependency>
+        <dependency>
+            <groupId>com.squareup.okhttp3</groupId>
+            <artifactId>okhttp</artifactId>
+            <version>${okhttp-version}</version>
+        </dependency>
     </dependencies>
 
     <build>

test/plugin/scenarios/kafka-scenario/src/main/java/org/apache/skywalking/apm/toolkit/kafka/KafkaPollAndInvoke.java

@@ -0,0 +1,30 @@
+/*
+ * 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.toolkit.kafka;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@Target(ElementType.METHOD)
+@Retention(RetentionPolicy.RUNTIME)
+public @interface KafkaPollAndInvoke {
+
+}

test/plugin/scenarios/kafka-scenario/src/main/java/test/org/apache/skywalking/apm/testcase/kafka/controller/CaseController.java

@@ -18,11 +18,16 @@
 
 package test.org.apache.skywalking.apm.testcase.kafka.controller;
 
+import java.io.IOException;
 import java.util.Arrays;
 import java.util.Properties;
 import java.util.function.Consumer;
 import java.util.regex.Pattern;
 import javax.annotation.PostConstruct;
+
+import okhttp3.OkHttpClient;
+import okhttp3.Request;
+import okhttp3.Response;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
@@ -34,6 +39,7 @@ import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.apache.skywalking.apm.toolkit.kafka.KafkaPollAndInvoke;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.annotation.PropertySource;
 import org.springframework.stereotype.Controller;
@@ -112,6 +118,12 @@ public class CaseController {
         throw new RuntimeException("kafka not ready");
     }
 
+    @RequestMapping("/kafka-thread2-ping")
+    @ResponseBody
+    public String kafkaThread2Ping() {
+        return SUCCESS;
+    }
+
     private static void wrapProducer(Consumer<Producer<String, String>> consFunc, String bootstrapServers) {
         Properties producerProperties = new Properties();
         producerProperties.put("bootstrap.servers", bootstrapServers);
@@ -223,29 +235,33 @@ public class CaseController {
             consumerProperties.put("value.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
             KafkaConsumer<String, String> consumer = new KafkaConsumer<>(consumerProperties);
             consumer.subscribe(topicPattern, new NoOpConsumerRebalanceListener());
-            int i = 0;
-            while (i++ <= 10) {
-                try {
-                    Thread.sleep(1 * 1000);
-                } catch (InterruptedException e) {
-                }
+            while (true) {
+                if (pollAndInvoke(consumer)) break;
+            }
+            consumer.close();
+        }
 
-                ConsumerRecords<String, String> records = consumer.poll(100);
-
-                if (!records.isEmpty()) {
-                    for (ConsumerRecord<String, String> record : records) {
-                        logger.info("header: {}", new String(record.headers()
-                          .headers("TEST")
-                          .iterator()
-                          .next()
-                          .value()));
-                        logger.info("offset = {}, key = {}, value = {}", record.offset(), record.key(), record.value());
-                    }
-                    break;
-                }
+        @KafkaPollAndInvoke
+        private boolean pollAndInvoke(KafkaConsumer<String, String> consumer) {
+            try {
+                Thread.sleep(1 * 1000);
+            } catch (InterruptedException e) {
             }
 
-            consumer.close();
+            ConsumerRecords<String, String> records = consumer.poll(100);
+
+            if (!records.isEmpty()) {
+                OkHttpClient client = new OkHttpClient.Builder().build();
+                Request request = new Request.Builder().url("http://localhost:8080/kafka-scenario/case/kafka-thread2-ping").build();
+                Response response = null;
+                try {
+                    response = client.newCall(request).execute();
+                } catch (IOException e) {
+                }
+                response.body().close();
+                return true;
+            }
+            return false;
         }
     }
 }