fix: propagate tracing context to downstream calls (#19)

Author: aaron-steinfeld

hypertrace-core-graphql-context/build.gradle.kts

@@ -10,6 +10,7 @@ dependencies {
   api("com.graphql-java-kickstart:graphql-java-servlet")
 
   implementation(project(":hypertrace-core-graphql-spi"))
+  implementation("com.google.guava:guava")
 
   testImplementation("org.junit.jupiter:junit-jupiter")
   testImplementation("org.mockito:mockito-core")

hypertrace-core-graphql-context/src/main/java/org/hypertrace/core/graphql/context/DefaultGraphQlRequestContextBuilder.java

@@ -1,12 +1,16 @@
 package org.hypertrace.core.graphql.context;
 
+import com.google.common.collect.Streams;
 import com.google.inject.Injector;
 import graphql.kickstart.servlet.context.DefaultGraphQLServletContext;
 import graphql.kickstart.servlet.context.DefaultGraphQLServletContextBuilder;
 import graphql.kickstart.servlet.context.GraphQLServletContext;
 import graphql.schema.DataFetcher;
 import java.util.Arrays;
+import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 import javax.inject.Inject;
 import javax.security.auth.Subject;
@@ -20,6 +24,8 @@ class DefaultGraphQlRequestContextBuilder extends DefaultGraphQLServletContextBu
   private static final String DEFAULT_CONTEXT_ID = "DEFAULT_CONTEXT_ID";
   static final String AUTHORIZATION_HEADER_KEY = "Authorization";
   static final String TENANT_ID_HEADER_KEY = "x-tenant-id";
+  static final Set<String> TRACING_CONTEXT_HEADER_KEY_PREFIXES =
+      Set.of("X-B3-", "traceparent", "tracestate");
 
   private final Injector injector;
   private final GraphQlServiceConfig serviceConfig;
@@ -75,6 +81,19 @@ public Optional<String> getTenantId() {
           .or(DefaultGraphQlRequestContextBuilder.this.serviceConfig::getDefaultTenantId);
     }
 
+    @Override
+    public Map<String, String> getTracingContextHeaders() {
+      return Streams.stream(
+              this.servletContext.getHttpServletRequest().getHeaderNames().asIterator())
+          .filter(
+              header ->
+                  TRACING_CONTEXT_HEADER_KEY_PREFIXES.stream()
+                      .anyMatch(prefix -> header.toLowerCase().startsWith(prefix.toLowerCase())))
+          .collect(
+              Collectors.toUnmodifiableMap(
+                  String::toLowerCase, this.servletContext.getHttpServletRequest()::getHeader));
+    }
+
     @Nonnull
     @Override
     public ContextualCachingKey getCachingKey() {

hypertrace-core-graphql-context/src/main/java/org/hypertrace/core/graphql/context/GraphQlRequestContext.java

@@ -2,6 +2,7 @@
 
 import graphql.kickstart.execution.context.GraphQLContext;
 import graphql.schema.DataFetcher;
+import java.util.Map;
 import java.util.Optional;
 import javax.annotation.Nonnull;
 
@@ -17,6 +18,8 @@ public interface GraphQlRequestContext extends GraphQLContext {
 
   Optional<String> getTenantId();
 
+  Map<String, String> getTracingContextHeaders();
+
   @Nonnull
   ContextualCachingKey getCachingKey();
 }

hypertrace-core-graphql-context/src/test/java/org/hypertrace/core/graphql/context/DefaultGraphQlRequestContextBuilderTest.java

@@ -1,5 +1,6 @@
 package org.hypertrace.core.graphql.context;
 
+import static java.util.Collections.emptyMap;
 import static org.hypertrace.core.graphql.context.DefaultGraphQlRequestContextBuilder.AUTHORIZATION_HEADER_KEY;
 import static org.hypertrace.core.graphql.context.DefaultGraphQlRequestContextBuilder.TENANT_ID_HEADER_KEY;
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -15,6 +16,9 @@
 
 import com.google.inject.Injector;
 import graphql.schema.DataFetcher;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
@@ -103,4 +107,32 @@ void returnsCachingKeysEqualForSameTenant() {
     var thirdKey = this.contextBuilder.build(this.mockRequest, this.mockResponse).getCachingKey();
     assertNotEquals(firstKey, thirdKey);
   }
+
+  @Test
+  void returnsEmptyMapIfNoTracingHeadersPresent() {
+    when(this.mockRequest.getHeaderNames()).thenReturn(Collections.enumeration(List.of("foo")));
+    assertEquals(emptyMap(), this.requestContext.getTracingContextHeaders());
+  }
+
+  @Test
+  void returnsLowerCasedTracingHeadersIfAnyMatches() {
+    when(this.mockRequest.getHeaderNames())
+        .thenReturn(
+            Collections.enumeration(
+                List.of(
+                    "traceSTATE", "traceparent", "other", "X-B3-traceid", "x-b3-parent-trace-id")));
+    when(this.mockRequest.getHeader(any(String.class)))
+        .thenAnswer(invocation -> invocation.getArgument(0) + " value");
+    assertEquals(
+        Map.of(
+            "tracestate",
+            "traceSTATE value",
+            "traceparent",
+            "traceparent value",
+            "x-b3-traceid",
+            "X-B3-traceid value",
+            "x-b3-parent-trace-id",
+            "x-b3-parent-trace-id value"),
+        this.requestContext.getTracingContextHeaders());
+  }
 }

hypertrace-core-graphql-grpc-utils/src/main/java/org/hypertrace/core/graphql/utils/grpc/DefaultGraphQlGrpcContextBuilder.java

@@ -6,6 +6,8 @@
 import io.grpc.Context;
 import io.grpc.stub.StreamObserver;
 import io.reactivex.rxjava3.core.Observable;
+import java.util.Arrays;
+import java.util.Collection;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Optional;
@@ -37,10 +39,12 @@ private DefaultGraphQlGrpcContext(
         GraphQlRequestContext requestContext,
         PlatformRequestContextBuilder platformRequestContextBuilder) {
       Map<String, String> grpcHeaders =
-          this.flattenOptionalMap(
-              Map.of(
-                  AUTHORIZATION_HEADER, this.extractAuthorizationHeader(requestContext),
-                  TENANT_ID_HEADER_KEY, this.extractTenantId(requestContext)));
+          this.mergeMaps(
+              requestContext.getTracingContextHeaders(),
+              this.flattenOptionalMap(
+                  Map.of(
+                      AUTHORIZATION_HEADER, this.extractAuthorizationHeader(requestContext),
+                      TENANT_ID_HEADER_KEY, this.extractTenantId(requestContext))));
 
       RequestContext platformContext = platformRequestContextBuilder.build(grpcHeaders);
       this.grpcContext = this.buildGrpcContext(platformContext);
@@ -87,5 +91,13 @@ private Map<String, String> flattenOptionalMap(Map<String, Optional<String>> opt
           .flatMap(Optional::stream)
           .collect(Collectors.toUnmodifiableMap(Entry::getKey, Entry::getValue));
     }
+
+    @SafeVarargs
+    private Map<String, String> mergeMaps(Map<String, String>... maps) {
+      return Arrays.stream(maps)
+          .map(Map::entrySet)
+          .flatMap(Collection::stream)
+          .collect(Collectors.toUnmodifiableMap(Entry::getKey, Entry::getValue));
+    }
   }
 }

hypertrace-core-graphql-grpc-utils/src/test/java/org/hypertrace/core/graphql/utils/grpc/DefaultGraphQlGrpcContextBuilderTest.java

@@ -140,4 +140,14 @@ void addsTenantIdToContext() {
     this.builder.build(this.mockRequestContext);
     verify(this.mockPlatformRequestContextBuilder).build(Map.of(TENANT_ID_HEADER_KEY, "tenant id"));
   }
+
+  @Test
+  void addsTracingHeadersToContext() {
+    when(this.mockRequestContext.getTenantId()).thenReturn(Optional.of("tenant id"));
+    when(this.mockRequestContext.getTracingContextHeaders())
+        .thenReturn(Map.of("traceid", "traceid value"));
+    this.builder.build(this.mockRequestContext);
+    verify(this.mockPlatformRequestContextBuilder)
+        .build(Map.of(TENANT_ID_HEADER_KEY, "tenant id", "traceid", "traceid value"));
+  }
 }