Timing header reader (#829)

* cleanup

* add server header timing reader

custom/build.gradle

@@ -11,6 +11,7 @@ dependencies {
   compileOnly("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure")
   compileOnly("io.opentelemetry.javaagent:opentelemetry-javaagent-extension-api")
   compileOnly("io.opentelemetry.javaagent:opentelemetry-javaagent-tooling")
+  compileOnly("io.opentelemetry.instrumentation:opentelemetry-instrumentation-api")
 
   testImplementation("io.opentelemetry.semconv:opentelemetry-semconv")
   testImplementation("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure-spi")

custom/src/main/java/com/grafana/extensions/filter/ServerAddressConfig.java

@@ -11,6 +11,7 @@
 import io.opentelemetry.sdk.metrics.View;
 import io.opentelemetry.sdk.metrics.ViewBuilder;
 import java.util.HashSet;
+import java.util.Set;
 
 public class ServerAddressConfig {
   public static final String SERVER_ADDRESS_OPT_IN =
@@ -22,7 +23,7 @@ static void configure(
       SdkMeterProviderBuilder sdkMeterProviderBuilder, ConfigProperties properties) {
     ViewBuilder builder = View.builder();
 
-    HashSet<String> keys = new HashSet<>();
+    Set<String> keys = new HashSet<>();
     keys.add("http.route");
     keys.add("http.request.method");
     keys.add("http.response.status_code");

custom/src/main/java/com/grafana/extensions/servertiming/ServerTimingHeaderCustomizer.java

@@ -6,10 +6,14 @@
 package com.grafana.extensions.servertiming;
 
 import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.api.trace.TraceFlags;
 import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseCustomizer;
 import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseMutator;
+import java.util.Set;
+import java.util.concurrent.ConcurrentSkipListSet;
 
 /**
  * Adds {@code Server-Timing} header (and {@code Access-Control-Expose-Headers}) to the HTTP
@@ -22,6 +26,8 @@ public class ServerTimingHeaderCustomizer implements HttpServerResponseCustomize
   // not using volatile because this field is set only once during agent initialization
   static boolean enabled = false;
 
+  public static Set<String> sampledTraces = new ConcurrentSkipListSet<>();
+
   @Override
   public <RESPONSE> void customize(
       Context context, RESPONSE response, HttpServerResponseMutator<RESPONSE> responseMutator) {
@@ -33,10 +39,21 @@ public <RESPONSE> void customize(
     responseMutator.appendHeader(response, EXPOSE_HEADERS, SERVER_TIMING);
   }
 
-  private static String toHeaderValue(Context context) {
+  static String toHeaderValue(Context context) {
+    SpanContext c = Span.fromContext(context).getSpanContext();
+    boolean sampled = sampledTraces.remove(c.getTraceId());
     TraceParentHolder traceParentHolder = new TraceParentHolder();
     W3CTraceContextPropagator.getInstance()
-        .inject(context, traceParentHolder, TraceParentHolder::set);
+        .inject(
+            context.with(
+                Span.wrap(
+                    SpanContext.create(
+                        c.getTraceId(),
+                        c.getSpanId(),
+                        sampled ? TraceFlags.getSampled() : TraceFlags.getDefault(),
+                        c.getTraceState()))),
+            traceParentHolder,
+            TraceParentHolder::set);
     return "traceparent;desc=\"" + traceParentHolder.traceParent + "\"";
   }
 

custom/src/main/java/com/grafana/extensions/servertiming/ServerTimingHeaderReader.java

@@ -0,0 +1,60 @@
+/*
+ * Copyright Grafana Labs
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package com.grafana.extensions.servertiming;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.TextMapGetter;
+import io.opentelemetry.instrumentation.api.semconv.http.HttpClientResponseConsumer;
+import io.opentelemetry.instrumentation.api.semconv.http.HttpCommonAttributesGetter;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+public class ServerTimingHeaderReader implements HttpClientResponseConsumer {
+  private static final TextMapGetter<Map<String, String>> GETTER =
+      new TextMapGetter<Map<String, String>>() {
+        @Override
+        public Iterable<String> keys(Map<String, String> carrier) {
+          return carrier.keySet();
+        }
+
+        @Override
+        public String get(Map<String, String> carrier, String key) {
+          return carrier.get(key);
+        }
+      };
+
+  @Override
+  public <REQUEST, RESPONSE> void consume(
+      HttpCommonAttributesGetter<REQUEST, RESPONSE> getter, REQUEST request, RESPONSE response) {
+    List<String> timings =
+        getter.getHttpResponseHeader(request, response, ServerTimingHeaderCustomizer.SERVER_TIMING);
+
+    for (String timing : timings) {
+      if (timing.startsWith("traceparent")) {
+        String[] parts = timing.split(";");
+        for (String part : parts) {
+          if (part.startsWith("desc=")) {
+            String traceParent = part.substring(6, part.length() - 1);
+            Context traceparent =
+                W3CTraceContextPropagator.getInstance()
+                    .extract(
+                        Context.current(),
+                        Collections.singletonMap("traceparent", traceParent),
+                        GETTER);
+            SpanContext spanContext = Span.fromContext(traceparent).getSpanContext();
+            if (spanContext.getTraceFlags().isSampled()) {
+              ServerTimingHeaderCustomizer.sampledTraces.add(spanContext.getTraceId());
+            }
+          }
+        }
+      }
+    }
+  }
+}

custom/src/main/java/io/opentelemetry/instrumentation/api/semconv/http/HttpClientResponseConsumer.java

@@ -0,0 +1,12 @@
+/*
+ * Copyright Grafana Labs
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.semconv.http;
+
+// copied from instrumentation-api - use when released
+public interface HttpClientResponseConsumer {
+  <REQUEST, RESPONSE> void consume(
+      HttpCommonAttributesGetter<REQUEST, RESPONSE> getter, REQUEST request, RESPONSE response);
+}

custom/src/main/resources/META-INF/services/io.opentelemetry.instrumentation.api.semconv.http.HttpClientResponseConsumer

@@ -0,0 +1 @@
+com.grafana.extensions.servertiming.ServerTimingHeaderReader

custom/src/test/java/com/grafana/extensions/servertiming/ServerTimingHeaderReaderTest.java

@@ -0,0 +1,58 @@
+/*
+ * Copyright Grafana Labs
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package com.grafana.extensions.servertiming;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
+import io.opentelemetry.instrumentation.testing.junit.LibraryInstrumentationExtension;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+class ServerTimingHeaderReaderTest {
+
+  private ServerTimingHeaderReader serverTimingHeaderReader = new ServerTimingHeaderReader();
+
+  @RegisterExtension InstrumentationExtension testing = LibraryInstrumentationExtension.create();
+
+  @BeforeEach
+  void setUp() {
+    ServerTimingHeaderCustomizer.sampledTraces.clear();
+  }
+
+  @Test
+  void notSampled() {
+    testing.runWithSpan(
+        "server",
+        () -> {
+          String serverTiming = ServerTimingHeaderCustomizer.toHeaderValue(Context.current());
+
+          serverTimingHeaderReader.consume(
+              new StringHttpCommonAttributesGetter(serverTiming), "request", "response");
+          assertThat(ServerTimingHeaderCustomizer.sampledTraces).isEmpty();
+        });
+  }
+
+  @Test
+  void sampled() {
+    testing.runWithSpan(
+        "server",
+        () -> {
+          String traceId = Span.current().getSpanContext().getTraceId();
+          ServerTimingHeaderCustomizer.sampledTraces.add(traceId);
+          String serverTiming = ServerTimingHeaderCustomizer.toHeaderValue(Context.current());
+
+          // remove the traceId to see that it is added back by the reader
+          ServerTimingHeaderCustomizer.sampledTraces.clear();
+          serverTimingHeaderReader.consume(
+              new StringHttpCommonAttributesGetter(serverTiming), "request", "response");
+          assertThat(ServerTimingHeaderCustomizer.sampledTraces).contains(traceId);
+        });
+  }
+}

custom/src/test/java/com/grafana/extensions/servertiming/ServerTimingHeaderTest.java

@@ -7,8 +7,7 @@
 
 import static com.grafana.extensions.servertiming.ServerTimingHeaderCustomizer.EXPOSE_HEADERS;
 import static com.grafana.extensions.servertiming.ServerTimingHeaderCustomizer.SERVER_TIMING;
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.assertj.core.api.Assertions.assertThat;
 
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.context.Context;
@@ -33,40 +32,38 @@ static void setUp() {
 
   @Test
   void shouldNotSetAnyHeadersWithoutValidCurrentSpan() {
-    // given
     var headers = new HashMap<String, String>();
 
-    // when
     serverTiming.customize(Context.root(), headers, Map::put);
 
-    // then
-    assertTrue(headers.isEmpty());
+    assertThat(headers).isEmpty();
   }
 
   @Test
   void shouldSetHeaders() {
-    // given
+
     var headers = new HashMap<String, String>();
 
-    // when
     var spanContext =
         testing.runWithSpan(
             "server",
             () -> {
+              ServerTimingHeaderCustomizer.sampledTraces.add(
+                  Span.current().getSpanContext().getTraceId());
               serverTiming.customize(Context.current(), headers, Map::put);
               return Span.current().getSpanContext();
             });
 
-    // then
-    assertEquals(2, headers.size());
+    assertThat(headers).hasSize(2);
 
     var serverTimingHeaderValue =
         "traceparent;desc=\"00-"
             + spanContext.getTraceId()
             + "-"
             + spanContext.getSpanId()
             + "-01\"";
-    assertEquals(serverTimingHeaderValue, headers.get(SERVER_TIMING));
-    assertEquals(SERVER_TIMING, headers.get(EXPOSE_HEADERS));
+    assertThat(headers)
+        .containsEntry(SERVER_TIMING, serverTimingHeaderValue)
+        .containsEntry(EXPOSE_HEADERS, SERVER_TIMING);
   }
 }

custom/src/test/java/com/grafana/extensions/servertiming/StringHttpCommonAttributesGetter.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright Grafana Labs
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package com.grafana.extensions.servertiming;
+
+import io.opentelemetry.instrumentation.api.semconv.http.HttpCommonAttributesGetter;
+import java.util.List;
+import javax.annotation.Nullable;
+
+class StringHttpCommonAttributesGetter implements HttpCommonAttributesGetter<String, String> {
+  private String serverTiming;
+
+  public StringHttpCommonAttributesGetter(String serverTiming) {
+    this.serverTiming = serverTiming;
+  }
+
+  @Nullable
+  @Override
+  public String getHttpRequestMethod(String s) {
+    return "";
+  }
+
+  @Override
+  public List<String> getHttpRequestHeader(String s, String name) {
+    return List.of();
+  }
+
+  @Nullable
+  @Override
+  public Integer getHttpResponseStatusCode(String s, String s2, @Nullable Throwable error) {
+    return 0;
+  }
+
+  @Override
+  public List<String> getHttpResponseHeader(String s, String s2, String name) {
+    if (name.equals("Server-Timing")) {
+      return List.of(serverTiming);
+    }
+    return List.of();
+  }
+}