trace locals (#1007)

Introduce trace locals. Each trace gets its own independently settable copy of the variable.

Author: bavardage

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

@@ -0,0 +1,6 @@
+type: improvement
+improvement:
+  description: Introduce trace locals. Each trace gets its own independently settable
+    copy of the variable.
+  links:
+  - https://github.com/palantir/tracing-java/pull/1007

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

@@ -0,0 +1,152 @@
+/*
+ * (c) Copyright 2023 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;
+
+import com.palantir.logsafe.Preconditions;
+import com.palantir.logsafe.exceptions.SafeIllegalArgumentException;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.function.Supplier;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/**
+ * Provides trace local variables. Each trace gets its own independent copy of the variable.
+ *
+ * A trace local is either set (has a value) or unset (has no value).
+ *
+ * Outside of a trace (i.e. when {@link Tracer#hasTraceId()} is false) then the trace local is always unset.
+ */
+public final class TraceLocal<T> {
+
+    @Nullable
+    private final Function<? super TraceLocal<?>, T> initialValue;
+
+    private TraceLocal(@Nullable Supplier<T> initialValue) {
+        if (initialValue == null) {
+            this.initialValue = null;
+        } else {
+            // eagerly transform supplier to avoid allocation per invocation
+            // (computeIfAbsent takes a Function)
+            this.initialValue = _ignored ->
+                    Preconditions.checkNotNull(initialValue.get(), "TraceLocal initial value must not be null");
+        }
+    }
+
+    public static <T> TraceLocal<T> of() {
+        return new TraceLocal<>(null);
+    }
+
+    /**
+     * Creates a trace local variable, with a way of supplying an initial value.
+     *
+     * When not already set, and this variable is accessed in a trace with the {@link #get()} method, the supplier
+     * will be invoked to supply a value (which will then be stored for future accesses).
+     *
+     * The supplier is thus normally invoked once per trace, but may be invoked again in case of subsequent
+     * invocations of {@link #remove()} followed by get.
+     */
+    public static <T> TraceLocal<T> withInitialValue(@Nonnull Supplier<T> initialValue) {
+        return new TraceLocal<>(Preconditions.checkNotNull(initialValue, "initial value supplier must not be null"));
+    }
+
+    /**
+     * Retrieve the current value of the trace local, with respect to the current trace.
+     *
+     * If there is no current trace, i.e. {@link Tracer#hasTraceId()} is null, then this will return null.
+     *
+     * If the value of this trace local has not been set for the current trace, then the supplier passed in the
+     * constructor will be called to supply a value.
+     */
+    @Nullable
+    public T get() {
+        TraceState traceState = Tracer.getTraceState();
+
+        if (traceState == null) {
+            return null;
+        }
+
+        if (initialValue == null) {
+            // not potentially setting a value, so just grab any current set value
+
+            Map<TraceLocal<?>, Object> traceLocals = traceState.getTraceLocals();
+
+            if (traceLocals == null) {
+                return null;
+            }
+
+            return (T) traceLocals.get(this);
+        } else {
+            Map<TraceLocal<?>, Object> traceLocals = traceState.getOrCreateTraceLocals();
+            return (T) traceLocals.computeIfAbsent(this, initialValue);
+        }
+    }
+
+    /**
+     * Sets the value of this trace local for the current trace.
+     *
+     * Returns the previous value of this trace local if set, or null if the value was previously unset.
+     */
+    @Nullable
+    public T set(@Nonnull T value) {
+        if (value == null) {
+            throw new SafeIllegalArgumentException("value must not be null");
+        }
+
+        TraceState traceState = Tracer.getTraceState();
+
+        if (traceState == null) {
+            return null;
+        }
+
+        return (T) traceState.getOrCreateTraceLocals().put(this, value);
+    }
+
+    /**
+     * Unsets the value of this trace local.
+     *
+     * Returns the previous value of this trace local if set, or null if the value was previously unset.
+     */
+    @Nullable
+    public T remove() {
+        TraceState traceState = Tracer.getTraceState();
+
+        if (traceState == null) {
+            return null;
+        }
+
+        Map<TraceLocal<?>, Object> traceLocals = traceState.getTraceLocals();
+        if (traceLocals == null) {
+            // no trace locals ever set, short circuit (avoid creating the trace local map)
+            return null;
+        }
+
+        return (T) traceLocals.remove(this);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        // identity semantics
+        return super.equals(obj);
+    }
+
+    @Override
+    public int hashCode() {
+        // identity semantics
+        return super.hashCode();
+    }
+}

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

@@ -21,7 +21,10 @@
 
 import com.google.common.base.Strings;
 import java.io.Serializable;
+import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 import javax.annotation.Nullable;
 
 /**
@@ -39,6 +42,12 @@ final class TraceState implements Serializable {
     @Nullable
     private final String forUserAgent;
 
+    @Nullable
+    private volatile TraceLocalMap traceLocals;
+
+    private static final AtomicReferenceFieldUpdater<TraceState, TraceLocalMap> traceLocalsUpdater =
+            AtomicReferenceFieldUpdater.newUpdater(TraceState.class, TraceLocalMap.class, "traceLocals");
+
     static TraceState of(String traceId, Optional<String> requestId, Optional<String> forUserAgent) {
         checkArgument(!Strings.isNullOrEmpty(traceId), "traceId must be non-empty");
         checkNotNull(requestId, "requestId should be not-null");
@@ -50,6 +59,7 @@ private TraceState(String traceId, @Nullable String requestId, @Nullable String
         this.traceId = traceId;
         this.requestId = requestId;
         this.forUserAgent = forUserAgent;
+        this.traceLocals = null;
     }
 
     /**
@@ -79,6 +89,24 @@ String forUserAgent() {
         return forUserAgent;
     }
 
+    Map<TraceLocal<?>, Object> getOrCreateTraceLocals() {
+        TraceLocalMap result = traceLocalsUpdater.get(this);
+
+        if (result == null) {
+            result = new TraceLocalMap();
+            if (!traceLocalsUpdater.compareAndSet(this, null, result)) {
+                return traceLocalsUpdater.get(this);
+            }
+        }
+
+        return result;
+    }
+
+    @Nullable
+    Map<TraceLocal<?>, Object> getTraceLocals() {
+        return traceLocals;
+    }
+
     @Override
     public String toString() {
         return "TraceState{"
@@ -87,4 +115,6 @@ public String toString() {
                 + "forUserAgent='" + forUserAgent
                 + "'}";
     }
+
+    private static final class TraceLocalMap extends ConcurrentHashMap<TraceLocal<?>, Object> {}
 }

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

@@ -349,6 +349,17 @@ private static Optional<String> getParentSpanId(@Nullable Trace trace) {
         return Optional.empty();
     }
 
+    @Nullable
+    static TraceState getTraceState() {
+        Trace maybeCurrentTrace = currentTrace.get();
+
+        if (maybeCurrentTrace == null) {
+            return null;
+        }
+
+        return maybeCurrentTrace.getTraceState();
+    }
+
     private static TraceState getTraceState(@Nullable Trace maybeCurrentTrace, SpanType newSpanType) {
         if (maybeCurrentTrace != null) {
             return maybeCurrentTrace.getTraceState();

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

@@ -44,4 +44,14 @@ public void testToString() {
                 .contains(span.getOperation())
                 .contains(span.getSpanId());
     }
+
+    @Test
+    public void testToString_doesNotContainTraceLocals() {
+        Trace trace = Trace.of(true, TraceState.of("traceId", Optional.empty(), Optional.empty()));
+
+        TraceLocal<String> traceLocal = TraceLocal.of();
+        trace.getTraceState().getOrCreateTraceLocals().put(traceLocal, "secret-value");
+
+        assertThat(trace.toString()).doesNotContain("secret");
+    }
 }

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

@@ -420,6 +420,34 @@ public void testHasUnobservableTrace_unobservableTrace() {
         assertThat(Tracer.hasUnobservableTrace()).isFalse();
     }
 
+    @Test
+    public void testTraceLocals() {
+        TraceLocal<String> traceLocal = TraceLocal.withInitialValue(() -> "initial");
+
+        Tracer.setSampler(AlwaysSampler.INSTANCE);
+        Tracer.fastStartSpan("outer");
+
+        try (CloseableTracer trace = CloseableTracer.startSpan("operation")) {
+            assertThat(traceLocal.get()).isEqualTo("initial");
+
+            traceLocal.set("some-value");
+            assertThat(traceLocal.get()).isEqualTo("some-value");
+        }
+
+        try (CloseableTracer trace = CloseableTracer.startSpan("operation")) {
+            assertThat(traceLocal.get()).isEqualTo("some-value");
+            traceLocal.set("other-value");
+        }
+
+        assertThat(traceLocal.get()).isEqualTo("other-value");
+
+        // outer...
+        Tracer.fastCompleteSpan();
+
+        // outside of a trace, it has no value
+        assertThat(traceLocal.get()).isEqualTo(null);
+    }
+
     @Test
     public void testSimpleDetachedTrace() {
         assertThat(Tracer.hasTraceId()).isFalse();
@@ -671,6 +699,33 @@ public void testNewDetachedTrace_doesNotModifyCurrentState() {
         }
     }
 
+    @Test
+    public void testDetachedState_traceLocals() {
+        Tracer.setSampler(AlwaysSampler.INSTANCE);
+
+        TraceLocal<String> traceLocal = TraceLocal.withInitialValue(() -> "initial");
+
+        try (CloseableTracer ignored = CloseableTracer.startSpan("test")) {
+            traceLocal.set("outer");
+
+            DetachedSpan span =
+                    DetachedSpan.start(Observability.SAMPLE, "12345", Optional.empty(), "op", SpanType.LOCAL);
+            try (CloseableSpan attached = span.attach()) {
+                assertThat(traceLocal.get()).isEqualTo("initial");
+                traceLocal.set("inner");
+            }
+
+            assertThat(traceLocal.get()).isEqualTo("outer");
+
+            try (CloseableSpan attached = span.attach()) {
+                assertThat(traceLocal.get()).isEqualTo("inner");
+            }
+
+            traceLocal.remove();
+            assertThat(traceLocal.get()).isEqualTo("initial");
+        }
+    }
+
     private static void startAndFastCompleteSpan() {
         Tracer.fastStartSpan("operation");
         Tracer.fastCompleteSpan();