Allow registering ContextStorage wrappers with code to avoid requirin…

context/build.gradle

@@ -17,6 +17,8 @@ testSets {
     braveInOtelTest
     otelInBraveTest
     otelAsBraveTest
+
+    storageWrappersTest
 }
 
 dependencies {

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<? super ContextStorage, ? extends ContextStorage> 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

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<Function<? super ContextStorage, ? extends ContextStorage>> wrappers =
+      new ArrayList<>();
+
+  private static final Object mutex = new Object();
+
+  static void addWrapper(Function<? super ContextStorage, ? extends ContextStorage> 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<Function<? super ContextStorage, ? extends ContextStorage>> getWrappers() {
+    synchronized (mutex) {
+      return wrappers;
+    }
+  }
+
+  static void setStorageInitialized() {
+    synchronized (mutex) {
+      storageInitialized = true;
+    }
+  }
+
+  private ContextStorageWrappers() {}
+}

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<Throwable> deferredStorageFailure = new AtomicReference<>();
-    storage = createStorage(deferredStorageFailure);
+    ContextStorage created = createStorage(deferredStorageFailure);
+    for (Function<? super ContextStorage, ? extends ContextStorage> 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) {

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<String> ANIMAL = ContextKey.named("key");
+
+  private static final AtomicInteger scopeOpenedCount = new AtomicInteger();
+  private static final AtomicInteger scopeClosedCount = new AtomicInteger();
+
+  @SuppressWarnings("UnnecessaryLambda")
+  private static final Function<ContextStorage, ContextStorage> 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);
+  }
+}