change caller app name in RpcInternalContext (sofastack#1226)

taobaorun · Jul 28, 2022 · ffda5e2 · ffda5e2
1 parent 8af82b2
commit ffda5e2
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 8 deletions.
diff --git a/.../bootstrap-api/src/main/java/com/alipay/sofa/rpc/bootstrap/DefaultClientProxyInvoker.java b/.../bootstrap-api/src/main/java/com/alipay/sofa/rpc/bootstrap/DefaultClientProxyInvoker.java
@@ -134,15 +134,18 @@ protected void decorateRequest(SofaRequest request) {
         request.addRequestProp(RemotingConstants.HEAD_APP_NAME, consumerConfig.getAppName());
         request.addRequestProp(RemotingConstants.HEAD_PROTOCOL, consumerConfig.getProtocol());
 
-        customRequest(request);
+        customRequest(request, internalContext);
 
     }
 
-    protected void customRequest(SofaRequest request) {
+    protected void customRequest(SofaRequest request, RpcInternalContext internalContext) {
         RpcInvokeContext context = RpcInvokeContext.getContext();
         Object customCallerApp = context.get(CUSTOM_CALLER_APP);
         if (customCallerApp instanceof String && StringUtils.isNotBlank((String) customCallerApp)) {
             context.remove(CUSTOM_CALLER_APP);
+            if (RpcInternalContext.isAttachmentEnable()) {
+                internalContext.setAttachment(INTERNAL_KEY_APP_NAME, customCallerApp);
+            }
             request.addRequestProp(RemotingConstants.HEAD_APP_NAME, customCallerApp);
         }
     }

diff --git a/...tstrap-api/src/test/java/com/alipay/sofa/rpc/bootstrap/DefaultClientProxyInvokerTest.java b/...tstrap-api/src/test/java/com/alipay/sofa/rpc/bootstrap/DefaultClientProxyInvokerTest.java
@@ -17,6 +17,7 @@
 package com.alipay.sofa.rpc.bootstrap;
 
 import com.alipay.sofa.rpc.common.RemotingConstants;
+import com.alipay.sofa.rpc.context.RpcInternalContext;
 import com.alipay.sofa.rpc.context.RpcInvokeContext;
 import com.alipay.sofa.rpc.core.request.SofaRequest;
 import org.junit.Assert;
@@ -27,6 +28,7 @@
 import org.mockito.junit.MockitoJUnitRunner;
 
 import static com.alipay.sofa.rpc.common.RpcConstants.CUSTOM_CALLER_APP;
+import static com.alipay.sofa.rpc.common.RpcConstants.INTERNAL_KEY_APP_NAME;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doCallRealMethod;
 
@@ -44,39 +46,46 @@ public class DefaultClientProxyInvokerTest {
 
     @Before
     public void before() {
-        doCallRealMethod().when(invoker).customRequest(any());
+        doCallRealMethod().when(invoker).customRequest(any(), any());
     }
 
     @Test
     public void testCustomCtx() {
+        RpcInternalContext internalContext = RpcInternalContext.getContext();
         SofaRequest request = getRequest();
-        invoker.customRequest(request);
+        invoker.customRequest(request, internalContext);
         Assert.assertEquals(originCallerApp, request.getRequestProp(RemotingConstants.HEAD_APP_NAME));
+        Assert.assertEquals(originCallerApp, internalContext.getAttachment(INTERNAL_KEY_APP_NAME));
 
         request = getRequest();
         RpcInvokeContext.getContext().put(CUSTOM_CALLER_APP, customCallerApp);
-        invoker.customRequest(request);
+        invoker.customRequest(request, internalContext);
         Assert.assertEquals(customCallerApp, request.getRequestProp(RemotingConstants.HEAD_APP_NAME));
+        Assert.assertEquals(customCallerApp, internalContext.getAttachment(INTERNAL_KEY_APP_NAME));
 
         request = getRequest();
         RpcInvokeContext.getContext().put(CUSTOM_CALLER_APP, new Object());
-        invoker.customRequest(request);
+        invoker.customRequest(request, internalContext);
         Assert.assertEquals(originCallerApp, request.getRequestProp(RemotingConstants.HEAD_APP_NAME));
+        Assert.assertEquals(originCallerApp, internalContext.getAttachment(INTERNAL_KEY_APP_NAME));
 
         request = getRequest();
         RpcInvokeContext.getContext().put(CUSTOM_CALLER_APP, "");
-        invoker.customRequest(request);
+        invoker.customRequest(request, internalContext);
         Assert.assertEquals(originCallerApp, request.getRequestProp(RemotingConstants.HEAD_APP_NAME));
+        Assert.assertEquals(originCallerApp, internalContext.getAttachment(INTERNAL_KEY_APP_NAME));
 
         request = getRequest();
         RpcInvokeContext.getContext().put(CUSTOM_CALLER_APP, null);
-        invoker.customRequest(request);
+        invoker.customRequest(request, internalContext);
         Assert.assertEquals(originCallerApp, request.getRequestProp(RemotingConstants.HEAD_APP_NAME));
+        Assert.assertEquals(originCallerApp, internalContext.getAttachment(INTERNAL_KEY_APP_NAME));
     }
 
     private SofaRequest getRequest() {
         SofaRequest sofaRequest = new SofaRequest();
         sofaRequest.addRequestProp(RemotingConstants.HEAD_APP_NAME, originCallerApp);
+        RpcInternalContext.getContext().setAttachment(INTERNAL_KEY_APP_NAME, originCallerApp);
         return sofaRequest;
     }
 }