fix the triple header context transfer issue (#1467)

* fix triple header and usedId problem

* Increase the extensibility of the header conversion context.

* fix the triple header context transfer issue

---------

Co-authored-by: liujianjun.ljj <[email protected]>

Author: EvenLjj

remoting/remoting-triple/src/main/java/com/alipay/sofa/rpc/interceptor/ServerReqHeaderInterceptor.java

@@ -62,14 +62,8 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(final ServerCall<Re
         SofaResponse sofaResponse = new SofaResponse();
         final Throwable[] throwable = { null };
         SofaRequest sofaRequest = new SofaRequest();
-        TripleTracerAdapter.serverReceived(sofaRequest, serverServiceDefinition, call, requestHeaders);
-        SofaTraceContext sofaTraceContext = SofaTraceContextHolder.getSofaTraceContext();
-        SofaTracerSpan serverSpan = sofaTraceContext.getCurrentSpan();
 
-        Context ctxWithSpan = Context.current()
-            .withValue(TracingContextKey.getKey(), serverSpan)
-            .withValue(TracingContextKey.getSpanContextKey(), serverSpan.context())
-            .withValue(TracingContextKey.getKeySofaRequest(), sofaRequest);
+        Context ctxWithSpan = convertHeaderToContext(call, requestHeaders, sofaRequest, serverServiceDefinition);
 
         //这里和下面不在一个线程
         if (RpcRunningState.isDebugMode()) {
@@ -167,4 +161,19 @@ private StatusRuntimeException fromThrowable(Throwable t) {
         };
         return result;
     }
+
+    protected <ReqT, RespT> Context convertHeaderToContext(ServerCall<ReqT, RespT> call,
+                                                           Metadata requestHeaders, SofaRequest sofaRequest,
+                                                           ServerServiceDefinition serverServiceDefinition) {
+        TripleTracerAdapter.serverReceived(sofaRequest, serverServiceDefinition, call, requestHeaders);
+        String userId = TripleTracerAdapter.getUserId(requestHeaders);
+        SofaTraceContext sofaTraceContext = SofaTraceContextHolder.getSofaTraceContext();
+        SofaTracerSpan serverSpan = sofaTraceContext.getCurrentSpan();
+        return Context.current()
+            .withValue(TracingContextKey.getKey(), serverSpan)
+            .withValue(TracingContextKey.getSpanContextKey(), serverSpan.context())
+            .withValue(TracingContextKey.getKeySofaRequest(), sofaRequest)
+            .withValue(TracingContextKey.getKeyMetadata(), requestHeaders)
+            .withValue(TracingContextKey.getKeyUserId(), userId);
+    }
 }

remoting/remoting-triple/src/main/java/com/alipay/sofa/rpc/tracer/sofatracer/TracingContextKey.java

@@ -31,6 +31,8 @@ public class TracingContextKey {
     public static final String                    KEY_SOFA_REQUEST_NAME = "io.opentracing.sofa-request";
     private static final Context.Key<SofaRequest> keySofaRequest        = Context.key(KEY_SOFA_REQUEST_NAME);
     private static final Context.Key<Metadata>    keyMetadata           = Context.key("io.opentracing.metadata");
+    public static final String                    KEY_TRIPLE_USER_ID    = "io.opentracing.tri-user-id";
+    private static final Context.Key<String>      keyUserId             = Context.key(KEY_TRIPLE_USER_ID);
 
     /**
      * Retrieves the active span.
@@ -70,4 +72,8 @@ public static Context.Key<SofaRequest> getKeySofaRequest() {
     public static Context.Key<Metadata> getKeyMetadata() {
         return keyMetadata;
     }
+
+    public static Context.Key<String> getKeyUserId() {
+        return keyUserId;
+    }
 }

remoting/remoting-triple/src/main/java/com/alipay/sofa/rpc/tracer/sofatracer/TripleTracerAdapter.java

@@ -16,6 +16,8 @@
  */
 package com.alipay.sofa.rpc.tracer.sofatracer;
 
+import com.alibaba.fastjson.JSON;
+import com.alibaba.fastjson.TypeReference;
 import com.alipay.common.tracer.core.appender.self.SelfLog;
 import com.alipay.common.tracer.core.context.span.SofaTracerSpanContext;
 import com.alipay.common.tracer.core.context.trace.SofaTraceContext;
@@ -51,6 +53,7 @@
 import java.util.Set;
 
 import static com.alipay.sofa.rpc.common.RemotingConstants.HEAD_APP_NAME;
+import static com.alipay.sofa.rpc.server.triple.TripleHeadKeys.HEAD_KEY_UNIT_INFO;
 
 /**
  * 客户端：startRpc ——> filter --> beforeSend --> 存入tracer信息 --> clientReceived
@@ -114,7 +117,7 @@ public static void beforeSend(SofaRequest sofaRequest, ConsumerConfig consumerCo
         if (StringUtils.isNotBlank(route)) {
             Map<String, String> map = new HashMap<>();
             map.put(USERID_KEY, route);
-            header.put(TripleHeadKeys.HEAD_KEY_UNIT_INFO.name(), JSONUtils.toJSONString(map));
+            header.put(HEAD_KEY_UNIT_INFO.name(), JSONUtils.toJSONString(map));
         }
 
         if (StringUtils.isNotEmpty(consumerConfig.getUniqueId())) {
@@ -300,4 +303,22 @@ public static void serverSend(SofaRequest request, final Metadata requestHeaders
             EventBus.post(new ServerSendEvent(request, response, throwable));
         }
     }
+
+    public static String getUserId(Metadata requestHeaders) {
+        String unitInfo = requestHeaders.get(HEAD_KEY_UNIT_INFO);
+        if (unitInfo == null) {
+            return null;
+        }
+        try {
+            Map<String, String> unitInfoMap = JSON.parseObject(unitInfo,
+                new TypeReference<Map<String, String>>() {
+                });
+            if (unitInfoMap != null) {
+                return unitInfoMap.get(USERID_KEY);
+            }
+        } catch (Exception e) {
+            LOGGER.warn("Failed to parse tri-unit-info: " + unitInfo, e);
+        }
+        return null;
+    }
 }

remoting/remoting-triple/src/test/java/com/alipay/sofa/rpc/tracer/sofatracer/TripleTracerAdapterTest.java

@@ -53,4 +53,14 @@ public void testBeforeSend() {
         Assert.assertEquals("value2", metadata.get(TripleHeadKeys.getKey("triple.header.object.key2")));
     }
 
+    @Test
+    public void testGetUserId() {
+        Metadata metadata = new Metadata();
+        Assert.assertNull(TripleTracerAdapter.getUserId(metadata));
+        metadata.put(TripleHeadKeys.HEAD_KEY_UNIT_INFO, "test");
+        Assert.assertNull(TripleTracerAdapter.getUserId(metadata));
+        metadata.put(TripleHeadKeys.HEAD_KEY_UNIT_INFO, "{\"userid\":\"99\"}");
+        Assert.assertEquals("99", TripleTracerAdapter.getUserId(metadata));
+    }
+
 }