聊聊skywalking的jedis-pulgin

jedis-2.x=org.apache.skywalking.apm.plugin.jedis.v2.define.JedisClusterInstrumentation
jedis-2.x=org.apache.skywalking.apm.plugin.jedis.v2.define.JedisInstrumentation

public class JedisClusterInstrumentation extends ClassInstanceMethodsEnhancePluginDefine {
​
    private static final String ARGUMENT_TYPE_NAME = "redis.clients.jedis.HostAndPort";
    private static final String ENHANCE_CLASS = "redis.clients.jedis.JedisCluster";
    private static final String CONSTRUCTOR_WITH_LIST_HOSTANDPORT_ARG_INTERCEPT_CLASS = "org.apache.skywalking.apm.plugin.jedis.v2.JedisClusterConstructorWithListHostAndPortArgInterceptor";
    private static final String METHOD_INTERCEPT_CLASS = "org.apache.skywalking.apm.plugin.jedis.v2.JedisMethodInterceptor";
    private static final String CONSTRUCTOR_WITH_HOSTANDPORT_ARG_INTERCEPT_CLASS = "org.apache.skywalking.apm.plugin.jedis.v2.JedisClusterConstructorWithHostAndPortArgInterceptor";
​
    @Override
    public ClassMatch enhanceClass() {
        return byName(ENHANCE_CLASS);
    }
​
    @Override
    public ConstructorInterceptPoint[] getConstructorsInterceptPoints() {
        return new ConstructorInterceptPoint[] {
            new ConstructorInterceptPoint() {
                @Override
                public ElementMatcher<MethodDescription> getConstructorMatcher() {
                    return takesArgument(0, Set.class);
                }
​
                @Override
                public String getConstructorInterceptor() {
                    return CONSTRUCTOR_WITH_LIST_HOSTANDPORT_ARG_INTERCEPT_CLASS;
                }
            },
            new ConstructorInterceptPoint() {
                @Override
                public ElementMatcher<MethodDescription> getConstructorMatcher() {
                    return takesArgumentWithType(0, ARGUMENT_TYPE_NAME);
                }
​
                @Override
                public String getConstructorInterceptor() {
                    return CONSTRUCTOR_WITH_HOSTANDPORT_ARG_INTERCEPT_CLASS;
                }
            }
        };
    }
​
    @Override
    public InstanceMethodsInterceptPoint[] getInstanceMethodsInterceptPoints() {
        return new InstanceMethodsInterceptPoint[] {
            new InstanceMethodsInterceptPoint() {
                @Override
                public ElementMatcher<MethodDescription> getMethodsMatcher() {
                    return RedisMethodMatch.INSTANCE.getJedisClusterMethodMatcher();
                }
​
                @Override
                public String getMethodsInterceptor() {
                    return METHOD_INTERCEPT_CLASS;
                }
​
                @Override public boolean isOverrideArgs() {
                    return false;
                }
            }
        };
    }
}

public class JedisClusterConstructorWithListHostAndPortArgInterceptor implements InstanceConstructorInterceptor {
​
    @Override
    public void onConstruct(EnhancedInstance objInst, Object[] allArguments) {
        StringBuilder redisConnInfo = new StringBuilder();
        Set<HostAndPort> hostAndPorts = (Set<HostAndPort>)allArguments[0];
        for (HostAndPort hostAndPort : hostAndPorts) {
            redisConnInfo.append(hostAndPort.toString()).append(";");
        }
​
        objInst.setSkyWalkingDynamicField(PeerFormat.shorten(redisConnInfo.toString()));
    }
}

public class JedisClusterConstructorWithHostAndPortArgInterceptor implements InstanceConstructorInterceptor {
​
    @Override
    public void onConstruct(EnhancedInstance objInst, Object[] allArguments) {
        HostAndPort hostAndPort = (HostAndPort)allArguments[0];
        objInst.setSkyWalkingDynamicField(hostAndPort.getHost() + ":" + hostAndPort.getPort());
    }
}

