Skip to content

Commit 7831d54

Browse files
lucas-myxlucas-myx
authored
feat: record dubbo group and version for replay (#435)
* feat: record dubbo group and version * fix: unit test * fix: UT
1 parent 18c7b54 commit 7831d54

File tree

20 files changed

+241
-15
lines changed

20 files changed

+241
-15
lines changed

‎arex-instrumentation-api/src/main/java/io/arex/inst/runtime/model/ArexConstants.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,4 +53,5 @@ private ArexConstants() {}
5353
public static final String CALL_REPLAY_MAX = "callReplayMax";
5454
public static final String EXCEED_MAX_SIZE_TITLE = "exceed.max.size";
5555
public static final String EXCEED_MAX_SIZE_FLAG = "isExceedMaxSize";
56+
public static final String RECORD_SIZE_LIMIT = "arex.record.size.limit";
5657
}

‎arex-instrumentation-api/src/main/java/io/arex/inst/runtime/util/MockUtils.java

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -198,7 +198,7 @@ public static boolean checkResponseMocker(Mocker responseMocker) {
198198
if (MapUtils.getBoolean(targetResponse.getAttributes(), ArexConstants.EXCEED_MAX_SIZE_FLAG)) {
199199
exceedSizeLog = StringUtil.format(
200200
", because exceed memory max limit:%s, please check method return size, suggest replace it",
201-
AgentSizeOf.humanReadableUnits(ArexConstants.MEMORY_SIZE_1MB));
201+
AgentSizeOf.humanReadableUnits(getSizeLimit()));
202202
}
203203
LogManager.info(logTitle, StringUtil.format("operation: %s body of targetResponse is empty%s", operationName, exceedSizeLog));
204204
return false;
@@ -224,4 +224,8 @@ public static int methodRequestTypeHash(Mocker requestMocker) {
224224
requestMocker.getOperationName(),
225225
requestMocker.getTargetRequest().getType()));
226226
}
227+
228+
public static long getSizeLimit() {
229+
return Config.get().getLong(ArexConstants.RECORD_SIZE_LIMIT, ArexConstants.MEMORY_SIZE_1MB);
230+
}
227231
}

