From 3480c2929d647eb7d4bddcd69705162dd01ec492 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=A0=E5=91=88=E7=86=B9?= Date: Tue, 19 Jul 2022 15:32:00 +0800 Subject: [PATCH] allow custom caller app name (#1215) --- .../bootstrap/DefaultClientProxyInvoker.java | 14 ++++ .../DefaultClientProxyInvokerTest.java | 82 +++++++++++++++++++ .../alipay/sofa/rpc/common/RpcConstants.java | 6 ++ 3 files changed, 102 insertions(+) create mode 100644 bootstrap/bootstrap-api/src/test/java/com/alipay/sofa/rpc/bootstrap/DefaultClientProxyInvokerTest.java diff --git a/bootstrap/bootstrap-api/src/main/java/com/alipay/sofa/rpc/bootstrap/DefaultClientProxyInvoker.java b/bootstrap/bootstrap-api/src/main/java/com/alipay/sofa/rpc/bootstrap/DefaultClientProxyInvoker.java index af76a6a42..30968c677 100644 --- a/bootstrap/bootstrap-api/src/main/java/com/alipay/sofa/rpc/bootstrap/DefaultClientProxyInvoker.java +++ b/bootstrap/bootstrap-api/src/main/java/com/alipay/sofa/rpc/bootstrap/DefaultClientProxyInvoker.java @@ -20,6 +20,7 @@ import com.alipay.sofa.rpc.client.Cluster; import com.alipay.sofa.rpc.codec.SerializerFactory; import com.alipay.sofa.rpc.common.RemotingConstants; +import com.alipay.sofa.rpc.common.utils.StringUtils; import com.alipay.sofa.rpc.config.ConfigUniqueNameGenerator; import com.alipay.sofa.rpc.context.BaggageResolver; import com.alipay.sofa.rpc.context.RpcInternalContext; @@ -32,6 +33,7 @@ import com.alipay.sofa.rpc.log.LogCodes; import com.alipay.sofa.rpc.message.ResponseFuture; +import static com.alipay.sofa.rpc.common.RpcConstants.CUSTOM_CALLER_APP; import static com.alipay.sofa.rpc.common.RpcConstants.HIDDEN_KEY_INVOKE_CONTEXT; import static com.alipay.sofa.rpc.common.RpcConstants.HIDDEN_KEY_PINPOINT; import static com.alipay.sofa.rpc.common.RpcConstants.INTERNAL_KEY_APP_NAME; @@ -131,6 +133,18 @@ protected void decorateRequest(SofaRequest request) { // 额外属性通过HEAD传递给服务端 request.addRequestProp(RemotingConstants.HEAD_APP_NAME, consumerConfig.getAppName()); request.addRequestProp(RemotingConstants.HEAD_PROTOCOL, consumerConfig.getProtocol()); + + customRequest(request); + + } + + protected void customRequest(SofaRequest request) { + RpcInvokeContext context = RpcInvokeContext.getContext(); + Object customCallerApp = context.get(CUSTOM_CALLER_APP); + if (customCallerApp instanceof String && StringUtils.isNotBlank((String) customCallerApp)) { + context.remove(CUSTOM_CALLER_APP); + request.addRequestProp(RemotingConstants.HEAD_APP_NAME, customCallerApp); + } } @Override diff --git a/bootstrap/bootstrap-api/src/test/java/com/alipay/sofa/rpc/bootstrap/DefaultClientProxyInvokerTest.java b/bootstrap/bootstrap-api/src/test/java/com/alipay/sofa/rpc/bootstrap/DefaultClientProxyInvokerTest.java new file mode 100644 index 000000000..c9c3b0de9 --- /dev/null +++ b/bootstrap/bootstrap-api/src/test/java/com/alipay/sofa/rpc/bootstrap/DefaultClientProxyInvokerTest.java @@ -0,0 +1,82 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.alipay.sofa.rpc.bootstrap; + +import com.alipay.sofa.rpc.common.RemotingConstants; +import com.alipay.sofa.rpc.context.RpcInvokeContext; +import com.alipay.sofa.rpc.core.request.SofaRequest; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.Mock; +import org.mockito.junit.MockitoJUnitRunner; + +import static com.alipay.sofa.rpc.common.RpcConstants.CUSTOM_CALLER_APP; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doCallRealMethod; + +/** + * @author zhaowang + * @version : DefaultClientProxyInvokerTest.java, v 0.1 2022年06月13日 3:52 下午 zhaowang + */ +@RunWith(MockitoJUnitRunner.class) +public class DefaultClientProxyInvokerTest { + + private String originCallerApp = "originCallerApp"; + private String customCallerApp = "customCallerApp"; + @Mock + private DefaultClientProxyInvoker invoker; + + @Before + public void before() { + doCallRealMethod().when(invoker).customRequest(any()); + } + + @Test + public void testCustomCtx() { + SofaRequest request = getRequest(); + invoker.customRequest(request); + Assert.assertEquals(originCallerApp, request.getRequestProp(RemotingConstants.HEAD_APP_NAME)); + + request = getRequest(); + RpcInvokeContext.getContext().put(CUSTOM_CALLER_APP, customCallerApp); + invoker.customRequest(request); + Assert.assertEquals(customCallerApp, request.getRequestProp(RemotingConstants.HEAD_APP_NAME)); + + request = getRequest(); + RpcInvokeContext.getContext().put(CUSTOM_CALLER_APP, new Object()); + invoker.customRequest(request); + Assert.assertEquals(originCallerApp, request.getRequestProp(RemotingConstants.HEAD_APP_NAME)); + + request = getRequest(); + RpcInvokeContext.getContext().put(CUSTOM_CALLER_APP, ""); + invoker.customRequest(request); + Assert.assertEquals(originCallerApp, request.getRequestProp(RemotingConstants.HEAD_APP_NAME)); + + request = getRequest(); + RpcInvokeContext.getContext().put(CUSTOM_CALLER_APP, null); + invoker.customRequest(request); + Assert.assertEquals(originCallerApp, request.getRequestProp(RemotingConstants.HEAD_APP_NAME)); + } + + private SofaRequest getRequest() { + SofaRequest sofaRequest = new SofaRequest(); + sofaRequest.addRequestProp(RemotingConstants.HEAD_APP_NAME, originCallerApp); + return sofaRequest; + } +} \ No newline at end of file diff --git a/core/api/src/main/java/com/alipay/sofa/rpc/common/RpcConstants.java b/core/api/src/main/java/com/alipay/sofa/rpc/common/RpcConstants.java index 6517ec577..66998bae2 100644 --- a/core/api/src/main/java/com/alipay/sofa/rpc/common/RpcConstants.java +++ b/core/api/src/main/java/com/alipay/sofa/rpc/common/RpcConstants.java @@ -692,4 +692,10 @@ public class RpcConstants { */ public static final String INTERNAL_KEY_SERVER_SEND_TIME_MICRO = INTERNAL_KEY_PREFIX + "server_send_time_micro"; + + /** + * RpcInvokeContext 中自定义 caller app 的 Key + */ + public static final String CUSTOM_CALLER_APP = "custom_caller_app"; + }