Update Jedis 4.x plugin to support Sharding and Cluster models. (#677)

Author: pg-yang

CHANGES.md

@@ -19,6 +19,7 @@ Release Notes.
 * Fix a bug in Spring Cloud Gateway if HttpClientFinalizer#send does not invoke, the span created at NettyRoutingFilterInterceptor can not stop.
 * Fix not tracing in HttpClient v5 when HttpHost(arg[0]) is null but `RoutingSupport#determineHost` works.
 * Support across thread tracing for SOFA-RPC.
+* Update Jedis 4.x plugin to support Sharding and Cluster models.
 
 #### Documentation
 * Update docs to describe `expired-plugins`.

apm-sniffer/apm-sdk-plugin/jedis-plugins/jedis-4.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/jedis/v4/AbstractConnectionInterceptor.java

@@ -18,6 +18,7 @@
 package org.apache.skywalking.apm.plugin.jedis.v4;
 
 import org.apache.skywalking.apm.agent.core.context.ContextManager;
+import org.apache.skywalking.apm.agent.core.context.tag.StringTag;
 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;
@@ -39,6 +40,8 @@ public abstract class AbstractConnectionInterceptor implements InstanceMethodsAr
 
     private static final String CACHE_TYPE = "Redis";
 
+    private static final StringTag TAG_ARGS = new StringTag("actual_target");
+
     @Override
     public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes, MethodInterceptResult result) throws Throwable {
         Iterator<Rawable> iterator = getCommands(allArguments);
@@ -49,12 +52,15 @@ public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allAr
         // Use lowercase to make config compatible with jedis-2.x-3.x plugin
         // Refer to `plugin.jedis.operation_mapping_read`, `plugin.jedis.operation_mapping_write` config item in agent.config
         String cmd = protocolCommand == null ? UNKNOWN : protocolCommand.toLowerCase();
-        String peer = String.valueOf(objInst.getSkyWalkingDynamicField());
+        ConnectionInformation connectionData = (ConnectionInformation) objInst.getSkyWalkingDynamicField();
+        // Use cluster information to adapt Virtual Cache if exists, otherwise use real server host
+        String peer =  StringUtil.isBlank(connectionData.getClusterNodes()) ? connectionData.getActualTarget() : connectionData.getClusterNodes();
         AbstractSpan span = ContextManager.createExitSpan("Jedis/" + cmd, peer);
         span.setComponent(ComponentsDefine.JEDIS);
         readKeyIfNecessary(iterator).ifPresent(key -> Tags.CACHE_KEY.set(span, key));
         Tags.CACHE_CMD.set(span, cmd);
         Tags.CACHE_TYPE.set(span, CACHE_TYPE);
+        TAG_ARGS.set(span, connectionData.getActualTarget());
         parseOperation(cmd).ifPresent(op -> Tags.CACHE_OP.set(span, op));
         SpanLayer.asCache(span);
     }

apm-sniffer/apm-sdk-plugin/jedis-plugins/jedis-4.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/jedis/v4/ConnectionConstructorInterceptor.java

@@ -26,6 +26,8 @@ public class ConnectionConstructorInterceptor implements InstanceConstructorInte
     @Override
     public void onConstruct(EnhancedInstance objInst, Object[] allArguments) throws Throwable {
         HostAndPort hostAndPort = ((DefaultJedisSocketFactory) allArguments[0]).getHostAndPort();
-        objInst.setSkyWalkingDynamicField(hostAndPort.getHost() + ":" + hostAndPort.getPort());
+        ConnectionInformation connectionData = new ConnectionInformation();
+        connectionData.setActualTarget(hostAndPort.toString());
+        objInst.setSkyWalkingDynamicField(connectionData);
     }
 }

apm-sniffer/apm-sdk-plugin/jedis-plugins/jedis-4.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/jedis/v4/ConnectionInformation.java

