Update to standard tracing headers, support tracing stacks

palantir · May 16, 2016 · c2e5f38 · c2e5f38
1 parent bb5e427
commit c2e5f38
Show file tree

Hide file tree

Showing 9 changed files with 194 additions and 96 deletions.
diff --git a/http-clients/src/main/java/com/palantir/remoting/http/TraceRequestInterceptor.java b/http-clients/src/main/java/com/palantir/remoting/http/TraceRequestInterceptor.java
@@ -25,10 +25,12 @@ public final class TraceRequestInterceptor implements RequestInterceptor {
 
     @Override
     public void apply(RequestTemplate template) {
-        TraceState traceState = Traces.getTrace();
-        if (traceState != null) {
-            template.header(Traces.TRACE_HEADER, traceState.getTraceId());
+        TraceState callState = Traces.deriveTrace(template.url());
+        template.header(Traces.Headers.TRACE_ID, callState.getTraceId());
+        if (callState.getParentSpanId().isPresent()) {
+            template.header(Traces.Headers.PARENT_SPAN_ID, callState.getParentSpanId().get());
         }
+        template.header(Traces.Headers.SPAN_ID, callState.getSpanId());
     }
 
 }
diff --git a/http-clients/src/test/java/com/palantir/remoting/http/TraceRequestInterceptorTest.java b/http-clients/src/test/java/com/palantir/remoting/http/TraceRequestInterceptorTest.java
@@ -17,9 +17,11 @@
 package com.palantir.remoting.http;
 
 import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.not;
 import static org.junit.Assert.assertThat;
 
 import com.google.common.base.Optional;
+import com.palantir.tracing.TraceState;
 import com.palantir.tracing.Traces;
 import com.squareup.okhttp.mockwebserver.MockResponse;
 import com.squareup.okhttp.mockwebserver.MockWebServer;
@@ -56,18 +58,19 @@ public void testTraceRequestInterceptor_sendsAValidTraceId() throws InterruptedE
         service.get();
         RecordedRequest request = server.takeRequest();
 
-        String traceId = request.getHeader(Traces.TRACE_HEADER);
+        String traceId = request.getHeader(Traces.Headers.TRACE_ID);
         assertThat(UUID.fromString(traceId).toString(), is(traceId));
     }
 
     @Test
     public void testTraceRequestInterceptor_sendsExplicitTraceId() throws InterruptedException {
-        Traces.createTrace("op", "my trace");
+        TraceState state = Traces.deriveTrace("operation");
         service.get();
         RecordedRequest request = server.takeRequest();
 
-        String traceId = request.getHeader(Traces.TRACE_HEADER);
-        assertThat(traceId, is("my trace"));
+        assertThat(request.getHeader(Traces.Headers.TRACE_ID), is(state.getTraceId()));
+        assertThat(request.getHeader(Traces.Headers.PARENT_SPAN_ID), is(state.getSpanId()));
+        assertThat(request.getHeader(Traces.Headers.SPAN_ID), not(state.getSpanId()));
     }
 
     @Path("/")

