Diff for traceid in VmApiProxyEnvironment

VmApiProxyEnvironment.diff

diff --git a/appengine-managed-runtime/src/main/java/com/google/apphosting/vmruntime/VmApiProxyEnvironment.java b/appengine-managed-runtime/src/main/java/com/google/apphosting/vmruntime/VmApiProxyEnvironment.java
index 2a0f347..9c776b8 100644
--- a/appengine-managed-runtime/src/main/java/com/google/apphosting/vmruntime/VmApiProxyEnvironment.java
+++ b/appengine-managed-runtime/src/main/java/com/google/apphosting/vmruntime/VmApiProxyEnvironment.java
@@ -329,6 +329,7 @@ public class VmApiProxyEnvironment implements ApiProxy.Environment {
             server,
             ticket,
             longAppId,
+            null,
             partition,
             module,
             majorVersion,
@@ -394,7 +395,8 @@ public class VmApiProxyEnvironment implements ApiProxy.Environment {
     }
     final String authDomain = request.getHeader(AUTH_DOMAIN_HEADER);
     boolean trustedApp = request.getHeader(IS_TRUSTED_IP_HEADER) != null;
-
+    String traceId = null;
+    
     Map<String, Object> attributes = new HashMap<>();
     // Fill in the attributes from the AttributeMapping.
     for (AttributeMapping mapping : AttributeMapping.values()) {
@@ -405,6 +407,10 @@ public class VmApiProxyEnvironment implements ApiProxy.Environment {
       String headerValue = request.getHeader(mapping.headerKey);
       if (headerValue != null) {
         attributes.put(mapping.attributeKey, headerValue);
+        if (mapping == AttributeMapping.CLOUD_TRACE_CONTEXT && !headerValue.isEmpty()) {
+          int slash = headerValue.indexOf("/");
+          traceId = (slash > 0) ? headerValue.substring(0, slash) : headerValue;
+        }
       } else if (mapping.defaultValue != null) {
         attributes.put(mapping.attributeKey, mapping.defaultValue);
       } // else: The attribute is expected to be missing if the header is not set.
@@ -429,6 +435,7 @@ public class VmApiProxyEnvironment implements ApiProxy.Environment {
             server,
             ticket,
             longAppId,
+            traceId,
             partition,
             module,
             majorVersion,
@@ -456,6 +463,7 @@ public class VmApiProxyEnvironment implements ApiProxy.Environment {
   private volatile String globalTicket; // global ticket is always valid
   private final String partition;
   private final String appId;
+  private final String traceId;
   private final String module;
   private final String majorVersion;
   private final String minorVersion;
@@ -499,6 +507,7 @@ public class VmApiProxyEnvironment implements ApiProxy.Environment {
       String server,
       String ticket,
       String appId,
+      String traceId,
       String partition,
       String module,
       String majorVersion,
@@ -541,6 +550,7 @@ public class VmApiProxyEnvironment implements ApiProxy.Environment {
     this.server = server;
     this.partition = partition;
     this.appId = partition + "~" + appId;
+    this.traceId = traceId;
     this.module = module == null ? "default" : module;
     this.majorVersion = majorVersion == null ? "" : majorVersion;
     this.minorVersion = minorVersion == null ? "" : minorVersion;
@@ -770,25 +780,6 @@ public class VmApiProxyEnvironment implements ApiProxy.Environment {
   }
 
   public String getTraceId() {
-    Object value =
-        getAttributes()
-            .get(VmApiProxyEnvironment.AttributeMapping.CLOUD_TRACE_CONTEXT.attributeKey);
-    if (!(value instanceof String)) {
-      return null;
-    }
-    String fullTraceId = (String) value;
-
-    // Extract the trace id from the header.
-    // TODO(user, qike): Use the code from the Trace SDK when it's available in /third_party.
-    if (fullTraceId.isEmpty() || Character.digit(fullTraceId.charAt(0), 16) < 0) {
-      return null;
-    }
-    for (int index = 1; index < fullTraceId.length(); index++) {
-      char ch = fullTraceId.charAt(index);
-      if (Character.digit(ch, 16) < 0) {
-        return fullTraceId.substring(0, index);
-      }
-    }
-    return null;
+    return traceId;
   }
 }