Add tracing

barista/build.gradle.kts

@@ -14,6 +14,7 @@ dependencies {
     implementation("com.fasterxml.jackson.datatype:jackson-datatype-jdk8")
     implementation("com.fasterxml.jackson.datatype:jackson-datatype-jsr310")
     implementation("com.google.guava:guava")
+    implementation("io.github.markelliot.barista-tracing:barista-tracing")
     implementation("io.undertow:undertow-core")
 
     implementation("org.apache.logging.log4j:log4j-core")

barista/src/main/java/barista/Server.java

@@ -5,14 +5,18 @@
 import barista.handlers.DispatchFromIoThreadHandler;
 import barista.handlers.EndpointHandlerBuilder;
 import barista.handlers.HandlerChain;
+import barista.handlers.TracingHandler;
 import barista.tls.TransportLayerSecurity;
+import barista.tracing.Spans;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import io.undertow.Undertow;
 import java.nio.file.Paths;
 import java.util.LinkedHashSet;
 import java.util.Objects;
 import java.util.Set;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public final class Server {
     private final Undertow undertow;
@@ -51,6 +55,7 @@ public static final class Builder {
         private SerDe serde = new SerDe.ObjectMapperSerDe();
         private Authz authz = Authz.denyAll();
         private boolean tls = true;
+        private double tracingRate = 0.2;
 
         private Builder() {}
 
@@ -97,15 +102,37 @@ public Builder disableTls() {
             return this;
         }
 
+        /**
+         * Sets the sample rate to run tracing for incoming requests without a traceId header.
+         *
+         * <p>A value of 1.0 means trace every request.
+         *
+         * <p>A value of 0.0 means trace no requests.
+         */
+        public Builder tracingRate(double rate) {
+            if (rate < 0.0 || rate > 1.0) {
+                throw new IllegalArgumentException("Cannot set a rate outside of range [0, 1]");
+            }
+            this.tracingRate = rate;
+            return this;
+        }
+
         public Server start() {
             Preconditions.checkNotNull(authz);
 
+            if (tracingRate > 0.0) {
+                // TODO(markelliot): use a custom format, perhaps emit to a specific log file
+                Logger tracing = LoggerFactory.getLogger("tracing");
+                Spans.register("barista", span -> tracing.info("TRACING {}", span));
+            }
+
             EndpointHandlerBuilder handler = new EndpointHandlerBuilder(serde, authz);
             Undertow.Builder builder =
                     Undertow.builder()
                             .setHandler(
                                     HandlerChain.of(DispatchFromIoThreadHandler::new)
                                             .then(h -> new CorsHandler(allowedOrigins, h))
+                                            .then(h -> new TracingHandler(tracingRate, h))
                                             .last(handler.build(authEndpoints, openEndpoints)));
             if (tls) {
                 builder.addHttpsListener(

barista/src/main/java/barista/handlers/TracingHandler.java

@@ -0,0 +1,59 @@
+package barista.handlers;
+
+import barista.tracing.Ids;
+import barista.tracing.Span;
+import barista.tracing.Trace;
+import barista.tracing.Traces;
+import io.undertow.server.HttpHandler;
+import io.undertow.server.HttpServerExchange;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.function.Supplier;
+
+public record TracingHandler(double rate, HttpHandler delegate) implements HttpHandler {
+    @Override
+    public void handleRequest(HttpServerExchange exchange) throws Exception {
+        Trace trace = getTraceForRequest(exchange);
+        Span span = getSpanForRequest(exchange, trace);
+        exchange.addExchangeCompleteListener(
+                (ignored, next) -> {
+                    span.close();
+                    next.proceed();
+                });
+        delegate.handleRequest(exchange);
+    }
+
+    private Trace getTraceForRequest(HttpServerExchange exchange) {
+        String traceId = getId(exchange, "X-B3-TraceId");
+        if (traceId != null) {
+            return Traces.create(traceId, isSampled(exchange));
+        }
+        return Traces.create(Ids.randomId(), ThreadLocalRandom.current().nextDouble() < rate);
+    }
+
+    private static boolean isSampled(HttpServerExchange exchange) {
+        String val = exchange.getRequestHeaders().getFirst("X-B3-Sampled");
+        return val != null && (val.equals("1") || val.equalsIgnoreCase("true"));
+    }
+
+    private Span getSpanForRequest(HttpServerExchange exchange, Trace trace) {
+        Supplier<String> opName = opName(exchange);
+        String spanId = getId(exchange, "X-B3-SpanId");
+        if (spanId != null) {
+            return trace.withParent(spanId, opName);
+        }
+        return trace.rootSpan(opName);
+    }
+
+    private static Supplier<String> opName(HttpServerExchange exchange) {
+        return () -> exchange.getRequestMethod() + " " + exchange.getRelativePath();
+    }
+
+    /** Gets the value of the named header the request and returns it if safe to handle. */
+    private static String getId(HttpServerExchange exchange, String header) {
+        String val = exchange.getRequestHeaders().getFirst(header);
+        if (val != null && val.length() < 32) {
+            return val;
+        }
+        return null;
+    }
+}

build.gradle.kts

@@ -30,6 +30,13 @@ allprojects {
 
     // lives in allprojects because of consistent-versions
     repositories {
+        maven {
+            url = uri("https://maven.pkg.github.com/markelliot/barista-tracing")
+            credentials {
+                username = System.getenv("GITHUB_ACTOR")
+                password = System.getenv("GITHUB_TOKEN")
+            }
+        }
         mavenCentral()
     }
 

versions.lock

@@ -15,7 +15,8 @@ com.google.guava:failureaccess:1.0.1 (1 constraints: 140ae1b4)
 com.google.guava:guava:27.0.1-jre (2 constraints: e71c01c8)
 com.google.guava:listenablefuture:9999.0-empty-to-avoid-conflict-with-guava (1 constraints: bd17c918)
 com.google.j2objc:j2objc-annotations:1.1 (1 constraints: b609eba0)
-com.jakewharton.nopen:nopen-annotations:1.0.1 (1 constraints: 0405f135)
+com.jakewharton.nopen:nopen-annotations:1.0.1 (2 constraints: e61a7e3d)
+io.github.markelliot.barista-tracing:barista-tracing:0.1.3 (1 constraints: 0605f135)
 io.undertow:undertow-core:2.1.0.Final (1 constraints: 1d07a45a)
 org.apache.logging.log4j:log4j-api:2.14.1 (3 constraints: ea2dae0e)
 org.apache.logging.log4j:log4j-core:2.14.1 (2 constraints: 0d16b624)

versions.props

@@ -4,6 +4,7 @@ com.fasterxml.jackson.*:* = 2.12.2
 com.google.errorprone:error_prone_core = 2.7.1
 com.google.guava:guava = 27.0.1-jre
 com.jakewharton.nopen:* = 1.0.1
+io.github.markelliot.barista-tracing:barista-tracing = 0.1.3
 io.undertow:undertow-core = 2.1.0.Final
 org.apache.logging.log4j:* = 2.14.1
 org.assertj:assertj-core = 3.12.2