Provide Tracing support for the full request lifecycle (#400)

Provide Tracing support for the full Undertow request lifecycle

changelog/@unreleased/pr-400.v2.yml

@@ -0,0 +1,5 @@
+type: feature
+feature:
+  description: Provide Tracing support for the full Undertow request lifecycle
+  links:
+  - https://github.com/palantir/tracing-java/pull/400

tracing-undertow/src/main/java/com/palantir/tracing/undertow/TracedOperationHandler.java

@@ -18,17 +18,11 @@
 
 import static com.palantir.logsafe.Preconditions.checkNotNull;
 
-import com.google.common.base.Strings;
-import com.palantir.tracing.Observability;
-import com.palantir.tracing.Tracer;
-import com.palantir.tracing.Tracers;
-import com.palantir.tracing.api.SpanType;
-import com.palantir.tracing.api.TraceHttpHeaders;
+import com.palantir.tracing.CloseableSpan;
+import com.palantir.tracing.DetachedSpan;
 import io.undertow.server.HttpHandler;
 import io.undertow.server.HttpServerExchange;
 import io.undertow.util.AttachmentKey;
-import io.undertow.util.HeaderMap;
-import io.undertow.util.HttpString;
 
 /**
  * Extracts Zipkin-style trace information from the given HTTP request and sets up a corresponding {@link
@@ -41,12 +35,12 @@
  * handler may register.
  */
 public final class TracedOperationHandler implements HttpHandler {
-    /** Attachment to check whether the current request is being traced. */
-    public static final AttachmentKey<Boolean> IS_SAMPLED_ATTACHMENT = AttachmentKey.create(Boolean.class);
-
-    private static final HttpString TRACE_ID = HttpString.tryFromString(TraceHttpHeaders.TRACE_ID);
-    private static final HttpString SPAN_ID = HttpString.tryFromString(TraceHttpHeaders.SPAN_ID);
-    private static final HttpString IS_SAMPLED = HttpString.tryFromString(TraceHttpHeaders.IS_SAMPLED);
+    /**
+     * Attachment to check whether the current request is being traced.
+     * @deprecated in favor of {@link TracingAttachments#IS_SAMPLED}
+     */
+    @Deprecated
+    public static final AttachmentKey<Boolean> IS_SAMPLED_ATTACHMENT = TracingAttachments.IS_SAMPLED;
 
     private final String operation;
     private final HttpHandler delegate;
@@ -58,63 +52,10 @@ public TracedOperationHandler(HttpHandler delegate, String operation) {
 
     @Override
     public void handleRequest(HttpServerExchange exchange) throws Exception {
-        String traceId = initializeTrace(exchange);
-
-        // Populate response before calling delegate since delegate might commit the response.
-        exchange.getResponseHeaders().put(TRACE_ID, traceId);
-        exchange.putAttachment(IS_SAMPLED_ATTACHMENT, Tracer.isTraceObservable());
-        try {
+        DetachedSpan detachedSpan = UndertowTracing.getOrInitializeRequestTrace(exchange);
+        try (CloseableSpan ignored = detachedSpan.childSpan(operation)) {
             delegate.handleRequest(exchange);
-        } finally {
-            Tracer.fastCompleteSpan();
-        }
-    }
-
-    // Force sample iff the context contains a "1" X-B3-Sampled header, force not sample if the header contains another
-    // non-empty value, or undecided if there is no such header or the header is empty.
-    private static Observability getObservabilityFromHeader(HeaderMap headers) {
-        String header = headers.getFirst(IS_SAMPLED);
-        if (Strings.isNullOrEmpty(header)) {
-            return Observability.UNDECIDED;
-        } else {
-            return "1".equals(header) ? Observability.SAMPLE : Observability.DO_NOT_SAMPLE;
-        }
-    }
-
-    /** Initializes trace state and a root span for this request, returning the traceId. */
-    private String initializeTrace(HttpServerExchange exchange) {
-        HeaderMap headers = exchange.getRequestHeaders();
-        // TODO(rfink): Log/warn if we find multiple headers?
-        String traceId = headers.getFirst(TRACE_ID); // nullable
-
-        // Set up thread-local span that inherits state from HTTP headers
-        if (Strings.isNullOrEmpty(traceId)) {
-            return initializeNewTrace(headers);
-        } else {
-            initializeTraceFromExisting(headers, traceId);
         }
-        return traceId;
-    }
-
-    /** Initializes trace state given a trace-id header from the client. */
-    private void initializeTraceFromExisting(HeaderMap headers, String traceId) {
-        Tracer.initTrace(getObservabilityFromHeader(headers), traceId);
-        String spanId = headers.getFirst(SPAN_ID); // nullable
-        if (spanId == null) {
-            Tracer.fastStartSpan(operation, SpanType.SERVER_INCOMING);
-        } else {
-            // caller's span is this span's parent.
-            Tracer.fastStartSpan(operation, spanId, SpanType.SERVER_INCOMING);
-        }
-    }
-
-    /** Initializes trace state for a request without tracing headers. */
-    private String initializeNewTrace(HeaderMap headers) {
-        // HTTP request did not indicate a trace; initialize trace state and create a span.
-        String newTraceId = Tracers.randomId();
-        Tracer.initTrace(getObservabilityFromHeader(headers), newTraceId);
-        Tracer.fastStartSpan(operation, SpanType.SERVER_INCOMING);
-        return newTraceId;
     }
 
     @Override

tracing-undertow/src/main/java/com/palantir/tracing/undertow/TracedRequestHandler.java

@@ -0,0 +1,52 @@
+/*
+ * (c) Copyright 2020 Palantir Technologies Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.palantir.tracing.undertow;
+
+import com.palantir.tracing.DetachedSpan;
+import io.undertow.server.HttpHandler;
+import io.undertow.server.HttpServerExchange;
+
+/**
+ * Extracts Zipkin-style trace information from the given HTTP request and sets up a corresponding {@link
+ * DetachedSpan} to span the entire request.
+ * See <a href="https://github.com/openzipkin/b3-propagation">b3-propagation</a>.
+ *
+ * This handler should be registered as early as possible in the request lifecycle to fully encapsulate
+ * all work.
+ *
+ * If this handler is registered multiple times in the handler chain, subsequent executions are
+ * ignored to preserve the first, most accurate span.
+ */
+public final class TracedRequestHandler implements HttpHandler {
+
+    private final HttpHandler delegate;
+
+    public TracedRequestHandler(HttpHandler delegate) {
+        this.delegate = delegate;
+    }
+
+    @Override
+    public void handleRequest(HttpServerExchange exchange) throws Exception {
+        UndertowTracing.getOrInitializeRequestTrace(exchange);
+        delegate.handleRequest(exchange);
+    }
+
+    @Override
+    public String toString() {
+        return "TracedRequestHandler{delegate=" + delegate + '}';
+    }
+}

tracing-undertow/src/main/java/com/palantir/tracing/undertow/TracingAttachments.java

@@ -0,0 +1,28 @@
+/*
+ * (c) Copyright 2020 Palantir Technologies Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.palantir.tracing.undertow;
+
+import io.undertow.util.AttachmentKey;
+
+/** Provides public tracing {@link AttachmentKey attachment keys}. */
+public final class TracingAttachments {
+
+    /** Attachment to check whether the current request is being traced. */
+    public static final AttachmentKey<Boolean> IS_SAMPLED = AttachmentKey.create(Boolean.class);
+
+    private TracingAttachments() {}
+}

tracing-undertow/src/main/java/com/palantir/tracing/undertow/UndertowTracing.java

@@ -0,0 +1,124 @@
+/*
+ * (c) Copyright 2020 Palantir Technologies Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.palantir.tracing.undertow;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Strings;
+import com.palantir.tracing.DetachedSpan;
+import com.palantir.tracing.InternalTracers;
+import com.palantir.tracing.Observability;
+import com.palantir.tracing.Tracers;
+import com.palantir.tracing.api.SpanType;
+import com.palantir.tracing.api.TraceHttpHeaders;
+import io.undertow.server.ExchangeCompletionListener;
+import io.undertow.server.HttpServerExchange;
+import io.undertow.util.AttachmentKey;
+import io.undertow.util.HeaderMap;
+import io.undertow.util.HttpString;
+import java.util.Optional;
+
+/**
+ * Internal utility functionality shared between {@link TracedOperationHandler} and {@link TracedRequestHandler}.
+ * Intentionally package private.
+ */
+final class UndertowTracing {
+
+    // Tracing header definitions
+    private static final HttpString TRACE_ID = HttpString.tryFromString(TraceHttpHeaders.TRACE_ID);
+    private static final HttpString SPAN_ID = HttpString.tryFromString(TraceHttpHeaders.SPAN_ID);
+    private static final HttpString IS_SAMPLED = HttpString.tryFromString(TraceHttpHeaders.IS_SAMPLED);
+
+    // Consider moving this to TracingAttachments and making it public. For now it's well encapsulated
+    // here because we expect the two handler implementations to be sufficient.
+    /** Detached span object representing the entire request including asynchronous components. */
+    @VisibleForTesting
+    static final AttachmentKey<DetachedSpan> REQUEST_SPAN = AttachmentKey.create(DetachedSpan.class);
+
+    private static final String OPERATION_NAME = "Undertow Request";
+
+    /**
+     * Apply detached tracing state to the provided {@link HttpServerExchange request}.
+     */
+    static DetachedSpan getOrInitializeRequestTrace(HttpServerExchange exchange) {
+        DetachedSpan detachedSpan = exchange.getAttachment(REQUEST_SPAN);
+        if (detachedSpan == null) {
+            return initializeRequestTrace(exchange);
+        }
+        return detachedSpan;
+    }
+
+    private static DetachedSpan initializeRequestTrace(HttpServerExchange exchange) {
+        HeaderMap requestHeaders = exchange.getRequestHeaders();
+        String maybeTraceId = requestHeaders.getFirst(TRACE_ID);
+        boolean newTraceId = maybeTraceId == null;
+        String traceId = newTraceId ? Tracers.randomId() : maybeTraceId;
+        DetachedSpan detachedSpan = detachedSpan(newTraceId, traceId, requestHeaders);
+        setExchangeState(exchange, detachedSpan, traceId);
+        return detachedSpan;
+    }
+
+    private static void setExchangeState(HttpServerExchange exchange, DetachedSpan detachedSpan, String traceId) {
+        // Populate response before proceeding since later operations might commit the response.
+        exchange.getResponseHeaders().put(TRACE_ID, traceId);
+        exchange.putAttachment(TracingAttachments.IS_SAMPLED, InternalTracers.isSampled(detachedSpan));
+        exchange.putAttachment(REQUEST_SPAN, detachedSpan);
+        exchange.addExchangeCompleteListener(DetachedTraceCompletionListener.INSTANCE);
+    }
+
+    private static DetachedSpan detachedSpan(
+            boolean newTrace,
+            String traceId,
+            HeaderMap requestHeaders) {
+        return DetachedSpan.start(
+                getObservabilityFromHeader(requestHeaders),
+                traceId,
+                newTrace ? Optional.empty() : Optional.ofNullable(requestHeaders.getFirst(SPAN_ID)),
+                OPERATION_NAME,
+                SpanType.SERVER_INCOMING);
+    }
+
+    private enum DetachedTraceCompletionListener implements ExchangeCompletionListener {
+        INSTANCE;
+
+        @Override
+        public void exchangeEvent(HttpServerExchange exchange, NextListener nextListener) {
+            try {
+                DetachedSpan detachedSpan = exchange.getAttachment(REQUEST_SPAN);
+                if (detachedSpan != null) {
+                    detachedSpan.complete();
+                }
+            } finally {
+                nextListener.proceed();
+            }
+        }
+    }
+
+    /**
+     * Force sample iff the context contains a "1" X-B3-Sampled header, force not sample if the header contains another
+     * non-empty value, or undecided if there is no such header or the header is empty.
+     */
+    private static Observability getObservabilityFromHeader(HeaderMap headers) {
+        String header = headers.getFirst(IS_SAMPLED);
+        if (Strings.isNullOrEmpty(header)) {
+            return Observability.UNDECIDED;
+        } else {
+            return "1".equals(header) ? Observability.SAMPLE : Observability.DO_NOT_SAMPLE;
+        }
+    }
+
+    private UndertowTracing() {}
+}