@@ -0,0 +1,26 @@
+/*
+ *   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.jedis.v4;
+
+import lombok.Data;
+
+@Data
+public class ConnectionInformation {
+    private String clusterNodes;
+    private String actualTarget;
+}

apm-sniffer/apm-sdk-plugin/jedis-plugins/jedis-4.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/jedis/v4/ConnectionProviderConstructorInterceptor.java

@@ -0,0 +1,43 @@
+/*
+ *   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.jedis.v4;
+
+import java.util.Collection;
+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.util.StringUtil;
+import redis.clients.jedis.HostAndPort;
+
+public class ConnectionProviderConstructorInterceptor implements InstanceConstructorInterceptor {
+    @Override
+    public void onConstruct(final EnhancedInstance objInst, final Object[] allArguments) throws Throwable {
+        if (objInst.getSkyWalkingDynamicField() != null) {
+            return;
+        }
+        Object arg = allArguments[0];
+        if (arg instanceof Collection) {
+            Collection<?> collection = (Collection<?>) arg;
+            final String[] array = collection.stream().map(Object::toString).toArray(String[]::new);
+            objInst.setSkyWalkingDynamicField(StringUtil.join(',', array));
+        }
+        if (arg instanceof HostAndPort) {
+            objInst.setSkyWalkingDynamicField(arg.toString());
+        }
+    }
+
+}

apm-sniffer/apm-sdk-plugin/jedis-plugins/jedis-4.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/jedis/v4/ConnectionProviderGetConnectionInterceptor.java

@@ -0,0 +1,60 @@
+/*
+ *   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.jedis.v4;
+
+import java.lang.reflect.Method;
+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;
+
+public class ConnectionProviderGetConnectionInterceptor implements InstanceMethodsAroundInterceptor {
+    @Override
+    public void beforeMethod(final EnhancedInstance objInst,
+                             final Method method,
+                             final Object[] allArguments,
+                             final Class<?>[] argumentsTypes,
+                             final MethodInterceptResult result) throws Throwable {
+
+    }
+
+    @Override
+    public Object afterMethod(final EnhancedInstance objInst,
+                              final Method method,
+                              final Object[] allArguments,
+                              final Class<?>[] argumentsTypes,
+                              final Object ret) throws Throwable {
+        if (ret instanceof EnhancedInstance) {
+            EnhancedInstance connection = (EnhancedInstance) ret;
+            if (connection.getSkyWalkingDynamicField() != null
+                && connection.getSkyWalkingDynamicField() instanceof ConnectionInformation) {
+                ((ConnectionInformation) connection.getSkyWalkingDynamicField()).setClusterNodes(
+                    (String) objInst.getSkyWalkingDynamicField());
+            }
+        }
+        return ret;
+    }
+
+    @Override
+    public void handleMethodException(final EnhancedInstance objInst,
+                                      final Method method,
+                                      final Object[] allArguments,
+                                      final Class<?>[] argumentsTypes,
+                                      final Throwable t) {
+
+    }
+}

apm-sniffer/apm-sdk-plugin/jedis-plugins/jedis-4.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/jedis/v4/JedisMethodInterceptor.java

@@ -18,6 +18,7 @@
 package org.apache.skywalking.apm.plugin.jedis.v4;
 
 import org.apache.skywalking.apm.agent.core.context.ContextManager;
+import org.apache.skywalking.apm.agent.core.context.tag.StringTag;
 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;
@@ -27,18 +28,23 @@
 import org.apache.skywalking.apm.network.trace.component.ComponentsDefine;
 
 import java.lang.reflect.Method;
+import org.apache.skywalking.apm.util.StringUtil;
 
 public class JedisMethodInterceptor implements InstanceMethodsAroundInterceptor {
+    private static final StringTag TAG_ARGS = new StringTag("actual_target");
 
     @Override
     public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
                              MethodInterceptResult result) throws Throwable {
-        String peer = String.valueOf(objInst.getSkyWalkingDynamicField());
+        final ConnectionInformation connectionData = (ConnectionInformation) objInst.getSkyWalkingDynamicField();
+        // Use cluster information to adapt Virtual Cache if exists, otherwise use real server host
+        String peer =  StringUtil.isBlank(connectionData.getClusterNodes()) ? connectionData.getActualTarget() : connectionData.getClusterNodes();
         AbstractSpan span = ContextManager.createExitSpan("Jedis/" + method.getName(), peer);
         span.setComponent(ComponentsDefine.JEDIS);
         SpanLayer.asCache(span);
         Tags.CACHE_TYPE.set(span, "Redis");
         Tags.CACHE_CMD.set(span, "BATCH_EXECUTE");
+        TAG_ARGS.set(span, connectionData.getActualTarget());
     }
 
     @Override

apm-sniffer/apm-sdk-plugin/jedis-plugins/jedis-4.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/jedis/v4/define/ConnectionProviderInstrumentation.java

@@ -0,0 +1,81 @@
+/*
+ *   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.jedis.v4.define;
+
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import net.bytebuddy.matcher.ElementMatchers;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.ConstructorInterceptPoint;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.InstanceMethodsInterceptPoint;
+import org.apache.skywalking.apm.agent.core.plugin.match.ClassMatch;
+
+import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static org.apache.skywalking.apm.agent.core.plugin.match.HierarchyMatch.byHierarchyMatch;
+
+public class ConnectionProviderInstrumentation extends AbstractWitnessInstrumentation {
+
+    private static final String ENHANCE_INTERFACE = "redis.clients.jedis.providers.ConnectionProvider";
+    private static final String CONNECTION_PROVIDER_CONSTRUCTION_INTERCEPT_CLASS = "org.apache.skywalking.apm.plugin.jedis.v4.ConnectionProviderConstructorInterceptor";
+    private static final String CONNECTION_PROVIDER_GET_CONNECTION_INTERCEPT_CLASS = "org.apache.skywalking.apm.plugin.jedis.v4.ConnectionProviderGetConnectionInterceptor";
+
+    @Override
+    protected ClassMatch enhanceClass() {
+        return byHierarchyMatch(ENHANCE_INTERFACE);
+    }
+
+    @Override
+    public ConstructorInterceptPoint[] getConstructorsInterceptPoints() {
+        return new ConstructorInterceptPoint[]{
+            new ConstructorInterceptPoint() {
+                @Override
+                public ElementMatcher<MethodDescription> getConstructorMatcher() {
+                    return ElementMatchers.takesArgument(0, named("redis.clients.jedis.HostAndPort"))
+                                          .or(ElementMatchers.takesArgument(0, hasSuperType(named("java.util.Collection"))));
+                }
+
+                @Override
+                public String getConstructorInterceptor() {
+                    return CONNECTION_PROVIDER_CONSTRUCTION_INTERCEPT_CLASS;
+                }
+            }
+        };
+    }
+
+    @Override
+    public InstanceMethodsInterceptPoint[] getInstanceMethodsInterceptPoints() {
+        return new InstanceMethodsInterceptPoint[] {
+            new InstanceMethodsInterceptPoint() {
+                @Override
+                public ElementMatcher<MethodDescription> getMethodsMatcher() {
+                    return named("getConnection");
+                }
+
+                @Override
+                public String getMethodsInterceptor() {
+                    return CONNECTION_PROVIDER_GET_CONNECTION_INTERCEPT_CLASS;
+                }
+
+                @Override
+                public boolean isOverrideArgs() {
+                    return false;
+                }
+            }
+        };
+    }
+}

apm-sniffer/apm-sdk-plugin/jedis-plugins/jedis-4.x-plugin/src/main/resources/skywalking-plugin.def

@@ -16,4 +16,5 @@
 
 jedis-4.x=org.apache.skywalking.apm.plugin.jedis.v4.define.ConnectionInstrumentation
 jedis-4.x=org.apache.skywalking.apm.plugin.jedis.v4.define.PipelineInstrumentation
-jedis-4.x=org.apache.skywalking.apm.plugin.jedis.v4.define.TransactionConstructorInstrumentation
+jedis-4.x=org.apache.skywalking.apm.plugin.jedis.v4.define.TransactionConstructorInstrumentation
+jedis-4.x=org.apache.skywalking.apm.plugin.jedis.v4.define.ConnectionProviderInstrumentation