‎arex-instrumentation-foundation/src/main/java/io/arex/foundation/config/ConfigManager.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -366,6 +366,7 @@ private void updateRuntimeConfig() {
366366
configMap.put(AGENT_VERSION, agentVersion);
367367
configMap.put(AGENT_ENABLED, Boolean.toString(agentEnabled));
368368
configMap.put(STORAGE_SERVICE_MODE, storageServiceMode);
369+
configMap.put(DISABLE_MODULE, System.getProperty(DISABLE_MODULE));
369370
Map<String, String> extendFieldMap = getExtendField();
370371
if (MapUtils.isNotEmpty(extendFieldMap)) {
371372
configMap.putAll(extendFieldMap);

‎arex-instrumentation/config/arex-apollo/src/main/java/io/arex/inst/config/apollo/ApolloConfigChecker.java

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,15 @@
11
package io.arex.inst.config.apollo;
22

3+
import io.arex.agent.bootstrap.constants.ConfigConstants;
4+
import io.arex.agent.bootstrap.util.StringUtil;
5+
import io.arex.inst.runtime.config.Config;
6+
37
public class ApolloConfigChecker {
48

59
private static boolean isLoadedApollo = false;
610

11+
private static final String APOLLO_MODULE = "apollo-config";
12+
713
static {
814
try {
915
Class.forName("com.ctrip.framework.apollo.ConfigService");
@@ -14,6 +20,6 @@ public class ApolloConfigChecker {
1420
}
1521

1622
public static boolean unloadApollo() {
17-
return !isLoadedApollo;
23+
return !isLoadedApollo || Config.get().getString(ConfigConstants.DISABLE_MODULE, StringUtil.EMPTY).contains(APOLLO_MODULE);
1824
}
1925
}

‎arex-instrumentation/config/arex-apollo/src/test/java/io/arex/inst/config/apollo/ApolloConfigCheckerTest.java

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,28 @@
11
package io.arex.inst.config.apollo;
22

3+
import io.arex.inst.runtime.config.Config;
4+
import org.junit.jupiter.api.AfterAll;
5+
import org.junit.jupiter.api.BeforeAll;
36
import org.junit.jupiter.api.Test;
7+
import org.mockito.Mockito;
48

59
import static org.junit.jupiter.api.Assertions.*;
10+
import static org.mockito.ArgumentMatchers.any;
611

712
class ApolloConfigCheckerTest {
813

14+
@BeforeAll
15+
static void setUp() {
16+
Mockito.mockStatic(Config.class);
17+
Mockito.when(Config.get()).thenReturn(Mockito.mock(Config.class));
18+
Mockito.when(Config.get().getString(any(), any())).thenReturn("mock");
19+
}
20+
21+
@AfterAll
22+
static void tearDown() {
23+
Mockito.clearAllCaches();
24+
}
25+
926
@Test
1027
void unloadApollo() {
1128
assertFalse(ApolloConfigChecker.unloadApollo());

‎arex-instrumentation/config/arex-apollo/src/test/java/io/arex/inst/config/apollo/ApolloDubboRequestHandlerTest.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package io.arex.inst.config.apollo;
22

33
import io.arex.agent.bootstrap.util.Assert;
4+
import io.arex.inst.runtime.config.Config;
45
import io.arex.inst.runtime.context.ContextManager;
56
import io.arex.inst.runtime.model.ArexConstants;
67
import org.junit.jupiter.api.AfterAll;
@@ -32,6 +33,9 @@ static void setUp() {
3233
target = new ApolloDubboRequestHandler();
3334
mockStaticHelper = Mockito.mockStatic(ApolloConfigHelper.class);
3435
Mockito.mockStatic(ContextManager.class);
36+
Mockito.mockStatic(Config.class);
37+
Mockito.when(Config.get()).thenReturn(Mockito.mock(Config.class));
38+
Mockito.when(Config.get().getString(any(), any())).thenReturn("mock");
3539
}
3640

3741
@AfterAll

‎arex-instrumentation/config/arex-apollo/src/test/java/io/arex/inst/config/apollo/ApolloServletV3RequestHandlerTest.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package io.arex.inst.config.apollo;
22

33
import io.arex.agent.bootstrap.util.Assert;
4+
import io.arex.inst.runtime.config.Config;
45
import io.arex.inst.runtime.context.ContextManager;
56
import io.arex.inst.runtime.model.ArexConstants;
67
import org.junit.jupiter.api.AfterAll;
@@ -32,6 +33,9 @@ static void setUp() {
3233
target = new ApolloServletV3RequestHandler();
3334
mockStaticHelper = Mockito.mockStatic(ApolloConfigHelper.class);
3435
Mockito.mockStatic(ContextManager.class);
36+
Mockito.mockStatic(Config.class);
37+
Mockito.when(Config.get()).thenReturn(Mockito.mock(Config.class));
38+
Mockito.when(Config.get().getString(any(), any())).thenReturn("mock");
3539
}
3640

3741
@AfterAll

‎arex-instrumentation/dubbo/arex-dubbo-alibaba/src/main/java/io/arex/inst/dubbo/alibaba/DubboAdapter.java

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import io.arex.agent.bootstrap.util.ArrayUtils;
1010
import io.arex.agent.bootstrap.util.StringUtil;
1111
import io.arex.inst.dubbo.common.AbstractAdapter;
12+
import io.arex.inst.dubbo.common.DubboConstants;
1213
import io.arex.inst.runtime.model.ArexConstants;
1314
import io.arex.inst.runtime.serializer.Serializer;
1415
import io.arex.inst.runtime.log.LogManager;
@@ -18,9 +19,12 @@
1819
import org.slf4j.Logger;
1920
import org.slf4j.LoggerFactory;
2021

22+
import java.util.Map;
2123
import java.util.concurrent.Future;
2224
import java.util.function.Function;
2325

26+
import static io.arex.inst.dubbo.common.DubboConstants.*;
27+
2428
public class DubboAdapter extends AbstractAdapter {
2529
private static final Logger LOGGER = LoggerFactory.getLogger(DubboAdapter.class);
2630
private final Invoker<?> invoker;
@@ -146,4 +150,24 @@ public String getProtocol() {
146150
public String getConfigVersion() {
147151
return invocation.getAttachment(ArexConstants.CONFIG_VERSION);
148152
}
153+
154+
@Override
155+
protected Map<String, String> getRequestHeaders() {
156+
Map<String, String> headerMap = RpcContext.getContext().getAttachments();
157+
headerMap.put(KEY_PROTOCOL, getProtocol());
158+
headerMap.put(KEY_GROUP, getValByKey(DubboConstants.KEY_GROUP));
159+
headerMap.put(KEY_VERSION, getValByKey(DubboConstants.KEY_VERSION));
160+
return headerMap;
161+
}
162+
163+
private String getValByKey(String key) {
164+
String value = invocation.getAttachment(key);
165+
if (StringUtil.isNotEmpty(value)) {
166+
return value;
167+
}
168+
if (invocation.getInvoker() != null && invocation.getInvoker().getUrl() != null) {
169+
return invocation.getInvoker().getUrl().getParameter(key);
170+
}
171+
return StringUtil.EMPTY;
172+
}
149173
}

‎arex-instrumentation/dubbo/arex-dubbo-alibaba/src/main/java/io/arex/inst/dubbo/alibaba/DubboProviderExtractor.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616
import java.util.HashMap;
1717
import java.util.Map;
1818

19+
import static io.arex.inst.dubbo.common.DubboConstants.KEY_HEADERS;
20+
1921
public class DubboProviderExtractor extends DubboExtractor {
2022
public static void onServiceEnter(Invoker<?> invoker, Invocation invocation) {
2123
CaseEventDispatcher.onEvent(CaseEvent.ofEnterEvent());
@@ -44,10 +46,8 @@ public static void onServiceExit(Invoker<?> invoker, Invocation invocation, Resu
4446
}
4547
private static Mocker makeMocker(DubboAdapter adapter) {
4648
Mocker mocker = MockUtils.createDubboProvider(adapter.getServiceOperation());
47-
Map<String, String> headerMap = RpcContext.getContext().getAttachments();
48-
headerMap.put("protocol", adapter.getProtocol());
4949
Map<String, Object> attributes = new HashMap<>();
50-
attributes.put("Headers", headerMap);
50+
attributes.put(KEY_HEADERS, adapter.getRequestHeaders());
5151
attributes.put(ArexConstants.CONFIG_VERSION, adapter.getConfigVersion());
5252
// alibaba dubbo < 2.6.3 not support responseAttributes
5353
return buildMocker(mocker, adapter, attributes, null);

‎arex-instrumentation/dubbo/arex-dubbo-alibaba/src/test/java/io/arex/inst/dubbo/alibaba/DubboAdapterTest.java

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import com.alibaba.dubbo.rpc.support.ProtocolUtils;
1010
import io.arex.agent.bootstrap.model.ArexMocker;
1111
import io.arex.agent.bootstrap.model.Mocker;
12+
import io.arex.inst.dubbo.common.DubboConstants;
1213
import io.arex.inst.runtime.context.ContextManager;
1314
import io.arex.inst.runtime.util.MockUtils;
1415
import org.junit.jupiter.api.AfterAll;
@@ -225,4 +226,34 @@ void getProtocol() {
225226
void getConfigVersion() {
226227
assertNull(adapter.getConfigVersion());
227228
}
229+
230+
@ParameterizedTest
231+
@MethodSource("getRequestHeadersCase")
232+
void getRequestHeaders(Runnable mocker) {
233+
mocker.run();
234+
assertNotNull(adapter.getRequestHeaders());
235+
}
236+
237+
static Stream<Arguments> getRequestHeadersCase() {
238+
Runnable emptyMocker = () -> {};
239+
Runnable invocationMocker = () -> {
240+
Mockito.when(invocation.getAttachment(DubboConstants.KEY_GROUP)).thenReturn("mock");
241+
Mockito.when(invocation.getAttachment(DubboConstants.KEY_VERSION)).thenReturn("mock");
242+
};
243+
Runnable invokerMocker = () -> {
244+
Mockito.when(invocation.getAttachment(DubboConstants.KEY_GROUP)).thenReturn("");
245+
Mockito.when(invocation.getAttachment(DubboConstants.KEY_VERSION)).thenReturn("");
246+
Invoker invoker = Mockito.mock(Invoker.class);
247+
Mockito.when(invocation.getInvoker()).thenReturn(invoker);
248+
URL url = Mockito.mock(URL.class);
249+
Mockito.when(invoker.getUrl()).thenReturn(url);
250+
Mockito.when(url.getParameter(any())).thenReturn("mock");
251+
};
252+
253+
return Stream.of(
254+
arguments(emptyMocker),
255+
arguments(invocationMocker),
256+
arguments(invokerMocker)
257+
);
258+
}
228259
}

‎arex-instrumentation/dubbo/arex-dubbo-apache-v2/src/main/java/io/arex/inst/dubbo/apache/v2/DubboAdapter.java

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import io.arex.agent.bootstrap.util.ArrayUtils;
66
import io.arex.agent.bootstrap.util.StringUtil;
77
import io.arex.inst.dubbo.common.AbstractAdapter;
8+
import io.arex.inst.dubbo.common.DubboConstants;
89
import io.arex.inst.runtime.log.LogManager;
910
import io.arex.inst.runtime.model.ArexConstants;
1011
import io.arex.inst.runtime.serializer.Serializer;
@@ -15,8 +16,11 @@
1516
import org.slf4j.Logger;
1617
import org.slf4j.LoggerFactory;
1718

19+
import java.util.Map;
1820
import java.util.function.Function;
1921

22+
import static io.arex.inst.dubbo.common.DubboConstants.*;
23+
2024
public class DubboAdapter extends AbstractAdapter {
2125
private static final Logger LOGGER = LoggerFactory.getLogger(DubboAdapter.class);
2226
private final Invoker<?> invoker;
@@ -123,4 +127,24 @@ public String getProtocol() {
123127
public String getConfigVersion() {
124128
return invocation.getAttachment(ArexConstants.CONFIG_VERSION);
125129
}
130+
131+
@Override
132+
protected Map<String, String> getRequestHeaders() {
133+
Map<String, String> headerMap = RpcContext.getContext().getAttachments();
134+
headerMap.put(KEY_PROTOCOL, getProtocol());
135+
headerMap.put(KEY_GROUP, getValByKey(DubboConstants.KEY_GROUP));
136+
headerMap.put(KEY_VERSION, getValByKey(DubboConstants.KEY_VERSION));
137+
return headerMap;
138+
}
139+
140+
private String getValByKey(String key) {
141+
String value = invocation.getAttachment(key);
142+
if (StringUtil.isNotEmpty(value)) {
143+
return value;
144+
}
145+
if (invocation.getInvoker() != null && invocation.getInvoker().getUrl() != null) {
146+
return invocation.getInvoker().getUrl().getParameter(key);
147+
}
148+
return StringUtil.EMPTY;
149+
}
126150
}

‎arex-instrumentation/dubbo/arex-dubbo-apache-v2/src/main/java/io/arex/inst/dubbo/apache/v2/DubboProviderExtractor.java

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import io.arex.agent.bootstrap.model.MockCategoryType;
44
import io.arex.agent.bootstrap.model.Mocker;
5+
import io.arex.inst.dubbo.common.DubboConstants;
56
import io.arex.inst.dubbo.common.DubboExtractor;
67
import io.arex.inst.runtime.context.ContextManager;
78
import io.arex.inst.runtime.listener.CaseEvent;
@@ -19,6 +20,8 @@
1920
import java.util.HashMap;
2021
import java.util.Map;
2122

23+
import static io.arex.inst.dubbo.common.DubboConstants.*;
24+
2225
public class DubboProviderExtractor extends DubboExtractor {
2326
public static void onServiceEnter(Invoker<?> invoker, Invocation invocation) {
2427
CaseEventDispatcher.onEvent(CaseEvent.ofEnterEvent());
@@ -48,13 +51,11 @@ public static void onServiceExit(Invoker<?> invoker, Invocation invocation, Resu
4851
}
4952
private static Mocker makeMocker(DubboAdapter adapter) {
5053
Mocker mocker = MockUtils.createDubboProvider(adapter.getServiceOperation());
51-
Map<String, String> headerMap = RpcContext.getContext().getAttachments();
52-
headerMap.put("protocol", adapter.getProtocol());
5354
Map<String, Object> requestAttributes = new HashMap<>();
54-
requestAttributes.put("Headers", headerMap);
55+
requestAttributes.put(KEY_HEADERS, adapter.getRequestHeaders());
5556
requestAttributes.put(ArexConstants.CONFIG_VERSION, adapter.getConfigVersion());
5657
Map<String, Object> responseAttributes = new HashMap<>();
57-
responseAttributes.put("Headers", RpcContext.getServerContext().getAttachments());
58+
responseAttributes.put(KEY_HEADERS, RpcContext.getServerContext().getAttachments());
5859
return buildMocker(mocker, adapter, requestAttributes, responseAttributes);
5960
}
6061
}

‎arex-instrumentation/dubbo/arex-dubbo-apache-v2/src/test/java/io/arex/inst/dubbo/apache/v2/DubboAdapterTest.java

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import io.arex.agent.bootstrap.model.ArexMocker;
44
import io.arex.agent.bootstrap.model.Mocker;
5+
import io.arex.inst.dubbo.common.DubboConstants;
56
import io.arex.inst.runtime.context.ContextManager;
67
import io.arex.inst.runtime.util.MockUtils;
78
import org.apache.dubbo.common.URL;
@@ -195,4 +196,34 @@ void getProtocol() {
195196
void getConfigVersion() {
196197
assertNull(adapter.getConfigVersion());
197198
}
199+
200+
@ParameterizedTest
201+
@MethodSource("getRequestHeadersCase")
202+
void getRequestHeaders(Runnable mocker) {
203+
mocker.run();
204+
assertNotNull(adapter.getRequestHeaders());
205+
}
206+
207+
static Stream<Arguments> getRequestHeadersCase() {
208+
Runnable emptyMocker = () -> {};
209+
Runnable invocationMocker = () -> {
210+
Mockito.when(invocation.getAttachment(DubboConstants.KEY_GROUP)).thenReturn("mock");
211+
Mockito.when(invocation.getAttachment(DubboConstants.KEY_VERSION)).thenReturn("mock");
212+
};
213+
Runnable invokerMocker = () -> {
214+
Mockito.when(invocation.getAttachment(DubboConstants.KEY_GROUP)).thenReturn("");
215+
Mockito.when(invocation.getAttachment(DubboConstants.KEY_VERSION)).thenReturn("");
216+
Invoker invoker = Mockito.mock(Invoker.class);
217+
Mockito.when(invocation.getInvoker()).thenReturn(invoker);
218+
URL url = Mockito.mock(URL.class);
219+
Mockito.when(invoker.getUrl()).thenReturn(url);
220+
Mockito.when(url.getParameter(any())).thenReturn("mock");
221+
};
222+
223+
return Stream.of(
224+
arguments(emptyMocker),
225+
arguments(invocationMocker),
226+
arguments(invokerMocker)
227+
);
228+
}
198229
}

0 commit comments

Comments
 (0)