From 83176e3d075096b772d521496291be7a7f04c2f5 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Wed, 18 Dec 2024 18:23:12 +0000 Subject: [PATCH] Introduce a shared context component, independent of tracing. Co-authored-by: Bruce Bujon --- components/context/build.gradle.kts | 9 + .../main/java/datadog/context/Context.java | 100 ++++++++ .../java/datadog/context/ContextBinder.java | 34 +++ .../main/java/datadog/context/ContextKey.java | 47 ++++ .../java/datadog/context/ContextManager.java | 50 ++++ .../datadog/context/ContextProviders.java | 16 ++ .../java/datadog/context/ContextScope.java | 12 + .../java/datadog/context/EmptyContext.java | 16 ++ .../datadog/context/ImplicitContextKeyed.java | 12 + .../java/datadog/context/IndexedContext.java | 46 ++++ .../datadog/context/SingletonContext.java | 51 ++++ .../context/ThreadLocalContextManager.java | 56 +++++ .../datadog/context/WeakMapContextBinder.java | 29 +++ .../datadog/context/ContextBinderTest.java | 22 ++ .../java/datadog/context/ContextKeyTest.java | 34 +++ .../datadog/context/ContextManagerTest.java | 219 ++++++++++++++++++ .../context/ContextProviderForkedTest.java | 86 +++++++ .../java/datadog/context/ContextTest.java | 113 +++++++++ settings.gradle | 2 + 19 files changed, 954 insertions(+) create mode 100644 components/context/build.gradle.kts create mode 100644 components/context/src/main/java/datadog/context/Context.java create mode 100644 components/context/src/main/java/datadog/context/ContextBinder.java create mode 100644 components/context/src/main/java/datadog/context/ContextKey.java create mode 100644 components/context/src/main/java/datadog/context/ContextManager.java create mode 100644 components/context/src/main/java/datadog/context/ContextProviders.java create mode 100644 components/context/src/main/java/datadog/context/ContextScope.java create mode 100644 components/context/src/main/java/datadog/context/EmptyContext.java create mode 100644 components/context/src/main/java/datadog/context/ImplicitContextKeyed.java create mode 100644 components/context/src/main/java/datadog/context/IndexedContext.java create mode 100644 components/context/src/main/java/datadog/context/SingletonContext.java create mode 100644 components/context/src/main/java/datadog/context/ThreadLocalContextManager.java create mode 100644 components/context/src/main/java/datadog/context/WeakMapContextBinder.java create mode 100644 components/context/src/test/java/datadog/context/ContextBinderTest.java create mode 100644 components/context/src/test/java/datadog/context/ContextKeyTest.java create mode 100644 components/context/src/test/java/datadog/context/ContextManagerTest.java create mode 100644 components/context/src/test/java/datadog/context/ContextProviderForkedTest.java create mode 100644 components/context/src/test/java/datadog/context/ContextTest.java diff --git a/components/context/build.gradle.kts b/components/context/build.gradle.kts new file mode 100644 index 00000000000..4dca7fc3036 --- /dev/null +++ b/components/context/build.gradle.kts @@ -0,0 +1,9 @@ +plugins { + id("me.champeau.jmh") +} + +apply(from = "$rootDir/gradle/java.gradle") + +jmh { + version = "1.28" +} diff --git a/components/context/src/main/java/datadog/context/Context.java b/components/context/src/main/java/datadog/context/Context.java new file mode 100644 index 00000000000..fcdc4d418d1 --- /dev/null +++ b/components/context/src/main/java/datadog/context/Context.java @@ -0,0 +1,100 @@ +package datadog.context; + +import static datadog.context.ContextProviders.binder; +import static datadog.context.ContextProviders.manager; + +import javax.annotation.Nullable; + +/** Immutable context scoped to an execution unit or carrier object. */ +public interface Context { + + /** Returns the root context. */ + static Context root() { + return manager().root(); + } + + /** + * Returns the context attached to the current execution unit. + * + * @return Attached context; {@link #root()} if there is none + */ + static Context current() { + return manager().current(); + } + + /** + * Attaches this context to the current execution unit. + * + * @return Scope to be closed when the context is invalid. + */ + default ContextScope attach() { + return manager().attach(this); + } + + /** + * Swaps this context with the one attached to current execution unit. + * + * @return Previously attached context; {@link #root()} if there was none + */ + default Context swap() { + return manager().swap(this); + } + + /** + * Detaches the context attached to the current execution unit, leaving it context-less. + * + *

WARNING: prefer {@link ContextScope#close()} to properly restore the surrounding context. + * + * @return Previously attached context; {@link #root()} if there was none + */ + static Context detach() { + return manager().detach(); + } + + /** + * Returns the context attached to the given carrier object. + * + * @return Attached context; {@link #root()} if there is none + */ + static Context from(Object carrier) { + return binder().from(carrier); + } + + /** Attaches this context to the given carrier object. */ + default void attachTo(Object carrier) { + binder().attachTo(carrier, this); + } + + /** + * Detaches the context attached to the given carrier object, leaving it context-less. + * + * @return Previously attached context; {@link #root()} if there was none + */ + static Context detachFrom(Object carrier) { + return binder().detachFrom(carrier); + } + + /** + * Gets the value stored in this context under the given key. + * + * @return Value stored under the key; {@code null} if there is no value. + */ + @Nullable + T get(ContextKey key); + + /** + * Creates a new context with the given key-value mapping. + * + * @return New context with the key-value mapping. + */ + Context with(ContextKey key, T value); + + /** + * Creates a new context with a value that has its own implicit key. + * + * @return New context with the implicitly keyed value. + */ + default Context with(ImplicitContextKeyed value) { + return value.storeInto(this); + } +} diff --git a/components/context/src/main/java/datadog/context/ContextBinder.java b/components/context/src/main/java/datadog/context/ContextBinder.java new file mode 100644 index 00000000000..bfcc85a887a --- /dev/null +++ b/components/context/src/main/java/datadog/context/ContextBinder.java @@ -0,0 +1,34 @@ +package datadog.context; + +/** Binds context to carrier objects. */ +public interface ContextBinder { + + /** + * Returns the context attached to the given carrier object. + * + * @return Attached context; {@link Context#root()} if there is none + */ + Context from(Object carrier); + + /** Attaches the given context to the given carrier object. */ + void attachTo(Object carrier, Context context); + + /** + * Detaches the context attached to the given carrier object, leaving it context-less. + * + * @return Previously attached context; {@link Context#root()} if there was none + */ + Context detachFrom(Object carrier); + + /** Requests use of a custom {@link ContextBinder}. */ + static void register(ContextBinder binder) { + ContextProviders.customBinder = binder; + } + + final class Provided { + static final ContextBinder INSTANCE = + null != ContextProviders.customBinder + ? ContextProviders.customBinder + : new WeakMapContextBinder(); + } +} diff --git a/components/context/src/main/java/datadog/context/ContextKey.java b/components/context/src/main/java/datadog/context/ContextKey.java new file mode 100644 index 00000000000..41b79a42c52 --- /dev/null +++ b/components/context/src/main/java/datadog/context/ContextKey.java @@ -0,0 +1,47 @@ +package datadog.context; + +import java.util.concurrent.atomic.AtomicInteger; + +/** + * Key for indexing values of type {@link T} stored in a {@link Context}. + * + *

Keys are compared by identity rather than by name. Each stored context type should either + * share its key for re-use or implement {@link ImplicitContextKeyed} to keep its key private. + */ +public final class ContextKey { + private static final AtomicInteger NEXT_INDEX = new AtomicInteger(0); + + private final String name; + final int index; + + private ContextKey(String name) { + this.name = name; + this.index = NEXT_INDEX.getAndIncrement(); + } + + /** Creates a new key with the given name. */ + public static ContextKey named(String name) { + return new ContextKey<>(name); + } + + @Override + public int hashCode() { + return index; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } else if (o == null || getClass() != o.getClass()) { + return false; + } else { + return index == ((ContextKey) o).index; + } + } + + @Override + public String toString() { + return name; + } +} diff --git a/components/context/src/main/java/datadog/context/ContextManager.java b/components/context/src/main/java/datadog/context/ContextManager.java new file mode 100644 index 00000000000..e8d45a8ec86 --- /dev/null +++ b/components/context/src/main/java/datadog/context/ContextManager.java @@ -0,0 +1,50 @@ +package datadog.context; + +/** Manages context across execution units. */ +public interface ContextManager { + + /** Returns the root context. */ + Context root(); + + /** + * Returns the context attached to the current execution unit. + * + * @return Attached context; {@link #root()} if there is none + */ + Context current(); + + /** + * Attaches the given context to the current execution unit. + * + * @return Scope to be closed when the context is invalid. + */ + ContextScope attach(Context context); + + /** + * Swaps the given context with the one attached to current execution unit. + * + * @return Previously attached context; {@link #root()} if there was none + */ + Context swap(Context context); + + /** + * Detaches the context attached to the current execution unit, leaving it context-less. + * + *

WARNING: prefer {@link ContextScope#close()} to properly restore the surrounding context. + * + * @return Previously attached context; {@link #root()} if there was none + */ + Context detach(); + + /** Requests use of a custom {@link ContextManager}. */ + static void register(ContextManager manager) { + ContextProviders.customManager = manager; + } + + final class Provided { + static final ContextManager INSTANCE = + null != ContextProviders.customManager + ? ContextProviders.customManager + : new ThreadLocalContextManager(); + } +} diff --git a/components/context/src/main/java/datadog/context/ContextProviders.java b/components/context/src/main/java/datadog/context/ContextProviders.java new file mode 100644 index 00000000000..e43d989bedc --- /dev/null +++ b/components/context/src/main/java/datadog/context/ContextProviders.java @@ -0,0 +1,16 @@ +package datadog.context; + +/** Provides {@link ContextManager} and {@link ContextBinder} implementations. */ +final class ContextProviders { + + static volatile ContextManager customManager; + static volatile ContextBinder customBinder; + + static ContextManager manager() { + return ContextManager.Provided.INSTANCE; // may be overridden by instrumentation + } + + static ContextBinder binder() { + return ContextBinder.Provided.INSTANCE; // may be overridden by instrumentation + } +} diff --git a/components/context/src/main/java/datadog/context/ContextScope.java b/components/context/src/main/java/datadog/context/ContextScope.java new file mode 100644 index 00000000000..3048d00b37a --- /dev/null +++ b/components/context/src/main/java/datadog/context/ContextScope.java @@ -0,0 +1,12 @@ +package datadog.context; + +/** Controls the validity of context attached to an execution unit. */ +public interface ContextScope extends AutoCloseable { + + /** Returns the context controlled by this scope. */ + Context context(); + + /** Detaches the context from the execution unit. */ + @Override + void close(); +} diff --git a/components/context/src/main/java/datadog/context/EmptyContext.java b/components/context/src/main/java/datadog/context/EmptyContext.java new file mode 100644 index 00000000000..ff1599fa4ee --- /dev/null +++ b/components/context/src/main/java/datadog/context/EmptyContext.java @@ -0,0 +1,16 @@ +package datadog.context; + +/** {@link Context} containing no values. */ +final class EmptyContext implements Context { + static final Context INSTANCE = new EmptyContext(); + + @Override + public T get(ContextKey key) { + return null; + } + + @Override + public Context with(ContextKey key, T value) { + return new SingletonContext(key.index, value); + } +} diff --git a/components/context/src/main/java/datadog/context/ImplicitContextKeyed.java b/components/context/src/main/java/datadog/context/ImplicitContextKeyed.java new file mode 100644 index 00000000000..a822158984b --- /dev/null +++ b/components/context/src/main/java/datadog/context/ImplicitContextKeyed.java @@ -0,0 +1,12 @@ +package datadog.context; + +/** {@link Context} value that has its own implicit {@link ContextKey}. */ +public interface ImplicitContextKeyed { + + /** + * Creates a new context with this value under its chosen key. + * + * @return New context with the implicitly keyed value. + */ + Context storeInto(Context context); +} diff --git a/components/context/src/main/java/datadog/context/IndexedContext.java b/components/context/src/main/java/datadog/context/IndexedContext.java new file mode 100644 index 00000000000..89bd29b6200 --- /dev/null +++ b/components/context/src/main/java/datadog/context/IndexedContext.java @@ -0,0 +1,46 @@ +package datadog.context; + +import static java.lang.Math.max; +import static java.util.Arrays.copyOfRange; + +import java.util.Arrays; + +/** {@link Context} containing many values. */ +final class IndexedContext implements Context { + private final Object[] store; + + IndexedContext(Object[] store) { + this.store = store; + } + + @Override + @SuppressWarnings("unchecked") + public T get(ContextKey key) { + int index = key.index; + return index < store.length ? (T) store[index] : null; + } + + @Override + public Context with(ContextKey key, T value) { + int index = key.index; + Object[] newStore = copyOfRange(store, 0, max(store.length, index + 1)); + newStore[index] = value; + + return new IndexedContext(newStore); + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + IndexedContext that = (IndexedContext) o; + return Arrays.equals(store, that.store); + } + + @Override + public int hashCode() { + int result = 31; + result = 31 * result + Arrays.hashCode(store); + return result; + } +} diff --git a/components/context/src/main/java/datadog/context/SingletonContext.java b/components/context/src/main/java/datadog/context/SingletonContext.java new file mode 100644 index 00000000000..b65dffc63ae --- /dev/null +++ b/components/context/src/main/java/datadog/context/SingletonContext.java @@ -0,0 +1,51 @@ +package datadog.context; + +import static java.lang.Math.max; + +import java.util.Objects; + +/** {@link Context} containing a single value. */ +final class SingletonContext implements Context { + private final int index; + private final Object value; + + SingletonContext(int index, Object value) { + this.index = index; + this.value = value; + } + + @Override + @SuppressWarnings("unchecked") + public V get(ContextKey key) { + return index == key.index ? (V) value : null; + } + + @Override + public Context with(ContextKey secondKey, V secondValue) { + int secondIndex = secondKey.index; + if (index == secondIndex) { + return new SingletonContext(index, secondValue); + } else { + Object[] store = new Object[max(index, secondIndex) + 1]; + store[index] = value; + store[secondIndex] = secondValue; + return new IndexedContext(store); + } + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + SingletonContext that = (SingletonContext) o; + return index == that.index && Objects.equals(value, that.value); + } + + @Override + public int hashCode() { + int result = 31; + result = 31 * result + index; + result = 31 * result + Objects.hashCode(value); + return result; + } +} diff --git a/components/context/src/main/java/datadog/context/ThreadLocalContextManager.java b/components/context/src/main/java/datadog/context/ThreadLocalContextManager.java new file mode 100644 index 00000000000..d8fab846111 --- /dev/null +++ b/components/context/src/main/java/datadog/context/ThreadLocalContextManager.java @@ -0,0 +1,56 @@ +package datadog.context; + +/** {@link ContextManager} that uses a {@link ThreadLocal} to track context per thread. */ +final class ThreadLocalContextManager implements ContextManager { + + private static final ThreadLocal CURRENT_HOLDER = + ThreadLocal.withInitial(() -> new Context[] {EmptyContext.INSTANCE}); + + @Override + public Context root() { + return EmptyContext.INSTANCE; + } + + @Override + public Context current() { + return CURRENT_HOLDER.get()[0]; + } + + @Override + public ContextScope attach(Context context) { + + Context[] holder = CURRENT_HOLDER.get(); + Context previous = holder[0]; + holder[0] = context; + + return new ContextScope() { + private boolean closed; + + @Override + public Context context() { + return context; + } + + @Override + public void close() { + if (!closed && context == holder[0]) { + holder[0] = previous; + closed = true; + } + } + }; + } + + @Override + public Context swap(Context context) { + Context[] holder = CURRENT_HOLDER.get(); + Context previous = holder[0]; + holder[0] = context; + return previous; + } + + @Override + public Context detach() { + return swap(root()); + } +} diff --git a/components/context/src/main/java/datadog/context/WeakMapContextBinder.java b/components/context/src/main/java/datadog/context/WeakMapContextBinder.java new file mode 100644 index 00000000000..144a62c1c35 --- /dev/null +++ b/components/context/src/main/java/datadog/context/WeakMapContextBinder.java @@ -0,0 +1,29 @@ +package datadog.context; + +import java.util.Collections; +import java.util.Map; +import java.util.WeakHashMap; + +/** {@link ContextBinder} that uses a global weak map of carriers to contexts. */ +final class WeakMapContextBinder implements ContextBinder { + + private static final Map TRACKED = + Collections.synchronizedMap(new WeakHashMap<>()); + + @Override + public Context from(Object carrier) { + Context bound = TRACKED.get(carrier); + return null != bound ? bound : Context.root(); + } + + @Override + public void attachTo(Object carrier, Context context) { + TRACKED.put(carrier, context); + } + + @Override + public Context detachFrom(Object carrier) { + Context previous = TRACKED.remove(carrier); + return null != previous ? previous : Context.root(); + } +} diff --git a/components/context/src/test/java/datadog/context/ContextBinderTest.java b/components/context/src/test/java/datadog/context/ContextBinderTest.java new file mode 100644 index 00000000000..6aabffba095 --- /dev/null +++ b/components/context/src/test/java/datadog/context/ContextBinderTest.java @@ -0,0 +1,22 @@ +package datadog.context; + +import static datadog.context.Context.root; +import static datadog.context.ContextTest.STRING_KEY; +import static org.junit.jupiter.api.Assertions.assertEquals; + +import org.junit.jupiter.api.Test; + +class ContextBinderTest { + + @Test + void testAttachAndDetach() { + Context context = root().with(STRING_KEY, "value"); + Object carrier = new Object(); + assertEquals(root(), Context.from(carrier)); + context.attachTo(carrier); + assertEquals(context, Context.from(carrier)); + // Detaching removes all context + assertEquals(context, Context.detachFrom(carrier)); + assertEquals(root(), Context.from(carrier)); + } +} diff --git a/components/context/src/test/java/datadog/context/ContextKeyTest.java b/components/context/src/test/java/datadog/context/ContextKeyTest.java new file mode 100644 index 00000000000..8e65485f6c3 --- /dev/null +++ b/components/context/src/test/java/datadog/context/ContextKeyTest.java @@ -0,0 +1,34 @@ +package datadog.context; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.NullSource; +import org.junit.jupiter.params.provider.ValueSource; + +class ContextKeyTest { + + @ParameterizedTest + @NullSource + @ValueSource(strings = {"", "key"}) + void testConstructor(String name) { + ContextKey key = ContextKey.named(name); + assertNotNull(key); + assertEquals(name, key.toString()); + } + + @Test + void testKeyNameCollision() { + ContextKey key1 = ContextKey.named("same-name"); + ContextKey key2 = ContextKey.named("same-name"); + assertNotEquals(key1, key2); + String value = "value"; + Context context = Context.root().with(key1, value); + assertEquals(value, context.get(key1)); + assertNull(context.get(key2)); + } +} diff --git a/components/context/src/test/java/datadog/context/ContextManagerTest.java b/components/context/src/test/java/datadog/context/ContextManagerTest.java new file mode 100644 index 00000000000..41ca1e6859c --- /dev/null +++ b/components/context/src/test/java/datadog/context/ContextManagerTest.java @@ -0,0 +1,219 @@ +package datadog.context; + +import static datadog.context.Context.current; +import static datadog.context.Context.detach; +import static datadog.context.Context.root; +import static datadog.context.ContextTest.STRING_KEY; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; + +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.Phaser; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class ContextManagerTest { + + @BeforeEach + void init() { + // Ensure no current context prior starting test + assertEquals(root(), current()); + } + + @Test + void testContextAttachment() { + Context context1 = root().with(STRING_KEY, "value1"); + try (ContextScope scope1 = context1.attach()) { + // Test context1 is attached + assertEquals(context1, current()); + assertEquals(context1, scope1.context()); + Context context2 = context1.with(STRING_KEY, "value2"); + try (ContextScope scope2 = context2.attach()) { + // Test context2 is attached + assertEquals(context2, current()); + assertEquals(context2, scope2.context()); + // Can still access context1 from its scope + assertEquals(context1, scope1.context()); + } + // Test context1 is restored + assertEquals(context1, current()); + } + } + + @Test + void testContextSwapping() { + Context context1 = root().with(STRING_KEY, "value1"); + assertEquals(root(), current()); + assertEquals(root(), context1.swap()); + // Test context1 is attached + Context context2 = context1.with(STRING_KEY, "value2"); + assertEquals(context1, current()); + assertEquals(context1, context2.swap()); + // Test context2 is attached + assertEquals(context2, current()); + assertEquals(context2, detach()); + // Test we're now context-less + assertEquals(root(), current()); + } + + @Test + void testAttachSameContextMultipleTimes() { + Context context = root().with(STRING_KEY, "value1"); + try (ContextScope ignored1 = context.attach()) { + assertEquals(context, current()); + try (ContextScope ignored2 = context.attach()) { + try (ContextScope ignored3 = context.attach()) { + assertEquals(context, current()); + } + // Test closing a scope on the current context should not deactivate it if activated + // multiple times + assertEquals(context, current()); + } + } + // Test closing the same number of scope as activation should deactivate the context + assertEquals(root(), current()); + } + + @Test + void testOnlyCurrentScopeCanBeClosed() { + Context context1 = root().with(STRING_KEY, "value1"); + try (ContextScope scope1 = context1.attach()) { + Context context2 = context1.with(STRING_KEY, "value2"); + try (ContextScope ignored = context2.attach()) { + // Try closing the non-current scope + scope1.close(); + // Test context2 is still attached + assertEquals(context2, current()); + } + // Test context1 is restored + assertEquals(context1, current()); + } + } + + @Test + void testClosingMultipleTimes() { + Context context1 = root().with(STRING_KEY, "value1"); + try (ContextScope ignored = context1.attach()) { + Context context2 = context1.with(STRING_KEY, "value2"); + ContextScope scope = context2.attach(); + // Test current context + assertEquals(context2, current()); + // Test current context deactivation + scope.close(); + assertEquals(context1, current()); + // Test multiple context deactivations don’t change current context + scope.close(); + assertEquals(context1, current()); + } + } + + @Test + void testThreadIndependence() { + /* + * This test has 2 executors in addition to the main thread. + * They are synchronized using a Phaser, and arrived before each assert phase. + * If an assert fails in of one the executor, the executor is "deregister" to unblock the test, + * and the exception is restored at the end of the test using "Future.get()". + */ + ExecutorService executor = Executors.newFixedThreadPool(2); + Phaser phaser = new Phaser(3); + /* + * Create first executor. + */ + Future future1 = + executor.submit( + () -> { + try { + // Fist step: check empty context + phaser.arriveAndAwaitAdvance(); + assertEquals(root(), current()); + // Second step: set context on first executor + Context context1 = root().with(STRING_KEY, "executor1"); + try (ContextScope ignored1 = context1.attach()) { + phaser.arriveAndAwaitAdvance(); + assertEquals(context1, current()); + // Third step: set context on second executor + phaser.arriveAndAwaitAdvance(); + assertEquals(context1, current()); + // Fourth step: set child context on first executor + Context context11 = context1.with(STRING_KEY, "executor1.1"); + try (ContextScope ignored11 = context11.attach()) { + phaser.arriveAndAwaitAdvance(); + assertEquals(context11, current()); + } + } + } finally { + // Complete the execution + phaser.arriveAndDeregister(); + } + }); + /* + * Create second executor. + */ + Future future2 = + executor.submit( + () -> { + try { + // First step: check empty context + phaser.arriveAndAwaitAdvance(); + assertEquals(root(), current()); + // Second step: set context on first executor + phaser.arriveAndAwaitAdvance(); + assertEquals(root(), current()); + // Third step: set context on second executor + Context context2 = root().with(STRING_KEY, "executor2"); + try (ContextScope ignored2 = context2.attach()) { + phaser.arriveAndAwaitAdvance(); + assertEquals(context2, current()); + // Fourth step: set child context on first executor + phaser.arriveAndAwaitAdvance(); + assertEquals(context2, current()); + } + } finally { + // Complete the execution + phaser.arriveAndDeregister(); + } + }); + /* + * Run main thread. + */ + // First step: check empty context + phaser.arriveAndAwaitAdvance(); + assertEquals(root(), current()); + // Second step: set context on first executor + phaser.arriveAndAwaitAdvance(); + assertEquals(root(), current()); + // Third step: set context on second executor + phaser.arriveAndAwaitAdvance(); + assertEquals(root(), current()); + // Fourth step: set child context on first executor + phaser.arriveAndAwaitAdvance(); + assertEquals(root(), current()); + // Complete execution and wait for the others + phaser.arriveAndAwaitAdvance(); + executor.shutdown(); + // Check any test error in executors + assertDoesNotThrow(() -> future1.get()); + assertDoesNotThrow(() -> future2.get()); + } + + @Test + void testNonThreadInheritance() { + Context context = root().with(STRING_KEY, "value"); + try (ContextScope ignored = context.attach()) { + // Check new thread don't inherit from current context + ExecutorService executor = Executors.newSingleThreadExecutor(); + Future future = executor.submit(() -> assertEquals(root(), current())); + assertDoesNotThrow(() -> future.get()); + } + } + + @AfterEach + void tearDown() { + // Ensure no current context after ending test + assertEquals(root(), current()); + } +} diff --git a/components/context/src/test/java/datadog/context/ContextProviderForkedTest.java b/components/context/src/test/java/datadog/context/ContextProviderForkedTest.java new file mode 100644 index 00000000000..18059adf059 --- /dev/null +++ b/components/context/src/test/java/datadog/context/ContextProviderForkedTest.java @@ -0,0 +1,86 @@ +package datadog.context; + +import static datadog.context.Context.root; +import static datadog.context.ContextTest.STRING_KEY; +import static org.junit.jupiter.api.Assertions.assertEquals; + +import org.junit.jupiter.api.Test; + +class ContextProviderForkedTest { + + @Test + void testCustomBinder() { + // register a NOOP context binder + ContextBinder.register(new ContextBinder() { + @Override + public Context from(Object carrier) { + return root(); + } + + @Override + public void attachTo(Object carrier, Context context) { + // no-op + } + + @Override + public Context detachFrom(Object carrier) { + return root(); + } + }); + + Context context = root().with(STRING_KEY, "value"); + + // NOOP binder, context will always be root + Object carrier = new Object(); + context.attachTo(carrier); + assertEquals(root(), Context.from(carrier)); + } + + @Test + void testCustomManager() { + // register a NOOP context manager + ContextManager.register(new ContextManager() { + @Override + public Context root() { + return EmptyContext.INSTANCE; + } + + @Override + public Context current() { + return root(); + } + + @Override + public ContextScope attach(Context context) { + return new ContextScope() { + @Override + public Context context() { + return root(); + } + + @Override + public void close() { + // no-op + } + }; + } + + @Override + public Context swap(Context context) { + return root(); + } + + @Override + public Context detach() { + return root(); + } + }); + + Context context = root().with(STRING_KEY, "value"); + + // NOOP manager, context will always be root + try (ContextScope ignored = context.attach()) { + assertEquals(root(), Context.current()); + } + } +} diff --git a/components/context/src/test/java/datadog/context/ContextTest.java b/components/context/src/test/java/datadog/context/ContextTest.java new file mode 100644 index 00000000000..c1632163bb6 --- /dev/null +++ b/components/context/src/test/java/datadog/context/ContextTest.java @@ -0,0 +1,113 @@ +package datadog.context; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import javax.annotation.Nullable; +import org.junit.jupiter.api.Test; + +class ContextTest { + static final ContextKey STRING_KEY = ContextKey.named("string-key"); + static final ContextKey BOOLEAN_KEY = ContextKey.named("boolean-key"); + + // demonstrate how values can hide their context keys + static class ValueWithKey implements ImplicitContextKeyed { + static final ContextKey HIDDEN_KEY = ContextKey.named("hidden-key"); + + @Override + public Context storeInto(Context context) { + return context.with(HIDDEN_KEY, this); + } + + @Nullable + public static ValueWithKey from(Context context) { + return context.get(HIDDEN_KEY); + } + } + + @Test + void testEmpty() { + // Test empty is always the same + Context empty = Context.root(); + assertEquals(empty, Context.root()); + // Test empty is not mutated + String stringValue = "value"; + empty.with(STRING_KEY, stringValue); + assertEquals(empty, Context.root()); + } + + @Test + void testWith() { + Context empty = Context.root(); + // Test accessing non-set value + assertNull(empty.get(STRING_KEY)); + // Test retrieving value + String stringValue = "value"; + Context context1 = empty.with(STRING_KEY, stringValue); + assertEquals(stringValue, context1.get(STRING_KEY)); + // Test overriding value + String stringValue2 = "value2"; + Context context2 = context1.with(STRING_KEY, stringValue2); + assertEquals(stringValue2, context2.get(STRING_KEY)); + // Test clearing value + Context context3 = context2.with(STRING_KEY, null); + assertNull(context3.get(STRING_KEY)); + // Test null key handling + assertThrows(NullPointerException.class, () -> empty.with(null, "test")); + } + + @Test + void testGet() { + // Setup context + Context empty = Context.root(); + String value = "value"; + Context context = empty.with(STRING_KEY, value); + // Test null key handling + assertThrows(NullPointerException.class, () -> context.get(null)); + // Test unset key + assertNull(context.get(BOOLEAN_KEY)); + // Test set key + assertEquals(value, context.get(STRING_KEY)); + } + + @SuppressWarnings({ + "EqualsWithItself", + "SimplifiableAssertion", + "ConstantValue", + "EqualsBetweenInconvertibleTypes" + }) + @Test + void testEqualsAndHashCode() { + // Setup contexts + Context empty = Context.root(); + Context context1 = empty.with(STRING_KEY, "value"); + Context context2 = empty.with(STRING_KEY, "value "); + Context context3 = empty.with(STRING_KEY, "value ".trim()); + // Test equals on self + assertTrue(empty.equals(empty)); + assertTrue(context1.equals(context1)); + // Test equals on null + assertFalse(context1.equals(null)); + // Test equals on different object type + assertFalse(context1.equals("value")); + // Test equals on different contexts with the same values + assertTrue(context1.equals(context3)); + assertEquals(context1.hashCode(), context3.hashCode()); + // Test equals on different contexts + assertFalse(context1.equals(empty)); + assertFalse(context1.equals(context2)); + } + + @Test + void testImplicitKey() { + // Setup context + Context empty = Context.root(); + ValueWithKey valueWithKey = new ValueWithKey(); + Context context = empty.with(valueWithKey); + assertNull(ValueWithKey.from(empty)); + assertEquals(valueWithKey, ValueWithKey.from(context)); + } +} diff --git a/settings.gradle b/settings.gradle index 2d78784b023..67178ac923b 100644 --- a/settings.gradle +++ b/settings.gradle @@ -66,6 +66,7 @@ include ':dd-java-agent:agent-otel:otel-shim' include ':dd-java-agent:agent-otel:otel-tooling' include ':communication' +include ':components:context' include ':components:json' include ':telemetry' include ':remote-config:remote-config-api' @@ -512,3 +513,4 @@ include ':dd-java-agent:benchmark' include ':dd-java-agent:benchmark-integration' include ':dd-java-agent:benchmark-integration:jetty-perftest' include ':dd-java-agent:benchmark-integration:play-perftest' +