public class JedisMethodInterceptor implements InstanceMethodsAroundInterceptor {
​
    @Override
    public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments,
        Class<?>[] argumentsTypes, MethodInterceptResult result) throws Throwable {
        String peer = String.valueOf(objInst.getSkyWalkingDynamicField());
        AbstractSpan span = ContextManager.createExitSpan("Jedis/" + method.getName(), peer);
        span.setComponent(ComponentsDefine.JEDIS);
        Tags.DB_TYPE.set(span, "Redis");
        SpanLayer.asCache(span);
​
        if (allArguments.length > 0 && allArguments[0] instanceof String) {
            Tags.DB_STATEMENT.set(span, method.getName() + " " + allArguments[0]);
        } else if (allArguments.length > 0 && allArguments[0] instanceof byte[]) {
            Tags.DB_STATEMENT.set(span, method.getName());
        }
    }
​
    @Override
    public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments,
        Class<?>[] argumentsTypes, Object ret) throws Throwable {
        ContextManager.stopSpan();
        return ret;
    }
​
    @Override
    public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments,
        Class<?>[] argumentsTypes, Throwable t) {
        AbstractSpan span = ContextManager.activeSpan();
        span.errorOccurred();
        span.log(t);
    }
}

public class JedisInstrumentation extends ClassInstanceMethodsEnhancePluginDefine {
​
    private static final String HOST_AND_PORT_ARG_TYPE_NAME = "redis.clients.jedis.HostAndPort";
    private static final String JEDIS_SHARD_INFO_ARG_TYPE_NAME = "redis.clients.jedis.JedisShardInfo";
    private static final String ENHANCE_CLASS = "redis.clients.jedis.Jedis";
    private static final String CONSTRUCTOR_WITH_STRING_ARG_INTERCEPT_CLASS = "org.apache.skywalking.apm.plugin.jedis.v2.JedisConstructorWithStringArgInterceptor";
    private static final String CONSTRUCTOR_WITH_SHARD_INFO_ARG_INTERCEPT_CLASS = "org.apache.skywalking.apm.plugin.jedis.v2.JedisConstructorWithShardInfoArgInterceptor";
    private static final String CONSTRUCTOR_WITH_HOST_AND_PORT_INTERCEPT_CLASS = "org.apache.skywalking.apm.plugin.jedis.v2.JedisClusterConstructorWithHostAndPortArgInterceptor";
    private static final String CONSTRUCTOR_WITH_URI_ARG_INTERCEPT_CLASS = "org.apache.skywalking.apm.plugin.jedis.v2.JedisConstructorWithUriArgInterceptor";
    private static final String JEDIS_METHOD_INTERCET_CLASS = "org.apache.skywalking.apm.plugin.jedis.v2.JedisMethodInterceptor";
​
    @Override
    public ClassMatch enhanceClass() {
        return byName(ENHANCE_CLASS);
    }
​
    @Override
    public ConstructorInterceptPoint[] getConstructorsInterceptPoints() {
        return new ConstructorInterceptPoint[] {
            new ConstructorInterceptPoint() {
                @Override
                public ElementMatcher<MethodDescription> getConstructorMatcher() {
                    return takesArgument(0, String.class);
                }
​
                @Override
                public String getConstructorInterceptor() {
                    return CONSTRUCTOR_WITH_STRING_ARG_INTERCEPT_CLASS;
                }
            },
            new ConstructorInterceptPoint() {
                @Override
                public ElementMatcher<MethodDescription> getConstructorMatcher() {
                    return takesArgumentWithType(0, HOST_AND_PORT_ARG_TYPE_NAME);
                }
​
                @Override
                public String getConstructorInterceptor() {
                    return CONSTRUCTOR_WITH_HOST_AND_PORT_INTERCEPT_CLASS;
                }
            },
            new ConstructorInterceptPoint() {
                @Override
                 public ElementMatcher<MethodDescription> getConstructorMatcher() {
                    return takesArgumentWithType(0, JEDIS_SHARD_INFO_ARG_TYPE_NAME);
                }
​
                @Override
                public String getConstructorInterceptor() {
                    return CONSTRUCTOR_WITH_SHARD_INFO_ARG_INTERCEPT_CLASS;
                }
            },
            new ConstructorInterceptPoint() {
                @Override
                public ElementMatcher<MethodDescription> getConstructorMatcher() {
                    return takesArgument(0, URI.class);
                }
​
                @Override
                public String getConstructorInterceptor() {
                    return CONSTRUCTOR_WITH_URI_ARG_INTERCEPT_CLASS;
                }
            }
        };
    }
​
    @Override
    public InstanceMethodsInterceptPoint[] getInstanceMethodsInterceptPoints() {
        return new InstanceMethodsInterceptPoint[] {
            new InstanceMethodsInterceptPoint() {
                @Override
                public ElementMatcher<MethodDescription> getMethodsMatcher() {
                    return RedisMethodMatch.INSTANCE.getJedisMethodMatcher();
                }
​
                @Override
                public String getMethodsInterceptor() {
                    return JEDIS_METHOD_INTERCET_CLASS;
                }
​
                @Override public boolean isOverrideArgs() {
                    return false;
                }
            }
        };
    }
}

