optimize lettuce plugin

CHANGES.md

@@ -14,7 +14,8 @@ Release Notes.
 * Fix kafka-reporter-plugin shade package conflict
 * Add all config items to `agent.conf` file for convenient containerization use cases.
 * Advanced Kafka Producer configuration enhancement.
-* Suport mTLS for gRPC channel.
+* Support mTLS for gRPC channel.
+* fix the bug that plugin record wrong time elapse for lettuce plugin
 
 #### Documentation
 

apm-sniffer/apm-sdk-plugin/lettuce-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/lettuce/v5/DefaultEndpointChannelActiveInterceptor.java

@@ -0,0 +1,56 @@
+/*
+ * 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.lettuce.v5;
+
+import io.netty.channel.Channel;
+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 java.lang.reflect.Method;
+import java.net.InetSocketAddress;
+import java.net.SocketAddress;
+
+public class DefaultEndpointChannelActiveInterceptor implements InstanceMethodsAroundInterceptor {
+    public static final String IP_AND_PORT_DELIMITER = ":";
+
+    @Override
+    public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes, MethodInterceptResult result) throws Throwable {
+        Channel channel = (Channel) allArguments[0];
+        if (channel == null) {
+            return;
+        }
+        SocketAddress socketAddress = channel.remoteAddress();
+        if (socketAddress instanceof InetSocketAddress) {
+            InetSocketAddress inetSocketAddress = (InetSocketAddress) socketAddress;
+            String peer = inetSocketAddress.getHostString() + IP_AND_PORT_DELIMITER + inetSocketAddress.getPort();
+            objInst.setSkyWalkingDynamicField(peer);
+        }
+    }
+
+    @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/apm-sdk-plugin/lettuce-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/lettuce/v5/RedisChannelWriterInterceptor.java

@@ -18,68 +18,84 @@
 
 package org.apache.skywalking.apm.plugin.lettuce.v5;
 
-import java.lang.reflect.Method;
-import java.util.Collection;
-
+import io.lettuce.core.protocol.CommandArgs;
+import io.lettuce.core.protocol.DecoratedCommand;
+import io.lettuce.core.protocol.RedisCommand;
 import org.apache.skywalking.apm.agent.core.conf.Constants;
 import org.apache.skywalking.apm.agent.core.context.ContextManager;
 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.agent.core.context.trace.SpanLayer;
 import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance;
-import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.InstanceConstructorInterceptor;
 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 org.apache.skywalking.apm.util.StringUtil;
 
-import io.lettuce.core.protocol.CommandArgs;
-import io.lettuce.core.protocol.RedisCommand;
+import java.lang.reflect.Method;
+import java.util.Collection;
+import java.util.List;
 
-public class RedisChannelWriterInterceptor implements InstanceMethodsAroundInterceptor, InstanceConstructorInterceptor {
+@SuppressWarnings("unchecked")
+public class RedisChannelWriterInterceptor implements InstanceMethodsAroundInterceptor {
 
     private static final String PASSWORD_MASK = "******";
     private static final String ABBR = "...";
     private static final String DELIMITER_SPACE = " ";
     private static final String AUTH = "AUTH";
 
     @Override
-    public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
-        MethodInterceptResult result) throws Throwable {
+    public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes, MethodInterceptResult result) {
         String peer = (String) objInst.getSkyWalkingDynamicField();
+        RedisCommand<?, ?, ?> spanCarrierCommand = getSpanCarrierCommand(allArguments[0]);
+        if (spanCarrierCommand == null) {
+            return;
+        }
+        EnhancedInstance enhancedCommand = (EnhancedInstance) spanCarrierCommand;
+
+        // command has been handle by another channel writer (cluster or sentinel case)
+        if (enhancedCommand.getSkyWalkingDynamicField() != null) {
+            //set peer in last channel writer (delegate)
+            if (peer != null) {
+                AbstractSpan span = (AbstractSpan) enhancedCommand.getSkyWalkingDynamicField();
+                span.setPeer(peer);
+            }
+            return;
+        }
 
         StringBuilder dbStatement = new StringBuilder();
         String operationName = "Lettuce/";
-
         if (allArguments[0] instanceof RedisCommand) {
-            RedisCommand redisCommand = (RedisCommand) allArguments[0];
+            RedisCommand<?, ?, ?> redisCommand = (RedisCommand<?, ?, ?>) allArguments[0];
             String command = redisCommand.getType().name();
             operationName = operationName + command;
             dbStatement.append(command);
             if (LettucePluginConfig.Plugin.Lettuce.TRACE_REDIS_PARAMETERS) {
                 dbStatement.append(DELIMITER_SPACE).append(getArgsStatement(redisCommand));
             }
         } else if (allArguments[0] instanceof Collection) {
-            @SuppressWarnings("unchecked") Collection<RedisCommand> redisCommands = (Collection<RedisCommand>) allArguments[0];
+            Collection<RedisCommand<?, ?, ?>> redisCommands = (Collection<RedisCommand<?, ?, ?>>) allArguments[0];
             operationName = operationName + "BATCH_WRITE";
-            for (RedisCommand redisCommand : redisCommands) {
+            for (RedisCommand<?, ?, ?> redisCommand : redisCommands) {
                 dbStatement.append(redisCommand.getType().name()).append(";");
             }
         }
-
         AbstractSpan span = ContextManager.createExitSpan(operationName, peer);
         span.setComponent(ComponentsDefine.LETTUCE);
         Tags.DB_TYPE.set(span, "Redis");
         Tags.DB_STATEMENT.set(span, dbStatement.toString());
         SpanLayer.asCache(span);
+        span.prepareForAsync();
+        ContextManager.stopSpan();
+        enhancedCommand.setSkyWalkingDynamicField(span);
     }
-    
-    private String getArgsStatement(RedisCommand redisCommand) {
+
+    private String getArgsStatement(RedisCommand<?, ?, ?> redisCommand) {
         String statement;
         if (AUTH.equalsIgnoreCase(redisCommand.getType().name())) {
             statement = PASSWORD_MASK;
         } else {
-            CommandArgs args = redisCommand.getArgs();
+            CommandArgs<?, ?> args = redisCommand.getArgs();
             statement = (args != null) ? args.toCommandString() : Constants.EMPTY_STRING;
         }
         if (StringUtil.isNotEmpty(statement) && statement.length() > LettucePluginConfig.Plugin.Lettuce.REDIS_PARAMETER_MAX_LENGTH) {
@@ -89,22 +105,43 @@ private String getArgsStatement(RedisCommand redisCommand) {
     }
 
     @Override
-    public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
-        Object ret) throws Throwable {
-        ContextManager.stopSpan();
+    public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes, Object ret) {
         return ret;
     }
 
     @Override
-    public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments,
-        Class<?>[] argumentsTypes, Throwable t) {
-        AbstractSpan span = ContextManager.activeSpan();
-        span.log(t);
+    public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes, Throwable t) {
+        RedisCommand<?, ?, ?> redisCommand = getSpanCarrierCommand(allArguments[0]);
+        if (redisCommand instanceof EnhancedInstance && ((EnhancedInstance) redisCommand).getSkyWalkingDynamicField() != null) {
+            EnhancedInstance enhancedRedisCommand = (EnhancedInstance) redisCommand;
+            AbstractSpan abstractSpan = (AbstractSpan) enhancedRedisCommand.getSkyWalkingDynamicField();
+            enhancedRedisCommand.setSkyWalkingDynamicField(null);
+            abstractSpan.log(t);
+            abstractSpan.asyncFinish();
+        }
     }
 
-    @Override
-    public void onConstruct(EnhancedInstance objInst, Object[] allArguments) {
-        EnhancedInstance optionsInst = (EnhancedInstance) allArguments[0];
-        objInst.setSkyWalkingDynamicField(optionsInst.getSkyWalkingDynamicField());
+    private static RedisCommand<?, ?, ?> getSpanCarrierCommand(Object o) {
+        RedisCommand<?, ?, ?> command = null;
+        if (o instanceof RedisCommand) {
+            command = (RedisCommand<?, ?, ?>) o;
+        } else if (o instanceof List) {
+            List<?> list = (List<?>) o;
+            command = list.isEmpty() ? null : (RedisCommand<?, ?, ?>) list.get(list.size() - 1);
+        } else if (o instanceof Collection) {
+            Collection<RedisCommand<?, ?, ?>> redisCommands = (Collection<RedisCommand<?, ?, ?>>) o;
+            RedisCommand<?, ?, ?> last = null;
+            for (RedisCommand<?, ?, ?> redisCommand : redisCommands) {
+                last = redisCommand;
+            }
+            command = last;
+        }
+        if (command instanceof DecoratedCommand) {
+            while (command instanceof DecoratedCommand) {
+                DecoratedCommand<?, ?, ?> wrapper = (DecoratedCommand<?, ?, ?>) command;
+                command = wrapper.getDelegate();
+            }
+        }
+        return command;
     }
 }