diff --git a/http-servers/src/main/java/com/palantir/remoting/http/server/TraceEnrichingFilter.java b/http-servers/src/main/java/com/palantir/remoting/http/server/TraceEnrichingFilter.java
@@ -0,0 +1,78 @@
+/*
+ * Copyright 2016 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.remoting.http.server;
+
+import com.google.common.base.Optional;
+import com.google.common.base.Strings;
+import com.palantir.tracing.TraceState;
+import com.palantir.tracing.Traces;
+import java.io.IOException;
+import javax.ws.rs.container.ContainerRequestContext;
+import javax.ws.rs.container.ContainerRequestFilter;
+import javax.ws.rs.container.ContainerResponseContext;
+import javax.ws.rs.container.ContainerResponseFilter;
+import javax.ws.rs.core.MultivaluedMap;
+import javax.ws.rs.ext.Provider;
+
+@Provider
+public final class TraceEnrichingFilter implements ContainerRequestFilter, ContainerResponseFilter {
+
+    @Override
+    public void filter(ContainerRequestContext requestContext) throws IOException {
+        String operation = new StringBuilder(requestContext.getMethod())
+                .append(" /")
+                .append(requestContext.getUriInfo().getPath())
+                .toString();
+
+        String traceId = requestContext.getHeaderString(Traces.Headers.TRACE_ID);
+        String parentSpanId = requestContext.getHeaderString(Traces.Headers.PARENT_SPAN_ID);
+        String spanId = requestContext.getHeaderString(Traces.Headers.SPAN_ID);
+
+        if (Strings.isNullOrEmpty(traceId)) {
+            // no trace for this request, just derive a new one
+            Traces.deriveTrace(operation);
+        } else {
+            // defend against badly formed requests
+            if (spanId == null) {
+                spanId = TraceState.randomId();
+            }
+
+            Traces.setTrace(TraceState.builder()
+                    .traceId(traceId)
+                    .parentSpanId(Optional.fromNullable(parentSpanId))
+                    .spanId(spanId)
+                    .operation(operation)
+                    .build());
+        }
+    }
+
+    @Override
+    public void filter(ContainerRequestContext requestContext, ContainerResponseContext responseContext)
+            throws IOException {
+        MultivaluedMap<String, Object> headers = responseContext.getHeaders();
+        Optional<TraceState> maybeTrace = Traces.complete();
+        if (maybeTrace.isPresent()) {
+            TraceState trace = maybeTrace.get();
+            headers.putSingle(Traces.Headers.TRACE_ID, trace.getTraceId());
+            headers.putSingle(Traces.Headers.SPAN_ID, trace.getSpanId());
+            if (trace.getParentSpanId().isPresent()) {
+                headers.putSingle(Traces.Headers.PARENT_SPAN_ID, trace.getParentSpanId().get());
+            }
+        }
+    }
+
+}
diff --git a/http-servers/src/main/java/com/palantir/remoting/http/server/TraceInheritingFilter.java b/http-servers/src/main/java/com/palantir/remoting/http/server/TraceInheritingFilter.java
diff --git a/...ttp/server/TraceInheritingFilterTest.java → ...http/server/TraceEnrichingFilterTest.java b/...ttp/server/TraceInheritingFilterTest.java → ...http/server/TraceEnrichingFilterTest.java
@@ -17,15 +17,15 @@
 package com.palantir.remoting.http.server;
 
 import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.not;
+import static org.hamcrest.Matchers.nullValue;
 import static org.junit.Assert.assertThat;
 
-import com.palantir.tracing.TraceState;
 import com.palantir.tracing.Traces;
 import io.dropwizard.Application;
 import io.dropwizard.Configuration;
 import io.dropwizard.setup.Environment;
 import io.dropwizard.testing.junit.DropwizardAppRule;
-import java.util.UUID;
 import javax.ws.rs.Consumes;
 import javax.ws.rs.GET;
 import javax.ws.rs.POST;
@@ -42,7 +42,7 @@
 import org.junit.ClassRule;
 import org.junit.Test;
 
-public final class TraceInheritingFilterTest {
+public final class TraceEnrichingFilterTest {
 
     @ClassRule
     public static final DropwizardAppRule<Configuration> APP = new DropwizardAppRule<>(TracingTestServer.class,
@@ -60,53 +60,69 @@ public void before() {
 
     @Test
     public void testTraceState_withHeaderUsesTraceId() {
-        Response response = target.path("/trace").request().header(Traces.TRACE_HEADER, "abc").get();
+        Response response = target.path("/trace").request()
+                .header(Traces.Headers.TRACE_ID, "traceId")
+                .header(Traces.Headers.PARENT_SPAN_ID, "parentSpanId")
+                .header(Traces.Headers.SPAN_ID, "spanId")
+                .get();
         assertThat(response.getStatus(), is(Status.OK.getStatusCode()));
-        assertThat(response.readEntity(TraceState.class),
-                is((TraceState) TraceState.builder().traceId("abc").operation("GET /trace").build()));
+        assertThat(response.readEntity(String.class), is("GET /trace"));
+        assertThat(response.getHeaderString(Traces.Headers.TRACE_ID), is("traceId"));
+        assertThat(response.getHeaderString(Traces.Headers.PARENT_SPAN_ID), is("parentSpanId"));
+        assertThat(response.getHeaderString(Traces.Headers.SPAN_ID), is("spanId"));
     }
 
     @Test
     public void testTraceState_respectsMethod() {
-        Response response = target.path("/trace").request().header(Traces.TRACE_HEADER, "abc").post(Entity.json(""));
+        Response response = target.path("/trace").request()
+                .header(Traces.Headers.TRACE_ID, "traceId")
+                .header(Traces.Headers.PARENT_SPAN_ID, "parentSpanId")
+                .header(Traces.Headers.SPAN_ID, "spanId")
+                .post(Entity.json(""));
         assertThat(response.getStatus(), is(Status.OK.getStatusCode()));
-        assertThat(response.readEntity(TraceState.class),
-                is((TraceState) TraceState.builder().traceId("abc").operation("POST /trace").build()));
+        assertThat(response.readEntity(String.class), is("POST /trace"));
+        assertThat(response.getHeaderString(Traces.Headers.TRACE_ID), is("traceId"));
+        assertThat(response.getHeaderString(Traces.Headers.PARENT_SPAN_ID), is("parentSpanId"));
+        assertThat(response.getHeaderString(Traces.Headers.SPAN_ID), is("spanId"));
     }
 
     @Test
-    public void testTraceState_withoutHeaderGeneratesValidUuid() {
+    public void testTraceState_withoutRequestHeadersGeneratesValidTraceResponseHeaders() {
         Response response = target.path("/trace").request().get();
         assertThat(response.getStatus(), is(Status.OK.getStatusCode()));
-        TraceState state = response.readEntity(TraceState.class);
-        assertThat(UUID.fromString(state.getTraceId()).toString(), is(state.getTraceId()));
+        assertThat(response.readEntity(String.class), is("GET /trace"));
+        assertThat(response.getHeaderString(Traces.Headers.TRACE_ID), not(nullValue()));
+        assertThat(response.getHeaderString(Traces.Headers.PARENT_SPAN_ID), is(nullValue()));
+        assertThat(response.getHeaderString(Traces.Headers.SPAN_ID), not(nullValue()));
     }
 
     @Test
-    public void testTraceState_withEmptyHeaderGeneratesValidUuid() {
-        Response response = target.path("/trace").request().header(Traces.TRACE_HEADER, "").get();
+    public void testTraceState_withEmptyTraceIdGeneratesValidTraceResponseHeaders() {
+        Response response = target.path("/trace").request().header(Traces.Headers.TRACE_ID, "").get();
         assertThat(response.getStatus(), is(Status.OK.getStatusCode()));
-        TraceState state = response.readEntity(TraceState.class);
-        assertThat(UUID.fromString(state.getTraceId()).toString(), is(state.getTraceId()));
+        assertThat(response.readEntity(String.class), is("GET /trace"));
+        assertThat(response.getHeaderString(Traces.Headers.TRACE_ID), not(nullValue()));
+        assertThat(response.getHeaderString(Traces.Headers.PARENT_SPAN_ID), is(nullValue()));
+        assertThat(response.getHeaderString(Traces.Headers.SPAN_ID), not(nullValue()));
     }
 
     public static class TracingTestServer extends Application<Configuration> {
         @Override
         public final void run(Configuration config, final Environment env) throws Exception {
-            env.jersey().register(new TraceInheritingFilter());
+            env.jersey().register(new TraceEnrichingFilter());
             env.jersey().register(new TracingTestResource());
         }
     }
 
     public static final class TracingTestResource implements TracingTestService {
         @Override
-        public TraceState getTrace() {
-            return Traces.getTrace();
+        public String getTraceOperation() {
+            return Traces.getTrace().get().getOperation();
         }
 
         @Override
-        public TraceState postTrace() {
-            return Traces.getTrace();
+        public String postTraceOperation() {
+            return Traces.getTrace().get().getOperation();
         }
     }
 
@@ -116,11 +132,11 @@ public TraceState postTrace() {
     public interface TracingTestService {
         @GET
         @Path("/trace")
-        TraceState getTrace();
+        String getTraceOperation();
 
         @POST
         @Path("/trace")
-        TraceState postTrace();
+        String postTraceOperation();
     }
 
 }
diff --git a/retrofit-clients/src/main/java/com/palantir/remoting/retrofit/TraceRequestInterceptor.java b/retrofit-clients/src/main/java/com/palantir/remoting/retrofit/TraceRequestInterceptor.java
@@ -24,10 +24,12 @@ public final class TraceRequestInterceptor implements RequestInterceptor {
 
     @Override
     public void intercept(RequestFacade request) {
-        TraceState traceState = Traces.getTrace();
-        if (traceState != null) {
-            request.addHeader(Traces.TRACE_HEADER, traceState.getTraceId());
+        TraceState callState = Traces.deriveTrace("");
+        request.addHeader(Traces.Headers.TRACE_ID, callState.getTraceId());
+        if (callState.getParentSpanId().isPresent()) {
+            request.addHeader(Traces.Headers.PARENT_SPAN_ID, callState.getParentSpanId().get());
         }
+        request.addHeader(Traces.Headers.SPAN_ID, callState.getSpanId());
     }
 
 }
diff --git a/...t-clients/src/test/java/com/palantir/remoting/retrofit/TracingRequestInterceptorTest.java b/...t-clients/src/test/java/com/palantir/remoting/retrofit/TracingRequestInterceptorTest.java
@@ -17,9 +17,11 @@
 package com.palantir.remoting.retrofit;
 
 import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.not;
 import static org.junit.Assert.assertThat;
 
 import com.google.common.base.Optional;
+import com.palantir.tracing.TraceState;
 import com.palantir.tracing.Traces;
 import com.squareup.okhttp.mockwebserver.MockResponse;
 import com.squareup.okhttp.mockwebserver.MockWebServer;
@@ -56,18 +58,19 @@ public void testTraceRequestInterceptor_sendsAValidTraceId() throws InterruptedE
         service.get();
         RecordedRequest request = server.takeRequest();
 
-        String traceId = request.getHeader(Traces.TRACE_HEADER);
+        String traceId = request.getHeader(Traces.Headers.TRACE_ID);
         assertThat(UUID.fromString(traceId).toString(), is(traceId));
     }
 
     @Test
     public void testTraceRequestInterceptor_sendsExplicitTraceId() throws InterruptedException {
-        Traces.createTrace("op", "my trace");
+        TraceState state = Traces.deriveTrace("operation");
         service.get();
         RecordedRequest request = server.takeRequest();
 
-        String traceId = request.getHeader(Traces.TRACE_HEADER);
-        assertThat(traceId, is("my trace"));
+        assertThat(request.getHeader(Traces.Headers.TRACE_ID), is(state.getTraceId()));
+        assertThat(request.getHeader(Traces.Headers.PARENT_SPAN_ID), is(state.getSpanId()));
+        assertThat(request.getHeader(Traces.Headers.SPAN_ID), not(state.getSpanId()));
     }
 
     public interface TestRequestInterceptorService {

diff --git a/tracing/src/main/java/com/palantir/tracing/TraceState.java b/tracing/src/main/java/com/palantir/tracing/TraceState.java
@@ -18,6 +18,7 @@
 
 import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
 import com.fasterxml.jackson.databind.annotation.JsonSerialize;
+import com.google.common.base.Optional;
 import java.nio.ByteBuffer;
 import java.util.Random;
 import java.util.UUID;
@@ -48,11 +49,25 @@ public String getTraceId() {
         return randomId();
     }
 
+    /**
+     * Return the identifier of the parent span for the current span, if one exists.
+     */
+    public abstract Optional<String> getParentSpanId();
+
+    /**
+     * Return a globally unique identifier representing a single span within the call trace.
+     */
+    @SuppressWarnings("checkstyle:designforextension")
+    @Value.Default
+    public String getSpanId() {
+        return randomId();
+    }
+
     public static final Builder builder() {
         return new Builder();
     }
 
-    private static String randomId() {
+    public static String randomId() {
         // non-secure random generated UUID because speed is important here and security is not
         byte[] randomBytes = new byte[16];
         RANDOM.nextBytes(randomBytes);