fix: add response attachments on service enter (#518)

fix: add response attachments on service enter
arextest · Jun 28, 2024 · a61784b · a61784b
1 parent fd8501f
commit a61784b
Showing 1 changed file with 1 addition and 9 deletions.
diff --git a/...ex-dubbo-apache-v3/src/main/java/io/arex/inst/dubbo/apache/v3/DubboProviderExtractor.java b/...ex-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);
-        }
-    }
 }