[improvement] Refactor Tracer.initTrace to take enum instead of optional boolean

tracing-jersey/src/main/java/com/palantir/tracing/jersey/TraceEnrichingFilter.java

@@ -17,6 +17,7 @@
 package com.palantir.tracing.jersey;
 
 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.Span;
@@ -66,10 +67,10 @@ public void filter(ContainerRequestContext requestContext) throws IOException {
         // Set up thread-local span that inherits state from HTTP headers
         if (Strings.isNullOrEmpty(traceId)) {
             // HTTP request did not indicate a trace; initialize trace state and create a span.
-            Tracer.initTrace(hasSampledHeader(requestContext), Tracers.randomId());
+            Tracer.initTrace(getObservabilityFromHeader(requestContext), Tracers.randomId());
             Tracer.startSpan(operation, SpanType.SERVER_INCOMING);
         } else {
-            Tracer.initTrace(hasSampledHeader(requestContext), traceId);
+            Tracer.initTrace(getObservabilityFromHeader(requestContext), traceId);
             if (spanId == null) {
                 Tracer.startSpan(operation, SpanType.SERVER_INCOMING);
             } else {
@@ -101,13 +102,14 @@ public void filter(ContainerRequestContext requestContext, ContainerResponseCont
         }
     }
 
-    // Returns true iff the context contains a "1" X-B3-Sampled header, or absent if there is no such header.
-    private static Optional<Boolean> hasSampledHeader(ContainerRequestContext context) {
+    // 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(ContainerRequestContext context) {
         String header = context.getHeaderString(TraceHttpHeaders.IS_SAMPLED);
-        if (header == null) {
-            return Optional.empty();
+        if (Strings.isNullOrEmpty(header)) {
+            return Observability.UNDECIDED;
         } else {
-            return Optional.of(header.equals("1"));
+            return "1".equals(header) ? Observability.SAMPLE : Observability.DO_NOT_SAMPLE;
         }
     }
 }

tracing-okhttp3/src/test/java/com/palantir/tracing/okhttp3/OkhttpTraceInterceptorTest.java

@@ -22,6 +22,7 @@
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 
+import com.palantir.tracing.Observability;
 import com.palantir.tracing.Tracer;
 import com.palantir.tracing.Tracers;
 import com.palantir.tracing.api.OpenSpan;
@@ -30,7 +31,6 @@
 import com.palantir.tracing.api.SpanType;
 import com.palantir.tracing.api.TraceHttpHeaders;
 import java.io.IOException;
-import java.util.Optional;
 import okhttp3.Interceptor;
 import okhttp3.Request;
 import org.junit.After;
@@ -65,7 +65,7 @@ public void before() {
 
     @After
     public void after() {
-        Tracer.initTrace(Optional.of(true), Tracers.randomId());
+        Tracer.initTrace(Observability.SAMPLE, Tracers.randomId());
         Tracer.unsubscribe("");
     }
 
@@ -105,15 +105,15 @@ public void testPopulatesNewTrace_whenParentTraceIsPresent() throws IOException
 
     @Test
     public void testAddsIsSampledHeader_whenTraceIsObservable() throws IOException {
-        Tracer.initTrace(Optional.of(true), Tracers.randomId());
+        Tracer.initTrace(Observability.SAMPLE, Tracers.randomId());
         OkhttpTraceInterceptor.INSTANCE.intercept(chain);
         verify(chain).proceed(requestCaptor.capture());
         assertThat(requestCaptor.getValue().header(TraceHttpHeaders.IS_SAMPLED)).isEqualTo("1");
     }
 
     @Test
     public void testHeaders_whenTraceIsNotObservable() throws IOException {
-        Tracer.initTrace(Optional.of(false), Tracers.randomId());
+        Tracer.initTrace(Observability.DO_NOT_SAMPLE, Tracers.randomId());
         String traceId = Tracer.getTraceId();
         OkhttpTraceInterceptor.INSTANCE.intercept(chain);
         verify(chain).proceed(requestCaptor.capture());

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

@@ -19,6 +19,7 @@
 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;
@@ -28,7 +29,6 @@
 import io.undertow.util.AttachmentKey;
 import io.undertow.util.HeaderMap;
 import io.undertow.util.HttpString;
-import java.util.Optional;
 
 /**
  * Extracts Zipkin-style trace information from the given HTTP request and sets up a corresponding
@@ -50,13 +50,6 @@ public final class TracedOperationHandler implements HttpHandler {
     private static final HttpString SPAN_ID = HttpString.tryFromString(TraceHttpHeaders.SPAN_ID);
     private static final HttpString IS_SAMPLED = HttpString.tryFromString(TraceHttpHeaders.IS_SAMPLED);
 
-
-    // Pre-compute sampled values, there's no need to do this work for each request
-    @SuppressWarnings("OptionalUsedAsFieldOrParameterType")
-    private static final Optional<Boolean> SAMPLED = Optional.of(Boolean.TRUE);
-    @SuppressWarnings("OptionalUsedAsFieldOrParameterType")
-    private static final Optional<Boolean> NOT_SAMPLED = Optional.of(Boolean.FALSE);
-
     private final String operation;
     private final HttpHandler delegate;
 
@@ -79,16 +72,14 @@ public void handleRequest(HttpServerExchange exchange) throws Exception {
         }
     }
 
-    // Returns true iff the context contains a "1" X-B3-Sampled header, false if the header contains another value,
-    // or absent if there is no such header.
-    private static Optional<Boolean> hasSampledHeader(HeaderMap headers) {
+    // 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 (header == null) {
-            return Optional.empty();
+        if (Strings.isNullOrEmpty(header)) {
+            return Observability.UNDECIDED;
         } else {
-            // No need to box the resulting boolean and allocate
-            // a new Optional wrapper for each invocation.
-            return header.equals("1") ? SAMPLED : NOT_SAMPLED;
+            return "1".equals(header) ? Observability.SAMPLE : Observability.DO_NOT_SAMPLE;
         }
     }
 
@@ -109,7 +100,7 @@ private String initializeTrace(HttpServerExchange exchange) {
 
     /** Initializes trace state given a trace-id header from the client. */
     private void initializeTraceFromExisting(HeaderMap headers, String traceId) {
-        Tracer.initTrace(hasSampledHeader(headers), traceId);
+        Tracer.initTrace(getObservabilityFromHeader(headers), traceId);
         String spanId = headers.getFirst(SPAN_ID); // nullable
         if (spanId == null) {
             Tracer.startSpan(operation, SpanType.SERVER_INCOMING);
@@ -123,7 +114,7 @@ private void initializeTraceFromExisting(HeaderMap headers, String traceId) {
     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(hasSampledHeader(headers), newTraceId);
+        Tracer.initTrace(getObservabilityFromHeader(headers), newTraceId);
         Tracer.startSpan(operation, SpanType.SERVER_INCOMING);
         return newTraceId;
     }

tracing/src/main/java/com/palantir/tracing/Observability.java

@@ -0,0 +1,37 @@
+/*
+ * (c) Copyright 2019 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;
+
+/**
+ * Represents the desired observability of a new trace.
+ */
+public enum Observability {
+    /**
+     * Force the trace to be sampled.
+     */
+    SAMPLE,
+
+    /**
+     * Force the trace to not be sampled.
+     */
+    DO_NOT_SAMPLE,
+
+    /**
+     * Do not force, and let the tracer decide the observability.
+     */
+    UNDECIDED
+}

tracing/src/main/java/com/palantir/tracing/Tracer.java

@@ -24,6 +24,7 @@
 import com.google.errorprone.annotations.CheckReturnValue;
 import com.palantir.logsafe.SafeArg;
 import com.palantir.logsafe.UnsafeArg;
+import com.palantir.logsafe.exceptions.SafeIllegalArgumentException;
 import com.palantir.tracing.api.OpenSpan;
 import com.palantir.tracing.api.Span;
 import com.palantir.tracing.api.SpanObserver;
@@ -62,32 +63,46 @@ private Tracer() {}
     private static volatile TraceSampler sampler = new RandomSampler(0.01f);
 
     /**
-     * Creates a new trace, but does not set it as the current trace. The new trace is {@link Trace#isObservable
-     * observable} iff the given flag is true, or, iff {@code isObservable} is absent, if the {@link #setSampler
-     * configured sampler} returns true.
+     * Creates a new trace, but does not set it as the current trace.
      */
-    private static Trace createTrace(Optional<Boolean> isObservable, String traceId) {
+    private static Trace createTrace(Observability observability, String traceId) {
         checkArgument(!Strings.isNullOrEmpty(traceId), "traceId must be non-empty");
-        boolean observable = shouldObserve(isObservable);
+        boolean observable = shouldObserve(observability);
         return new Trace(observable, traceId);
     }
 
-    // Avoid lambda allocation on hot paths
-    @SuppressWarnings("OptionalIsPresent")
-    private static boolean shouldObserve(Optional<Boolean> isObservable) {
-        if (isObservable.isPresent()) {
-            return Boolean.TRUE.equals(isObservable.get());
+    private static boolean shouldObserve(Observability observability) {
+        switch (observability) {
+            case SAMPLE:
+                return true;
+            case DO_NOT_SAMPLE:
+                return false;
+            case UNDECIDED:
+                return sampler.sample();
         }
-        return sampler.sample();
+
+        throw new SafeIllegalArgumentException("Unknown observability", SafeArg.of("observability", observability));
     }
 
     /**
-     * Initializes the current thread's trace, erasing any previously accrued open spans. The new trace is {@link
-     * Trace#isObservable observable} iff the given flag is true, or, iff {@code isObservable} is absent, if the {@link
-     * #setSampler configured sampler} returns true.
+     * Deprecated.
+     *
+     * @deprecated Use {@link #initTrace(Observability, String)}
      */
+    @Deprecated
     public static void initTrace(Optional<Boolean> isObservable, String traceId) {
-        setTrace(createTrace(isObservable, traceId));
+        Observability observability = isObservable
+                .map(value -> Boolean.TRUE.equals(value) ? Observability.SAMPLE : Observability.DO_NOT_SAMPLE)
+                .orElse(Observability.UNDECIDED);
+
+        setTrace(createTrace(observability, traceId));
+    }
+
+    /**
+     * Initializes the current thread's trace, erasing any previously accrued open spans.
+     */
+    public static void initTrace(Observability observability, String traceId) {
+        setTrace(createTrace(observability, traceId));
     }
 
     /**
@@ -356,7 +371,7 @@ private static void setTraceSampledMdcIfObservable(boolean observable) {
     private static Trace getOrCreateCurrentTrace() {
         Trace trace = currentTrace.get();
         if (trace == null) {
-            trace = createTrace(Optional.empty(), Tracers.randomId());
+            trace = createTrace(Observability.UNDECIDED, Tracers.randomId());
             setTrace(trace);
         }
         return trace;

tracing/src/main/java/com/palantir/tracing/Tracers.java

@@ -229,7 +229,7 @@ public static <V> Callable<V> wrapWithNewTrace(String operation, Callable<V> del
             Optional<Trace> originalTrace = Tracer.getAndClearTraceIfPresent();
 
             try {
-                Tracer.initTrace(Optional.empty(), Tracers.randomId());
+                Tracer.initTrace(Observability.UNDECIDED, Tracers.randomId());
                 Tracer.startSpan(operation);
                 return delegate.call();
             } finally {
@@ -261,7 +261,7 @@ public static Runnable wrapWithNewTrace(String operation, Runnable delegate) {
             Optional<Trace> originalTrace = Tracer.getAndClearTraceIfPresent();
 
             try {
-                Tracer.initTrace(Optional.empty(), Tracers.randomId());
+                Tracer.initTrace(Observability.UNDECIDED, Tracers.randomId());
                 Tracer.startSpan(operation);
                 delegate.run();
             } finally {
@@ -295,7 +295,7 @@ public static Runnable wrapWithAlternateTraceId(String traceId, String operation
             Optional<Trace> originalTrace = Tracer.getAndClearTraceIfPresent();
 
             try {
-                Tracer.initTrace(Optional.empty(), traceId);
+                Tracer.initTrace(Observability.UNDECIDED, traceId);
                 Tracer.startSpan(operation);
                 delegate.run();
             } finally {

tracing/src/test/java/com/palantir/tracing/AsyncSlf4jSpanObserverTest.java

@@ -39,7 +39,6 @@
 import java.net.InetAddress;
 import java.nio.charset.StandardCharsets;
 import java.util.Map;
-import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 import java.util.stream.Collectors;
 import org.jmock.lib.concurrent.DeterministicScheduler;
@@ -99,7 +98,7 @@ public void testJsonFormatToLog() throws Exception {
         DeterministicScheduler executor = new DeterministicScheduler();
         Tracer.subscribe(TEST_OBSERVER, AsyncSlf4jSpanObserver.of(
                 "serviceName", Inet4Address.getLoopbackAddress(), logger, executor));
-        Tracer.initTrace(Optional.of(true), Tracers.randomId());
+        Tracer.initTrace(Observability.SAMPLE, Tracers.randomId());
         Tracer.startSpan("operation");
         Span span = Tracer.completeSpan().get();
         verify(appender, never()).doAppend(any(ILoggingEvent.class)); // async logger only fires when executor runs
@@ -120,7 +119,7 @@ public void testJsonFormatToLog() throws Exception {
     public void testDefaultConstructorDeterminesIpAddress() throws Exception {
         DeterministicScheduler executor = new DeterministicScheduler();
         Tracer.subscribe(TEST_OBSERVER, AsyncSlf4jSpanObserver.of("serviceName", executor));
-        Tracer.initTrace(Optional.of(true), Tracers.randomId());
+        Tracer.initTrace(Observability.SAMPLE, Tracers.randomId());
         Tracer.startSpan("operation");
         Span span = Tracer.completeSpan().get();
 

tracing/src/test/java/com/palantir/tracing/AsyncTracerTest.java

@@ -20,13 +20,12 @@
 
 import com.google.common.collect.Lists;
 import java.util.List;
-import java.util.Optional;
 import org.junit.Test;
 
 public class AsyncTracerTest {
     @Test
     public void doesNotLeakEnqueueSpan() {
-        Tracer.initTrace(Optional.empty(), "defaultTraceId");
+        Tracer.initTrace(Observability.UNDECIDED, "defaultTraceId");
         Trace originalTrace = getTrace();
         AsyncTracer deferredTracer = new AsyncTracer();
         assertThat(originalTrace.top()).isEmpty();
@@ -42,7 +41,7 @@ public void doesNotLeakEnqueueSpan() {
 
     @Test
     public void completesBothDeferredSpans() {
-        Tracer.initTrace(Optional.of(true), "defaultTraceId");
+        Tracer.initTrace(Observability.SAMPLE, "defaultTraceId");
         Tracer.startSpan("defaultSpan");
         AsyncTracer asyncTracer = new AsyncTracer();
         List<String> observedSpans = Lists.newArrayList();
@@ -57,7 +56,7 @@ public void completesBothDeferredSpans() {
 
     @Test
     public void preservesState() {
-        Tracer.initTrace(Optional.empty(), "defaultTraceId");
+        Tracer.initTrace(Observability.UNDECIDED, "defaultTraceId");
         Tracer.startSpan("foo");
         Tracer.startSpan("bar");
         Tracer.startSpan("baz");

tracing/src/test/java/com/palantir/tracing/DeferredTracerTest.java

@@ -23,14 +23,13 @@
 import java.io.IOException;
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
-import java.util.Optional;
 import org.junit.Test;
 
 public class DeferredTracerTest {
 
     @Test
     public void testIsSerializable() throws IOException, ClassNotFoundException {
-        Tracer.initTrace(Optional.empty(), "defaultTraceId");
+        Tracer.initTrace(Observability.UNDECIDED, "defaultTraceId");
 
         DeferredTracer deferredTracer = new DeferredTracer("operation");
 
@@ -39,7 +38,7 @@ public void testIsSerializable() throws IOException, ClassNotFoundException {
             objectOutputStream.writeObject(deferredTracer);
         }
 
-        Tracer.initTrace(Optional.empty(), "someOtherTraceId");
+        Tracer.initTrace(Observability.UNDECIDED, "someOtherTraceId");
 
         ByteArrayInputStream bais = new ByteArrayInputStream(baos.toByteArray());
         try (ObjectInputStream objectInputStream = new ObjectInputStream(bais)) {