Fix possible memory leaks in mybatis plugin. (#6913)

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/trace/AbstractSpan.java

@@ -104,8 +104,6 @@ public interface AbstractSpan extends AsyncSpan {
 
     String getOperationName();
 
-    int getComponentId();
-
     /**
      * Reference other trace segment.
      *

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/trace/AbstractTracingSpan.java

@@ -232,11 +232,6 @@ public abstract class AbstractTracingSpan implements AbstractSpan {
         return operationName;
     }
 
-    @Override
-    public int getComponentId() {
-        return componentId;
-    }
-
     @Override
     public AbstractTracingSpan setLayer(SpanLayer layer) {
         this.layer = layer;

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/trace/NoopSpan.java

@@ -100,11 +100,6 @@ public class NoopSpan implements AbstractSpan {
         return "";
     }
 
-    @Override
-    public int getComponentId() {
-        return 0;
-    }
-
     @Override
     public void ref(TraceSegmentRef ref) {
     }

apm-sniffer/optional-plugins/mybatis-3.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/mybatis/Constants.java

@@ -0,0 +1,24 @@
+/*
+ * 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.mybatis;
+
+public class Constants {
+    public static final String MYBATIS_SHELL_METHOD_NAME = "mybatis_shell_method_name";
+}

apm-sniffer/optional-plugins/mybatis-3.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/mybatis/MyBatisInterceptor.java

@@ -28,45 +28,33 @@ import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInt
 import org.apache.skywalking.apm.agent.core.util.MethodUtil;
 import org.apache.skywalking.apm.network.trace.component.ComponentsDefine;
 
-public class SqlSessionOperationInterceptor implements InstanceMethodsAroundInterceptor {
-
-    public static final String MYBATIS_ENTRY_METHOD_NAME = "mybatis_entry_method_name";
+public class MyBatisInterceptor implements InstanceMethodsAroundInterceptor {
 
     @Override
     public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
                              MethodInterceptResult result) throws Throwable {
-        String operationName = MethodUtil.generateOperationName(method);
-        AbstractSpan lastSpan = null;
-        if (ContextManager.isActive()) {
-            lastSpan = ContextManager.activeSpan();
-        }
-        if (lastSpan == null || lastSpan.getComponentId() != ComponentsDefine.MYBATIS.getId()) {
-            AbstractSpan span = ContextManager.createLocalSpan(operationName);
-            span.setComponent(ComponentsDefine.MYBATIS);
-            if (allArguments != null && allArguments.length != 0) {
-                Tags.MYBATIS_MAPPER.set(span, (String) allArguments[0]);
-            }
-            ContextManager.getRuntimeContext().put(MYBATIS_ENTRY_METHOD_NAME, operationName);
+        String operationName;
+        if (ContextManager.getRuntimeContext().get(Constants.MYBATIS_SHELL_METHOD_NAME) != null) {
+            operationName = String.valueOf(ContextManager.getRuntimeContext().get(Constants.MYBATIS_SHELL_METHOD_NAME));
+            ContextManager.getRuntimeContext().remove(Constants.MYBATIS_SHELL_METHOD_NAME);
+        } else {
+            operationName = MethodUtil.generateOperationName(method);
         }
+        AbstractSpan span = ContextManager.createLocalSpan(operationName);
+        span.setComponent(ComponentsDefine.MYBATIS);
+        Tags.MYBATIS_MAPPER.set(span, (String) allArguments[0]);
     }
 
     @Override
     public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
                               Object ret) throws Throwable {
-        String operationName = MethodUtil.generateOperationName(method);
-        if (String.valueOf(ContextManager.getRuntimeContext().get(MYBATIS_ENTRY_METHOD_NAME)).equals(operationName)) {
-            ContextManager.getRuntimeContext().remove(MYBATIS_ENTRY_METHOD_NAME);
-            ContextManager.stopSpan();
-        }
+        ContextManager.stopSpan();
         return ret;
     }
 
     @Override
     public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments,
                                       Class<?>[] argumentsTypes, Throwable t) {
-        String operationName = MethodUtil.generateOperationName(method);
-        if (String.valueOf(ContextManager.getRuntimeContext().get(MYBATIS_ENTRY_METHOD_NAME)).equals(operationName)) {
-            ContextManager.activeSpan().log(t);
-        }
+        ContextManager.activeSpan().log(t);
     }
 }

apm-sniffer/optional-plugins/mybatis-3.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/mybatis/MyBatisMethodMatch.java

@@ -22,16 +22,25 @@ import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
 import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 public enum MyBatisMethodMatch {
     INSTANCE;
 
     public ElementMatcher<MethodDescription> getMyBatisMethodMatcher() {
-        return named("selectOne").or(named("selectList"))
-                                 .or(named("selectMap"))
-                                 .or(named("select"))
+        return named("select").and(takesArguments(4))
+                              .or(named("selectList").and(takesArguments(3)))
+                              .or(named("update").and(takesArguments(2)));
+    }
+
+    public ElementMatcher<MethodDescription> getMyBatisShellMethodMatcher() {
+        return named("selectOne").or(named("selectMap"))
                                  .or(named("insert"))
-                                 .or(named("update"))
-                                 .or(named("delete"));
+                                 .or(named("delete"))
+                                 .or(named("select").and(takesArguments(2)))
+                                 .or(named("select").and(takesArguments(3)))
+                                 .or(named("selectList").and(takesArguments(1)))
+                                 .or(named("selectList").and(takesArguments(2)))
+                                 .or(named("update").and(takesArguments(1)));
     }
 }

apm-sniffer/optional-plugins/mybatis-3.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/mybatis/MyBatisShellMethodInterceptor.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.plugin.mybatis;
+
+import java.lang.reflect.Method;
+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.agent.core.util.MethodUtil;
+
+public class MyBatisShellMethodInterceptor implements InstanceMethodsAroundInterceptor {
+
+    @Override
+    public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
+                             MethodInterceptResult result) throws Throwable {
+        if (ContextManager.getRuntimeContext().get(Constants.MYBATIS_SHELL_METHOD_NAME) == null) {
+            String operationName = MethodUtil.generateOperationName(method);
+            ContextManager.getRuntimeContext().put(Constants.MYBATIS_SHELL_METHOD_NAME, operationName);
+        }
+    }
+
+    @Override
+    public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
+                              Object ret) throws Throwable {
+        return ret;
+    }
+
+    @Override
+    public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments,
+                                      Class<?>[] argumentsTypes, Throwable t) {
+
+    }
+}

apm-sniffer/optional-plugins/mybatis-3.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/mybatis/define/MyBatisInstrumentation.java

@@ -46,7 +46,7 @@ public class MyBatisInstrumentation extends ClassInstanceMethodsEnhancePluginDef
 
                 @Override
                 public String getMethodsInterceptor() {
-                    return "org.apache.skywalking.apm.plugin.mybatis.SqlSessionOperationInterceptor";
+                    return "org.apache.skywalking.apm.plugin.mybatis.MyBatisInterceptor";
                 }
 
                 @Override

apm-sniffer/optional-plugins/mybatis-3.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/mybatis/define/MyBatisShellMethodInstrumentation.java

@@ -28,7 +28,7 @@ import org.apache.skywalking.apm.plugin.mybatis.MyBatisMethodMatch;
 
 import static org.apache.skywalking.apm.agent.core.plugin.match.NameMatch.byName;
 
-public class MyBatisSpringInstrumentation extends ClassInstanceMethodsEnhancePluginDefine {
+public class MyBatisShellMethodInstrumentation extends ClassInstanceMethodsEnhancePluginDefine {
 
     @Override
     public ConstructorInterceptPoint[] getConstructorsInterceptPoints() {
@@ -41,12 +41,12 @@ public class MyBatisSpringInstrumentation extends ClassInstanceMethodsEnhancePlu
             new InstanceMethodsInterceptPoint() {
                 @Override
                 public ElementMatcher<MethodDescription> getMethodsMatcher() {
-                    return MyBatisMethodMatch.INSTANCE.getMyBatisMethodMatcher();
+                    return MyBatisMethodMatch.INSTANCE.getMyBatisShellMethodMatcher();
                 }
 
                 @Override
                 public String getMethodsInterceptor() {
-                    return "org.apache.skywalking.apm.plugin.mybatis.SqlSessionOperationInterceptor";
+                    return "org.apache.skywalking.apm.plugin.mybatis.MyBatisShellMethodInterceptor";
                 }
 
                 @Override
@@ -59,6 +59,6 @@ public class MyBatisSpringInstrumentation extends ClassInstanceMethodsEnhancePlu
 
     @Override
     public ClassMatch enhanceClass() {
-        return byName("org.mybatis.spring.SqlSessionTemplate");
+        return byName("org.apache.ibatis.session.defaults.DefaultSqlSession");
     }
 }

apm-sniffer/optional-plugins/mybatis-3.x-plugin/src/main/resources/skywalking-plugin.def

@@ -15,4 +15,4 @@
 # limitations under the License.
 
 mybatis-3.x=org.apache.skywalking.apm.plugin.mybatis.define.MyBatisInstrumentation
-mybatis-3.x=org.apache.skywalking.apm.plugin.mybatis.define.MyBatisSpringInstrumentation
+mybatis-3.x=org.apache.skywalking.apm.plugin.mybatis.define.MyBatisShellMethodInstrumentation

test/plugin/scenarios/mybatis-3.x-scenario/config/expectedData.yaml

@@ -35,23 +35,7 @@ segmentItems:
         - {key: db.type, value: sql}
         - {key: db.instance, value: ''}
         - {key: db.statement, value: 'insert into test.table_demo (name) values (?)'}
-      - operationName: Mysql/JDBI/Connection/close
-        operationId: 0
-        parentSpanId: 1
-        spanId: 3
-        spanLayer: Database
-        startTime: not null
-        endTime: not null
-        componentId: 33
-        isError: false
-        spanType: Exit
-        peer: mysql:3306
-        skipAnalysis: 'false'
-        tags:
-        - {key: db.type, value: sql}
-        - {key: db.instance, value: ''}
-        - {key: db.statement, value: ''}
-      - operationName: org.mybatis.spring.SqlSessionTemplate.insert(java.lang.String,java.lang.Object)
+      - operationName: org.apache.ibatis.session.defaults.DefaultSqlSession.insert(java.lang.String,java.lang.Object)
         operationId: 0
         parentSpanId: 0
         spanId: 1
@@ -65,6 +49,22 @@ segmentItems:
         skipAnalysis: 'false'
         tags:
         - {key: mybatis.mapper, value: test.apache.skywalking.apm.testcase.mybatis.mapper.DemoMapper.insert}
+      - operationName: Mysql/JDBI/Connection/close
+        operationId: 0
+        parentSpanId: 0
+        spanId: 3
+        spanLayer: Database
+        startTime: not null
+        endTime: not null
+        componentId: 33
+        isError: false
+        spanType: Exit
+        peer: mysql:3306
+        skipAnalysis: 'false'
+        tags:
+          - {key: db.type, value: sql}
+          - {key: db.instance, value: ''}
+          - {key: db.statement, value: ''}
       - operationName: Mysql/JDBI/PreparedStatement/execute
         operationId: 0
         parentSpanId: 4
@@ -81,23 +81,7 @@ segmentItems:
         - {key: db.type, value: sql}
         - {key: db.instance, value: ''}
         - {key: db.statement, value: 'insert into test.table_demo (name) values (?)'}
-      - operationName: Mysql/JDBI/Connection/close
-        operationId: 0
-        parentSpanId: 4
-        spanId: 6
-        spanLayer: Database
-        startTime: not null
-        endTime: not null
-        componentId: 33
-        isError: false
-        spanType: Exit
-        peer: mysql:3306
-        skipAnalysis: 'false'
-        tags:
-        - {key: db.type, value: sql}
-        - {key: db.instance, value: ''}
-        - {key: db.statement, value: ''}
-      - operationName: org.mybatis.spring.SqlSessionTemplate.insert(java.lang.String,java.lang.Object)
+      - operationName: org.apache.ibatis.session.defaults.DefaultSqlSession.insert(java.lang.String,java.lang.Object)
         operationId: 0
         parentSpanId: 0
         spanId: 4
@@ -111,6 +95,22 @@ segmentItems:
         skipAnalysis: 'false'
         tags:
         - {key: mybatis.mapper, value: test.apache.skywalking.apm.testcase.mybatis.mapper.DemoMapper.insert}
+      - operationName: Mysql/JDBI/Connection/close
+        operationId: 0
+        parentSpanId: 0
+        spanId: 6
+        spanLayer: Database
+        startTime: not null
+        endTime: not null
+        componentId: 33
+        isError: false
+        spanType: Exit
+        peer: mysql:3306
+        skipAnalysis: 'false'
+        tags:
+        - {key: db.type, value: sql}
+        - {key: db.instance, value: ''}
+        - {key: db.statement, value: ''}
       - operationName: /case/mybatis-case
         operationId: 0
         parentSpanId: -1