diff --git a/context/build.gradle b/context/build.gradle index c88f2fd0ef5..57ff7991f13 100644 --- a/context/build.gradle +++ b/context/build.gradle @@ -17,6 +17,8 @@ testSets { braveInOtelTest otelInBraveTest otelAsBraveTest + + storageWrappersTest } dependencies { diff --git a/context/src/main/java/io/opentelemetry/context/ContextStorage.java b/context/src/main/java/io/opentelemetry/context/ContextStorage.java index cd13e7e2ea2..f9e94d2d768 100644 --- a/context/src/main/java/io/opentelemetry/context/ContextStorage.java +++ b/context/src/main/java/io/opentelemetry/context/ContextStorage.java @@ -22,6 +22,7 @@ package io.opentelemetry.context; +import java.util.function.Function; import javax.annotation.Nullable; /** @@ -78,6 +79,16 @@ static ContextStorage defaultStorage() { return ThreadLocalContextStorage.INSTANCE; } + /** + * Adds the {@code wrapper}, which will be executed with the {@link ContextStorage} is first used, + * i.e., by calling {@link Context#makeCurrent()}. This must be called as early in your + * application as possible to have effect, often as part of a static initialization block in your + * main class. + */ + static void addWrapper(Function wrapper) { + ContextStorageWrappers.addWrapper(wrapper); + } + /** * Sets the specified {@link Context} as the current {@link Context} and returns a {@link Scope} * representing the scope of execution. {@link Scope#close()} must be called when the current diff --git a/context/src/main/java/io/opentelemetry/context/ContextStorageWrappers.java b/context/src/main/java/io/opentelemetry/context/ContextStorageWrappers.java new file mode 100644 index 00000000000..9f8eb3276a1 --- /dev/null +++ b/context/src/main/java/io/opentelemetry/context/ContextStorageWrappers.java @@ -0,0 +1,55 @@ +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + */ + +package io.opentelemetry.context; + +import java.util.ArrayList; +import java.util.List; +import java.util.function.Function; +import java.util.logging.Level; +import java.util.logging.Logger; + +/** + * Holder of functions to wrap the used {@link ContextStorage}. Separate class from {@link + * LazyStorage} to allow registering wrappers before initializing storage. + */ +final class ContextStorageWrappers { + + private static final Logger log = Logger.getLogger(ContextStorageWrappers.class.getName()); + + private static boolean storageInitialized; + + private static final List> wrappers = + new ArrayList<>(); + + private static final Object mutex = new Object(); + + static void addWrapper(Function wrapper) { + synchronized (mutex) { + if (storageInitialized) { + log.log( + Level.FINE, + "ContextStorage has already been initialized, ignoring call to add wrapper.", + new Throwable()); + return; + } + wrappers.add(wrapper); + } + } + + static List> getWrappers() { + synchronized (mutex) { + return wrappers; + } + } + + static void setStorageInitialized() { + synchronized (mutex) { + storageInitialized = true; + } + } + + private ContextStorageWrappers() {} +} diff --git a/context/src/main/java/io/opentelemetry/context/LazyStorage.java b/context/src/main/java/io/opentelemetry/context/LazyStorage.java index 4b312df5c20..2aeef651e50 100644 --- a/context/src/main/java/io/opentelemetry/context/LazyStorage.java +++ b/context/src/main/java/io/opentelemetry/context/LazyStorage.java @@ -41,6 +41,7 @@ import java.util.List; import java.util.ServiceLoader; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; import java.util.logging.Level; import java.util.logging.Logger; @@ -72,7 +73,13 @@ static ContextStorage get() { static { AtomicReference deferredStorageFailure = new AtomicReference<>(); - storage = createStorage(deferredStorageFailure); + ContextStorage created = createStorage(deferredStorageFailure); + for (Function wrapper : + ContextStorageWrappers.getWrappers()) { + created = wrapper.apply(created); + } + storage = created; + ContextStorageWrappers.setStorageInitialized(); Throwable failure = deferredStorageFailure.get(); // Logging must happen after storage has been set, as loggers may use Context. if (failure != null) { diff --git a/context/src/storageWrappersTest/java/io/opentelemetry/context/StorageWrappersTest.java b/context/src/storageWrappersTest/java/io/opentelemetry/context/StorageWrappersTest.java new file mode 100644 index 00000000000..7be7b99f8b9 --- /dev/null +++ b/context/src/storageWrappersTest/java/io/opentelemetry/context/StorageWrappersTest.java @@ -0,0 +1,63 @@ +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + */ + +package io.opentelemetry.context; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Function; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.RepeatedTest; + +class StorageWrappersTest { + + private static final ContextKey ANIMAL = ContextKey.named("key"); + + private static final AtomicInteger scopeOpenedCount = new AtomicInteger(); + private static final AtomicInteger scopeClosedCount = new AtomicInteger(); + + @SuppressWarnings("UnnecessaryLambda") + private static final Function wrapper = + delegate -> + new ContextStorage() { + @Override + public Scope attach(Context toAttach) { + Scope scope = delegate.attach(toAttach); + scopeOpenedCount.incrementAndGet(); + return () -> { + scope.close(); + scopeClosedCount.incrementAndGet(); + }; + } + + @Override + public Context current() { + return delegate.current(); + } + }; + + @BeforeEach + void resetCounts() { + scopeOpenedCount.set(0); + scopeClosedCount.set(0); + } + + // Run twice to ensure second wrapping has no effect. + @RepeatedTest(2) + void wrapAndInitialize() { + ContextStorage.addWrapper(wrapper); + + assertThat(scopeOpenedCount).hasValue(0); + assertThat(scopeClosedCount).hasValue(0); + + try (Scope ignored = Context.current().with(ANIMAL, "koala").makeCurrent()) { + assertThat(Context.current().get(ANIMAL)).isEqualTo("koala"); + } + + assertThat(scopeOpenedCount).hasValue(1); + assertThat(scopeClosedCount).hasValue(1); + } +}