public class JedisConstructorWithStringArgInterceptor implements InstanceConstructorInterceptor {
​
    @Override
    public void onConstruct(EnhancedInstance objInst, Object[] allArguments) {
        String host = (String)allArguments[0];
        String port = "6379";
        if (allArguments.length > 1) {
            port = String.valueOf(allArguments[1]);
        }
​
        objInst.setSkyWalkingDynamicField(host + ":" + port);
    }
}

public class JedisConstructorWithShardInfoArgInterceptor implements InstanceConstructorInterceptor {
​
    @Override
    public void onConstruct(EnhancedInstance objInst, Object[] allArguments) {
        String redisConnInfo;
        JedisShardInfo shardInfo = (JedisShardInfo)allArguments[0];
        redisConnInfo = shardInfo.getHost() + ":" + shardInfo.getPort();
        objInst.setSkyWalkingDynamicField(redisConnInfo);
    }
}

public class JedisConstructorWithUriArgInterceptor implements InstanceConstructorInterceptor {
​
    @Override
    public void onConstruct(EnhancedInstance objInst, Object[] allArguments) {
        URI uri = (URI)allArguments[0];
        objInst.setSkyWalkingDynamicField(uri.getHost() + ":" + uri.getPort());
    }
}

public class JedisMethodInterceptor implements InstanceMethodsAroundInterceptor {
​
    @Override
    public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments,
        Class<?>[] argumentsTypes, MethodInterceptResult result) throws Throwable {
        String peer = String.valueOf(objInst.getSkyWalkingDynamicField());
        AbstractSpan span = ContextManager.createExitSpan("Jedis/" + method.getName(), peer);
        span.setComponent(ComponentsDefine.JEDIS);
        Tags.DB_TYPE.set(span, "Redis");
        SpanLayer.asCache(span);
​
        if (allArguments.length > 0 && allArguments[0] instanceof String) {
            Tags.DB_STATEMENT.set(span, method.getName() + " " + allArguments[0]);
        } else if (allArguments.length > 0 && allArguments[0] instanceof byte[]) {
            Tags.DB_STATEMENT.set(span, method.getName());
        }
    }
​
    @Override
    public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments,
        Class<?>[] argumentsTypes, Object ret) throws Throwable {
        ContextManager.stopSpan();
        return ret;
    }
​
    @Override
    public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments,
        Class<?>[] argumentsTypes, Throwable t) {
        AbstractSpan span = ContextManager.activeSpan();
        span.errorOccurred();
        span.log(t);
    }
}

