fix custom header set (#1054)

sofastack · Jul 16, 2021 · 8763684 · 8763684
1 parent c933fac
commit 8763684
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 34 deletions.
diff --git a/core-impl/filter/src/main/java/com/alipay/sofa/rpc/filter/ConsumerCustomHeaderFilter.java b/core-impl/filter/src/main/java/com/alipay/sofa/rpc/filter/ConsumerCustomHeaderFilter.java
@@ -18,7 +18,7 @@
 
 import com.alipay.sofa.common.utils.Ordered;
 import com.alipay.sofa.rpc.common.utils.CommonUtils;
-import com.alipay.sofa.rpc.context.RpcInternalContext;
+import com.alipay.sofa.rpc.context.RpcInvokeContext;
 import com.alipay.sofa.rpc.core.exception.SofaRpcException;
 import com.alipay.sofa.rpc.core.request.SofaRequest;
 import com.alipay.sofa.rpc.core.response.SofaResponse;
@@ -43,7 +43,7 @@ public SofaResponse invoke(FilterInvoker invoker, SofaRequest request) throws So
     }
 
     private void setCustomHeader(SofaRequest sofaRequest) {
-        RpcInternalContext context = RpcInternalContext.getContext();
+        RpcInvokeContext context = RpcInvokeContext.getContext();
         Map customHeader = context.getCustomHeader();
         if (CommonUtils.isNotEmpty(customHeader)) {
             sofaRequest.addRequestProps(customHeader);

diff --git a/core-impl/filter/src/test/java/com/alipay/sofa/rpc/filter/CustomHeaderFilterTest.java b/core-impl/filter/src/test/java/com/alipay/sofa/rpc/filter/CustomHeaderFilterTest.java
@@ -17,7 +17,7 @@
 package com.alipay.sofa.rpc.filter;
 
 import com.alipay.sofa.rpc.config.AbstractInterfaceConfig;
-import com.alipay.sofa.rpc.context.RpcInternalContext;
+import com.alipay.sofa.rpc.context.RpcInvokeContext;
 import com.alipay.sofa.rpc.core.exception.SofaRpcException;
 import com.alipay.sofa.rpc.core.request.SofaRequest;
 import com.alipay.sofa.rpc.core.response.SofaResponse;
@@ -34,11 +34,11 @@ public class CustomHeaderFilterTest {
 
     @Test
     public void testCustomFilter() {
-        RpcInternalContext.getContext().addCustomHeader("a", "b");
+        RpcInvokeContext.getContext().addCustomHeader("a", "b");
         ConsumerCustomHeaderFilter filter = new ConsumerCustomHeaderFilter();
         SofaRequest request = new SofaRequest();
         filter.invoke(new EmptyInvoker(null), request);
-        Assert.assertTrue(RpcInternalContext.getContext().getCustomHeader().isEmpty());
+        Assert.assertTrue(RpcInvokeContext.getContext().getCustomHeader().isEmpty());
         Assert.assertEquals("b", request.getRequestProp("a"));
     }
 

diff --git a/core/api/src/main/java/com/alipay/sofa/rpc/context/RpcInternalContext.java b/core/api/src/main/java/com/alipay/sofa/rpc/context/RpcInternalContext.java
@@ -28,7 +28,6 @@
 import java.net.InetSocketAddress;
 import java.util.ArrayDeque;
 import java.util.Deque;
-import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
@@ -185,13 +184,6 @@ protected RpcInternalContext() {
      */
     private ProviderInfo        providerInfo;
 
-
-    /**
-     * 自定义 header ，用完一次即删
-     */
-    protected HashMap<String, String> customHeader = new HashMap<>();
-
-
     /**
      * Is provider side.
      *
@@ -507,25 +499,4 @@ static boolean isHiddenParamKey(String key) {
         char c = key.charAt(0);
         return c == RpcConstants.HIDE_KEY_PREFIX;
     }
-
-
-    public Map<String, String> getCustomHeader() {
-        return new HashMap<>(customHeader);
-    }
-
-    /**
-     * 设置请求头，与 RequestBaggage 相比
-     * 1. 不受 enable baggage 开关影响，始终生效
-     * 2. 仅对一次调用生效，调用完成之会被清空
-     *
-     * @param key header key
-     * @param value header value
-     */
-    public void addCustomHeader(String key, String value) {
-        customHeader.put(key, value);
-    }
-
-    public void clearCustomHeader() {
-        customHeader.clear();
-    }
 }
diff --git a/core/api/src/main/java/com/alipay/sofa/rpc/context/RpcInvokeContext.java b/core/api/src/main/java/com/alipay/sofa/rpc/context/RpcInvokeContext.java
@@ -38,6 +38,13 @@ public class RpcInvokeContext {
      */
     protected static final ThreadLocal<RpcInvokeContext> LOCAL = new ThreadLocal<RpcInvokeContext>();
 
+    /**
+     * 自定义 header ，用完一次即删
+     */
+    protected HashMap<String, String> customHeader = new HashMap<>();
+
+
+
     /**
      * 得到上下文，没有则初始化
      *
@@ -393,6 +400,27 @@ public RpcInvokeContext setFuture(ResponseFuture<?> future) {
         return this;
     }
 
+
+    public Map<String, String> getCustomHeader() {
+        return new HashMap<>(customHeader);
+    }
+
+    /**
+     * 设置请求头，与 RequestBaggage 相比
+     * 1. 不受 enable baggage 开关影响，始终生效
+     * 2. 仅对一次调用生效，调用完成之会被清空
+     *
+     * @param key header key
+     * @param value header value
+     */
+    public void addCustomHeader(String key, String value) {
+        customHeader.put(key, value);
+    }
+
+    public void clearCustomHeader() {
+        customHeader.clear();
+    }
+
     @Override
     public String toString() {
         final StringBuilder sb = new StringBuilder(128);
@@ -408,4 +436,6 @@ public String toString() {
         sb.append('}');
         return sb.toString();
     }
+
+
 }