diff --git a/arex-instrumentation/dubbo/arex-dubbo-apache-v3/src/main/java/io/arex/inst/dubbo/apache/v3/DubboProviderExtractor.java b/arex-instrumentation/dubbo/arex-dubbo-apache-v3/src/main/java/io/arex/inst/dubbo/apache/v3/DubboProviderExtractor.java index 27eb1b21f..cd8dcc47f 100644 --- a/arex-instrumentation/dubbo/arex-dubbo-apache-v3/src/main/java/io/arex/inst/dubbo/apache/v3/DubboProviderExtractor.java +++ b/arex-instrumentation/dubbo/arex-dubbo-apache-v3/src/main/java/io/arex/inst/dubbo/apache/v3/DubboProviderExtractor.java @@ -32,12 +32,12 @@ public static void onServiceEnter(Invoker invoker, Invocation invocation) { addAttachmentsToContext(adapter); RequestHandlerManager.handleAfterCreateContext(invocation.getAttachments(), MockCategoryType.DUBBO_PROVIDER.getName()); invocation.getAttributes().put(ArexConstants.ORIGINAL_REQUEST, Serializer.serialize(invocation.getArguments())); + setResponseHeader((k, v) -> RpcContext.getServerContext().setAttachment(k, v)); } public static void onServiceExit(Invoker invoker, Invocation invocation, Result result) { if (!ContextManager.needRecordOrReplay()) { return; } - setResponseHeader((k, v) -> setAttachment(invocation, k, v)); DubboAdapter adapter = DubboAdapter.of(invoker, invocation); RequestHandlerManager.postHandle(invocation.getAttachments(), RpcContext.getServerContext().getAttachments(), MockCategoryType.DUBBO_PROVIDER.getName()); @@ -54,12 +54,4 @@ private static Mocker makeMocker(DubboAdapter adapter) { responseAttributes.put(KEY_HEADERS, RpcContext.getServerContext().getObjectAttachments()); return buildMocker(mocker, adapter, requestAttributes, responseAttributes); } - - private static void setAttachment(Invocation invocation, String key, String value) { - RpcContext.getServerContext().setAttachment(key, value); - if (invocation instanceof RpcInvocation) { - RpcInvocation rpcInvocation = (RpcInvocation) invocation; - rpcInvocation.setAttachment(key, value); - } - } }