public enum RedisMethodMatch {
    INSTANCE;
​
    private ElementMatcher.Junction<MethodDescription> getIntersectionalMethodMacher() {
        return named("zcount").or(named("sunionstore")).or(named("zunionstore"))
            .or(named("del")).or(named("zinterstore")).or(named("echo"))
            .or(named("hscan")).or(named("psubscribe")).or(named("type"))
            .or(named("sinterstore")).or(named("setex")).or(named("zlexcount"))
            .or(named("brpoplpush")).or(named("bitcount")).or(named("llen"))
            .or(named("zscan")).or(named("lpushx")).or(named("bitpos"))
            .or(named("setnx")).or(named("hvals")).or(named("evalsha"))
            .or(named("substr")).or(named("geodist")).or(named("zrangeByLex"))
            .or(named("geoadd")).or(named("expire")).or(named("bitop"))
            .or(named("zrangeByScore")).or(named("smove")).or(named("lset"))
            .or(named("decrBy")).or(named("pttl")).or(named("scan"))
            .or(named("zrank")).or(named("blpop")).or(named("rpoplpush"))
            .or(named("zremrangeByLex")).or(named("get")).or(named("lpop"))
            .or(named("persist")).or(named("scriptExists")).or(named("georadius"))
            .or(named("set")).or(named("srandmember")).or(named("incr")).or(named("setbit"))
            .or(named("hexists")).or(named("expireAt")).or(named("pexpire")).or(named("zcard"))
            .or(named("bitfield")).or(named("zrevrangeByLex")).or(named("sinter")).or(named("srem"))
            .or(named("getrange")).or(named("rename")).or(named("zrevrank")).or(named("exists"))
            .or(named("setrange")).or(named("zremrangeByRank")).or(named("sadd")).or(named("sdiff"))
            .or(named("zrevrange")).or(named("getbit")).or(named("scard")).or(named("sdiffstore"))
            .or(named("zrevrangeByScore")).or(named("zincrby")).or(named("rpushx")).or(named("psetex"))
            .or(named("zrevrangeWithScores")).or(named("strlen")).or(named("hdel")).or(named("zremrangeByScore"))
            .or(named("geohash")).or(named("brpop")).or(named("lrem")).or(named("hlen")).or(named("decr"))
            .or(named("scriptLoad")).or(named("lpush")).or(named("lindex")).or(named("zrange")).or(named("incrBy"))
            .or(named("getSet")).or(named("ltrim")).or(named("incrByFloat")).or(named("rpop")).or(named("sort"))
            .or(named("zrevrangeByScoreWithScores")).or(named("pfadd")).or(named("eval")).or(named("linsert"))
            .or(named("pfcount")).or(named("hkeys")).or(named("hsetnx")).or(named("hincrBy")).or(named("hgetAll"))
            .or(named("hset")).or(named("spop")).or(named("zrangeWithScores")).or(named("hincrByFloat"))
            .or(named("hmset")).or(named("renamenx")).or(named("zrem")).or(named("msetnx")).or(named("hmget"))
            .or(named("sunion")).or(named("hget")).or(named("zadd")).or(named("move")).or(named("subscribe"))
            .or(named("geopos")).or(named("mset")).or(named("zrangeByScoreWithScores")).or(named("zscore"))
            .or(named("pexpireAt")).or(named("georadiusByMember")).or(named("ttl")).or(named("lrange"))
            .or(named("smembers")).or(named("pfmerge")).or(named("rpush")).or(named("publish"))
            .or(named("mget")).or(named("sscan")).or(named("append")).or(named("sismember"));
    }
​
    public ElementMatcher<MethodDescription> getJedisMethodMatcher() {
        return getIntersectionalMethodMacher().or(named("sentinelMasters")).or(named("clusterReplicate")).or(named("readonly"))
            .or(named("randomKey")).or(named("clusterInfo")).or(named("pubsubNumSub"))
            .or(named("sentinelSlaves")).or(named("clusterSetSlotImporting")).or(named("clusterSlaves"))
            .or(named("clusterFailover")).or(named("clusterSetSlotMigrating")).or(named("watch"))
            .or(named("clientKill")).or(named("clusterKeySlot")).or(named("clusterCountKeysInSlot"))
            .or(named("sentinelGetMasterAddrByName")).or(named("objectRefcount")).or(named("clusterMeet"))
            .or(named("sentinelSet")).or(named("clusterSetSlotNode")).or(named("clusterAddSlots"))
            .or(named("pubsubNumPat")).or(named("slowlogGet")).or(named("sentinelReset")).or(named("clusterNodes"))
            .or(named("sentinelMonitor")).or(named("configGet")).or(named("objectIdletime"))
            .or(named("pubsubChannels")).or(named("getParams")).or(named("sentinelRemove"))
            .or(named("migrate")).or(named("clusterForget")).or(named("asking")).or(named("keys"))
            .or(named("clientSetname")).or(named("clusterSaveConfig")).or(named("configSet"))
            .or(named("dump")).or(named("clusterFlushSlots")).or(named("clusterGetKeysInSlot"))
            .or(named("clusterReset")).or(named("restore")).or(named("clusterDelSlots"))
            .or(named("sentinelFailover")).or(named("clusterSetSlotStable")).or(named("objectEncoding"));
    }
​
    public ElementMatcher<MethodDescription> getJedisClusterMethodMatcher() {
        return getIntersectionalMethodMacher();
    }
}