Skip to content

Commit b1746ce

Browse files
DanLi39DanLi39
authored
feat: support methods provided by redisTemplate (#455)
1 parent 5b6fdb3 commit b1746ce

File tree

8 files changed

+367
-22
lines changed

8 files changed

+367
-22
lines changed

‎arex-instrumentation/redis/arex-spring-data-redis/src/main/java/io/arex/inst/spring/data/redis/MethodCollector.java

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package io.arex.inst.spring.data.redis;
22

33
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
4+
import static net.bytebuddy.matcher.ElementMatchers.named;
45
import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
56
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
67
import io.arex.inst.extension.MethodInstrumentation;
@@ -24,7 +25,9 @@ public class MethodCollector {
2425
"rangeByScoreWithScores", "rangeByLex", "reverseRangeByScore", "reverseRangeByScoreWithScores",
2526
"rank", "reverseRank", "removeRange", "removeRangeByScore", "score", "count", "zCard",
2627
// DefaultHashOperations
27-
"keys", "values", "entries", "randomKey", "randomEntry", "randomKeys", "randomEntries","delete"
28+
"keys", "values", "entries", "randomKey", "randomEntry", "randomKeys", "randomEntries", "delete",
29+
// RedisTemplate
30+
"hasKey", "unlink", "expire", "expireAt", "getExpire", "persist", "dump", "type"
2831
};
2932

3033
static final String[] TWO_OBJECT_KEY = new String[]{
@@ -33,14 +36,18 @@ public class MethodCollector {
3336
// DefaultSetOperations
3437
"difference", "differenceAndStore", "intersect", "intersectAndStore", "union", "unionAndStore",
3538
// DefaultHashOperations
36-
"hasKey", "increment", "lengthOfValue", "get", "put", "putIfAbsent"
39+
"hasKey", "increment", "lengthOfValue", "get", "put", "putIfAbsent",
40+
// RedisTemplate
41+
"renameIfAbsent"
3742
};
3843

3944
static final String[] COLLECTION_KEY = new String[]{
4045
// DefaultValueOperations
4146
"multiGet",
4247
// DefaultSetOperations
43-
"difference", "differenceAndStore", "intersect", "intersectAndStore", "union", "unionAndStore"
48+
"difference", "differenceAndStore", "intersect", "intersectAndStore", "union", "unionAndStore",
49+
// RedisTemplate
50+
"countExistingKeys", "delete", "unlink"
4451
};
4552

4653
static final String[] MAP_KEY = new String[]{
@@ -60,6 +67,9 @@ public class MethodCollector {
6067
"putAll"
6168
};
6269

70+
static final String EXECUTE = "execute";
71+
static final String SORT = "sort";
72+
6373
public static MethodInstrumentation arg1IsObjectKey(String adviceClassName) {
6474
return new MethodInstrumentation(isMethod().and(namedOneOf(ONE_OBJECT_KEY)).
6575
and(takesArgument(0, Object.class)), adviceClassName);
@@ -89,4 +99,12 @@ public static MethodInstrumentation arg1IsMapKey(String adviceClassName) {
8999
return new MethodInstrumentation(isMethod().and(namedOneOf(MAP_KEY)).
90100
and((takesArgument(0, Map.class))), adviceClassName);
91101
}
102+
103+
public static MethodInstrumentation execute(String adviceClassName) {
104+
return new MethodInstrumentation(isMethod().and(named(EXECUTE)).and(takesArgument(0, named("org.springframework.data.redis.core.script.RedisScript"))), adviceClassName);
105+
}
106+
107+
public static MethodInstrumentation sort(String adviceClassName) {
108+
return new MethodInstrumentation(isMethod().and(named(SORT)).and(takesArgument(0, named("org.springframework.data.redis.core.query.SortQuery"))), adviceClassName);
109+
}
92110
}

‎arex-instrumentation/redis/arex-spring-data-redis/src/main/java/io/arex/inst/spring/data/redis/SpringDataRedisInstrumentation.java renamed to ‎arex-instrumentation/redis/arex-spring-data-redis/src/main/java/io/arex/inst/spring/data/redis/OperationsInstrumentation.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,9 @@
1515
import org.springframework.data.redis.core.RedisTemplate;
1616

1717
/**
18-
* SpringDataRedisInstrumentation
18+
* OperationsInstrumentation
1919
*/
20-
public class SpringDataRedisInstrumentation extends TypeInstrumentation {
20+
public class OperationsInstrumentation extends TypeInstrumentation {
2121

2222
@Override
2323
protected ElementMatcher<TypeDescription> typeMatcher() {

‎arex-instrumentation/redis/arex-spring-data-redis/src/main/java/io/arex/inst/spring/data/redis/SpringDataRedisModuleInstrumentation.java renamed to ‎arex-instrumentation/redis/arex-spring-data-redis/src/main/java/io/arex/inst/spring/data/redis/RedisModuleInstrumentation.java

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,24 @@
11
package io.arex.inst.spring.data.redis;
22

3-
import static java.util.Collections.singletonList;
3+
import static java.util.Arrays.asList;
44
import com.google.auto.service.AutoService;
5-
import io.arex.inst.extension.ModuleDescription;
65
import io.arex.inst.extension.ModuleInstrumentation;
76
import io.arex.inst.extension.TypeInstrumentation;
87
import java.util.List;
98

109
/**
11-
* SpringDataRedisModuleInstrumentation
10+
* RedisModuleInstrumentation
1211
*/
1312
@AutoService(ModuleInstrumentation.class)
14-
public class SpringDataRedisModuleInstrumentation extends ModuleInstrumentation {
15-
public SpringDataRedisModuleInstrumentation() {
13+
public class RedisModuleInstrumentation extends ModuleInstrumentation {
14+
public RedisModuleInstrumentation() {
1615
super("spring-data-redis");
1716
}
1817

1918
@Override
2019
public List<TypeInstrumentation> instrumentationTypes() {
21-
return singletonList(new SpringDataRedisInstrumentation());
20+
return asList(
21+
new RedisTemplateInstrumentation(),
22+
new OperationsInstrumentation());
2223
}
2324
}

‎arex-instrumentation/redis/arex-spring-data-redis/src/main/java/io/arex/inst/spring/data/redis/RedisTemplateInstrumentation.java

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,173 @@
1+
package io.arex.inst.spring.data.redis;
2+
3+
import static java.util.Arrays.asList;
4+
import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
5+
import io.arex.agent.bootstrap.model.MockResult;
6+
import io.arex.inst.extension.MethodInstrumentation;
7+
import io.arex.inst.extension.TypeInstrumentation;
8+
import io.arex.inst.redis.common.RedisKeyUtil;
9+
import java.util.List;
10+
import java.util.Objects;
11+
import net.bytebuddy.asm.Advice;
12+
import net.bytebuddy.description.type.TypeDescription;
13+
import net.bytebuddy.implementation.bytecode.assign.Assigner;
14+
import net.bytebuddy.matcher.ElementMatcher;
15+
import org.springframework.data.redis.core.RedisTemplate;
16+
import org.springframework.data.redis.core.query.SortQuery;
17+
import org.springframework.data.redis.core.script.RedisScript;
18+
19+
/**
20+
* RedisTemplateInstrumentation
21+
*/
22+
public class RedisTemplateInstrumentation extends TypeInstrumentation {
23+
24+
@Override
25+
protected ElementMatcher<TypeDescription> typeMatcher() {
26+
return namedOneOf("org.springframework.data.redis.core.RedisTemplate");
27+
}
28+
29+
@Override
30+
public List<MethodInstrumentation> methodAdvices() {
31+
return asList(MethodCollector.arg1IsObjectKey(OneKeyAdvice.class.getName()),
32+
MethodCollector.arg1IsCollectionKey(OneKeyAdvice.class.getName()),
33+
MethodCollector.arg1AndArg2AreObjectKey(TwoKeysAdvice.class.getName()),
34+
MethodCollector.execute(ExecuteAdvice.class.getName()),
35+
MethodCollector.sort(SortAdvice.class.getName())
36+
);
37+
}
38+
39+
public static class OneKeyAdvice {
40+
41+
@Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class, suppress = Throwable.class)
42+
public static boolean onEnter(@Advice.This RedisTemplate template,
43+
@Advice.Origin("#m") String methodName,
44+
@Advice.Argument(0) Object key,
45+
@Advice.Local("mockResult") MockResult mockResult) {
46+
mockResult = RedisTemplateProvider.methodOnEnter(Objects.toString(template.getConnectionFactory()),
47+
methodName, RedisKeyUtil.generate(key));
48+
return mockResult != null && mockResult.notIgnoreMockResult();
49+
}
50+
51+
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
52+
public static void onExit(@Advice.This RedisTemplate template,
53+
@Advice.Origin("#m") String methodName,
54+
@Advice.Argument(0) Object key,
55+
@Advice.Return(readOnly = false, typing = Assigner.Typing.DYNAMIC) Object result,
56+
@Advice.Thrown(readOnly = false, typing = Assigner.Typing.DYNAMIC) Throwable throwable,
57+
@Advice.Local("mockResult") MockResult mockResult) {
58+
if (mockResult != null && mockResult.notIgnoreMockResult()) {
59+
if (mockResult.getThrowable() != null) {
60+
throwable = mockResult.getThrowable();
61+
} else {
62+
result = mockResult.getResult();
63+
}
64+
return;
65+
}
66+
RedisTemplateProvider.methodOnExit(Objects.toString(template.getConnectionFactory()), methodName,
67+
RedisKeyUtil.generate(key), result, throwable);
68+
}
69+
}
70+
71+
public static class TwoKeysAdvice {
72+
73+
@Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class, suppress = Throwable.class)
74+
public static boolean onEnter(@Advice.This RedisTemplate template,
75+
@Advice.Origin("#m") String methodName,
76+
@Advice.Argument(0) Object key,
77+
@Advice.Argument(1) Object otherKey,
78+
@Advice.Local("mockResult") MockResult mockResult) {
79+
mockResult = RedisTemplateProvider.methodOnEnter(Objects.toString(template.getConnectionFactory()),
80+
methodName, RedisKeyUtil.generate(key, otherKey));
81+
return mockResult != null && mockResult.notIgnoreMockResult();
82+
}
83+
84+
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
85+
public static void onExit(@Advice.This RedisTemplate template,
86+
@Advice.Origin("#m") String methodName,
87+
@Advice.Argument(0) Object key,
88+
@Advice.Argument(1) Object otherKey,
89+
@Advice.Return(readOnly = false, typing = Assigner.Typing.DYNAMIC) Object result,
90+
@Advice.Thrown(readOnly = false, typing = Assigner.Typing.DYNAMIC) Throwable throwable,
91+
@Advice.Local("mockResult") MockResult mockResult) {
92+
if (mockResult != null && mockResult.notIgnoreMockResult()) {
93+
if (mockResult.getThrowable() != null) {
94+
throwable = mockResult.getThrowable();
95+
} else {
96+
result = mockResult.getResult();
97+
}
98+
return;
99+
}
100+
RedisTemplateProvider.methodOnExit(Objects.toString(template.getConnectionFactory()), methodName,
101+
RedisKeyUtil.generate(key, otherKey), result, throwable);
102+
}
103+
}
104+
105+
106+
public static class ExecuteAdvice {
107+
108+
@Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class, suppress = Throwable.class)
109+
public static boolean onEnter(@Advice.This RedisTemplate template,
110+
@Advice.Origin("#m") String methodName,
111+
@Advice.AllArguments Object[] args,
112+
@Advice.Local("mockResult") MockResult mockResult) {
113+
List keys = args.length > 4 && args[3] instanceof List ? (List) args[3] : (List) args[1];
114+
RedisScript redisScript = (RedisScript) args[0];
115+
mockResult = RedisTemplateProvider.methodOnEnter(Objects.toString(template.getConnectionFactory()),
116+
methodName, RedisKeyUtil.generate(redisScript.getScriptAsString(), keys));
117+
return mockResult != null && mockResult.notIgnoreMockResult();
118+
}
119+
120+
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
121+
public static void onExit(@Advice.This RedisTemplate template,
122+
@Advice.Origin("#m") String methodName,
123+
@Advice.AllArguments Object[] args,
124+
@Advice.Return(readOnly = false, typing = Assigner.Typing.DYNAMIC) Object result,
125+
@Advice.Thrown(readOnly = false, typing = Assigner.Typing.DYNAMIC) Throwable throwable,
126+
@Advice.Local("mockResult") MockResult mockResult) {
127+
if (mockResult != null && mockResult.notIgnoreMockResult()) {
128+
if (mockResult.getThrowable() != null) {
129+
throwable = mockResult.getThrowable();
130+
} else {
131+
result = mockResult.getResult();
132+
}
133+
return;
134+
}
135+
RedisScript redisScript = (RedisScript) args[0];
136+
List keys = args.length > 4 && args[3] instanceof List ? (List) args[3] : (List) args[1];
137+
RedisTemplateProvider.methodOnExit(Objects.toString(template.getConnectionFactory()), methodName,
138+
RedisKeyUtil.generate(redisScript.getScriptAsString(), keys), result, throwable);
139+
}
140+
}
141+
142+
public static class SortAdvice {
143+
144+
@Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class, suppress = Throwable.class)
145+
public static boolean onEnter(@Advice.This RedisTemplate template,
146+
@Advice.Origin("#m") String methodName,
147+
@Advice.Argument(0) SortQuery sortQuery,
148+
@Advice.Local("mockResult") MockResult mockResult) {
149+
mockResult = RedisTemplateProvider.methodOnEnter(Objects.toString(template.getConnectionFactory()),
150+
methodName, RedisKeyUtil.generate(sortQuery));
151+
return mockResult != null && mockResult.notIgnoreMockResult();
152+
}
153+
154+
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
155+
public static void onExit(@Advice.This RedisTemplate template,
156+
@Advice.Origin("#m") String methodName,
157+
@Advice.Argument(0) SortQuery sortQuery,
158+
@Advice.Return(readOnly = false, typing = Assigner.Typing.DYNAMIC) Object result,
159+
@Advice.Thrown(readOnly = false, typing = Assigner.Typing.DYNAMIC) Throwable throwable,
160+
@Advice.Local("mockResult") MockResult mockResult) {
161+
if (mockResult != null && mockResult.notIgnoreMockResult()) {
162+
if (mockResult.getThrowable() != null) {
163+
throwable = mockResult.getThrowable();
164+
} else {
165+
result = mockResult.getResult();
166+
}
167+
return;
168+
}
169+
RedisTemplateProvider.methodOnExit(Objects.toString(template.getConnectionFactory()), methodName,
170+
RedisKeyUtil.generate(sortQuery), result, throwable);
171+
}
172+
}
173+
}

‎arex-instrumentation/redis/arex-spring-data-redis/src/test/java/io/arex/inst/spring/data/redis/MethodCollectorTest.java

Lines changed: 39 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,12 @@
22

33
import static io.arex.inst.spring.data.redis.MethodCollector.ONE_OBJECT_KEY;
44
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
5+
import static org.junit.jupiter.api.Assertions.assertEquals;
6+
7+
import java.lang.reflect.Method;
8+
import net.bytebuddy.description.method.MethodDescription.ForLoadedMethod;
59
import org.junit.jupiter.api.Test;
10+
import org.springframework.data.redis.core.RedisTemplate;
611

712
public class MethodCollectorTest {
813

@@ -14,16 +19,18 @@ public class MethodCollectorTest {
1419
"incrementScore", "reverseRange", "rangeWithScores", "reverseRangeWithScores", "rangeByScore",
1520
"rangeByScoreWithScores", "rangeByLex", "reverseRangeByScore", "reverseRangeByScoreWithScores", "rank",
1621
"reverseRank", "removeRange", "removeRangeByScore", "score", "count", "zCard", "keys", "values", "entries",
17-
"randomKey", "randomEntry", "randomKeys", "randomEntries", "delete"
22+
"randomKey", "randomEntry", "randomKeys", "randomEntries", "delete", "hasKey", "unlink", "expire", "expireAt",
23+
"getExpire", "persist", "dump", "type"
1824
};
1925

2026
private static final String[] TWO_OBJECT_KEY_TEST = new String[]{
2127
"rightPopAndLeftPush", "difference", "differenceAndStore", "intersect", "intersectAndStore", "union",
22-
"unionAndStore", "hasKey", "increment", "lengthOfValue", "get", "put", "putIfAbsent"
28+
"unionAndStore", "hasKey", "increment", "lengthOfValue", "get", "put", "putIfAbsent", "renameIfAbsent"
2329
};
2430

2531
private static final String[] COLLECTION_KEY_TEST = new String[]{
26-
"multiGet", "difference", "differenceAndStore", "intersect", "intersectAndStore", "union", "unionAndStore"
32+
"multiGet", "difference", "differenceAndStore", "intersect", "intersectAndStore", "union", "unionAndStore",
33+
"countExistingKeys", "delete", "unlink"
2734
};
2835

2936
private static final String[] MAP_KEY_TEST = new String[]{
@@ -47,4 +54,33 @@ public void keyMatcher() {
4754
assertArrayEquals(OBJECT_AND_COLLECTION_KEY_TEST, MethodCollector.OBJECT_AND_COLLECTION_KEY);
4855
assertArrayEquals(OBJECT_AND_MAP_KEY_TEST, MethodCollector.OBJECT_AND_MAP_KEY);
4956
}
57+
58+
@Test
59+
void execute() throws NoSuchMethodException {
60+
Method[] declaredMethods = RedisTemplate.class.getDeclaredMethods();
61+
int size = 0;
62+
for (Method method : declaredMethods) {
63+
boolean matches = MethodCollector.execute("").getMethodMatcher().matches(new ForLoadedMethod(method));
64+
if (matches) {
65+
size++;
66+
System.out.println(method);
67+
}
68+
}
69+
assertEquals(size, 2);
70+
71+
}
72+
73+
@Test
74+
void sort() throws NoSuchMethodException {
75+
Method[] declaredMethods = RedisTemplate.class.getDeclaredMethods();
76+
int size = 0;
77+
for (Method method : declaredMethods) {
78+
boolean matches = MethodCollector.sort("").getMethodMatcher().matches(new ForLoadedMethod(method));
79+
if (matches) {
80+
size++;
81+
System.out.println(method);
82+
}
83+
}
84+
assertEquals(size, 5);
85+
}
5086
}

‎arex-instrumentation/redis/arex-spring-data-redis/src/test/java/io/arex/inst/spring/data/redis/SpringDataRedisInstrumentationTest.java renamed to ‎arex-instrumentation/redis/arex-spring-data-redis/src/test/java/io/arex/inst/spring/data/redis/OperationsInstrumentationTest.java

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@
66
import static org.junit.jupiter.api.Assertions.assertNotNull;
77
import com.google.common.collect.Lists;
88
import io.arex.agent.bootstrap.model.MockResult;
9-
import io.arex.inst.spring.data.redis.SpringDataRedisInstrumentation.OneKeyAdvice;
10-
import io.arex.inst.spring.data.redis.SpringDataRedisInstrumentation.TwoKeysAdvice;
9+
import io.arex.inst.spring.data.redis.OperationsInstrumentation.OneKeyAdvice;
10+
import io.arex.inst.spring.data.redis.OperationsInstrumentation.TwoKeysAdvice;
1111
import java.util.List;
1212
import org.junit.jupiter.api.AfterAll;
1313
import org.junit.jupiter.api.BeforeAll;
@@ -18,9 +18,9 @@
1818
import org.springframework.data.redis.core.RedisTemplate;
1919

2020
@ExtendWith(MockitoExtension.class)
21-
class SpringDataRedisInstrumentationTest {
21+
class OperationsInstrumentationTest {
2222

23-
private static SpringDataRedisInstrumentation target;
23+
private static OperationsInstrumentation target;
2424
private static RedisTemplate redisTemplate = new RedisTemplate();
2525
private static final String KEY = "key";
2626
private static final String RESTULT = "result";
@@ -30,7 +30,7 @@ class SpringDataRedisInstrumentationTest {
3030
@BeforeAll
3131
static void setUp() {
3232
redisTemplate = Mockito.mock(RedisTemplate.class);
33-
target = new SpringDataRedisInstrumentation();
33+
target = new OperationsInstrumentation();
3434
}
3535

3636

‎arex-instrumentation/redis/arex-spring-data-redis/src/test/java/io/arex/inst/spring/data/redis/SpringDataRedisModuleInstrumentationTest.java renamed to ‎arex-instrumentation/redis/arex-spring-data-redis/src/test/java/io/arex/inst/spring/data/redis/RedisModuleInstrumentationTest.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,11 @@
33
import static org.junit.jupiter.api.Assertions.assertEquals;
44
import org.junit.jupiter.api.Test;
55

6-
class SpringDataRedisModuleInstrumentationTest {
6+
class RedisModuleInstrumentationTest {
77

88
@Test
99
void instrumentationTypes() {
10-
SpringDataRedisModuleInstrumentation instrumentation = new SpringDataRedisModuleInstrumentation();
11-
assertEquals(instrumentation.instrumentationTypes().size(),1);
10+
RedisModuleInstrumentation instrumentation = new RedisModuleInstrumentation();
11+
assertEquals(instrumentation.instrumentationTypes().size(), 2);
1212
}
1313
}

0 commit comments

Comments
 (0)