Add a way to enforce default context storage implementation (#2039)

* Add a way to enforce default context storage implementation

* Fix review

* Add testcase

* Use junit_pioneer

* Change the methods to be package-private

* Add some comments

* Remove DefaultContext.threadLocalStorage

* Fix review

context/build.gradle

@@ -28,6 +28,7 @@ dependencies {
     otelInBraveTestImplementation "io.zipkin.brave:brave:5.12.6"
 
     testImplementation libraries.awaitility
+    testImplementation libraries.junit_pioneer
 
     signature libraries.android_signature
 }

context/src/main/java/io/opentelemetry/context/DefaultContext.java

@@ -28,15 +28,6 @@ final class DefaultContext implements Context {
 
   private static final Context ROOT = new DefaultContext();
 
-  /**
-   * Returns the default {@link ContextStorage} used to attach {@link Context}s to scopes of
-   * execution. Should only be used when defining your own {@link ContextStorage} in case you want
-   * to delegate functionality to the default implementation.
-   */
-  static ContextStorage threadLocalStorage() {
-    return ThreadLocalContextStorage.INSTANCE;
-  }
-
   // Used by auto-instrumentation agent. Check with auto-instrumentation before making changes to
   // this method.
   //

context/src/main/java/io/opentelemetry/context/LazyStorage.java

@@ -65,6 +65,7 @@ final class LazyStorage {
 
   private static final String CONTEXT_STORAGE_PROVIDER_PROPERTY =
       "io.opentelemetry.context.contextStorageProvider";
+  private static final String ENFORCE_DEFAULT_STORAGE_VALUE = "default";
 
   private static final Logger logger = Logger.getLogger(LazyStorage.class.getName());
 
@@ -82,7 +83,12 @@ final class LazyStorage {
   }
 
   static ContextStorage createStorage(AtomicReference<Throwable> deferredStorageFailure) {
-    String providerClassName = System.getProperty(CONTEXT_STORAGE_PROVIDER_PROPERTY, "");
+    // Get the specified SPI implementation first here
+    final String providerClassName = System.getProperty(CONTEXT_STORAGE_PROVIDER_PROPERTY, "");
+    // Allow user to enforce default ThreadLocalContextStorage
+    if (ENFORCE_DEFAULT_STORAGE_VALUE.equals(providerClassName)) {
+      return ContextStorage.defaultStorage();
+    }
 
     List<ContextStorageProvider> providers = new ArrayList<>();
     for (ContextStorageProvider provider : ServiceLoader.load(ContextStorageProvider.class)) {
@@ -112,7 +118,7 @@ final class LazyStorage {
                   + "qualified class name of the provider to use. Falling back to default "
                   + "ContextStorage. Found providers: "
                   + providers));
-      return DefaultContext.threadLocalStorage();
+      return ContextStorage.defaultStorage();
     }
 
     for (ContextStorageProvider provider : providers) {
@@ -128,7 +134,7 @@ final class LazyStorage {
                 + providerClassName
                 + " but found providers: "
                 + providers));
-    return DefaultContext.threadLocalStorage();
+    return ContextStorage.defaultStorage();
   }
 
   private LazyStorage() {}

context/src/test/java/io/opentelemetry/context/LazyStorageTest.java

@@ -13,37 +13,36 @@ import java.io.IOException;
 import java.io.Writer;
 import java.net.URL;
 import java.util.concurrent.atomic.AtomicReference;
-import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
+import org.junitpioneer.jupiter.ClearSystemProperty;
+import org.junitpioneer.jupiter.SetSystemProperty;
 
 class LazyStorageTest {
 
   private static final String CONTEXT_STORAGE_PROVIDER_PROPERTY =
       "io.opentelemetry.context.contextStorageProvider";
+  private static final String MOCK_CONTEXT_STORAGE_PROVIDER =
+      "io.opentelemetry.context.LazyStorageTest$MockContextStorageProvider";
   private static final AtomicReference<Throwable> DEFERRED_STORAGE_FAILURE =
       new AtomicReference<>();
 
-  @AfterEach
-  public void after() {
-    System.clearProperty(CONTEXT_STORAGE_PROVIDER_PROPERTY);
-  }
-
   @Test
-  public void empty_provides() {
+  @ClearSystemProperty(key = CONTEXT_STORAGE_PROVIDER_PROPERTY)
+  void empty_providers() {
     assertThat(LazyStorage.createStorage(DEFERRED_STORAGE_FAILURE))
-        .isEqualTo(DefaultContext.threadLocalStorage());
+        .isEqualTo(ContextStorage.defaultStorage());
   }
 
   @Test
-  public void set_storage_provider_property_and_empty_provides() {
+  @SetSystemProperty(key = CONTEXT_STORAGE_PROVIDER_PROPERTY, value = MOCK_CONTEXT_STORAGE_PROVIDER)
-    System.setProperty(
+  void set_storage_provider_property_and_empty_providers() {
-        CONTEXT_STORAGE_PROVIDER_PROPERTY, MockContextStorageProvider.class.getName());
     assertThat(LazyStorage.createStorage(DEFERRED_STORAGE_FAILURE))
-        .isEqualTo(DefaultContext.threadLocalStorage());
+        .isEqualTo(ContextStorage.defaultStorage());
   }
 
   @Test
-  public void unset_storage_provider_property_and_one_provides() throws Exception {
+  @ClearSystemProperty(key = CONTEXT_STORAGE_PROVIDER_PROPERTY)
+  void unset_storage_provider_property_and_one_providers() throws Exception {
     File serviceFile = createContextStorageProvider();
     try {
       assertThat(LazyStorage.createStorage(DEFERRED_STORAGE_FAILURE)).isEqualTo(mockContextStorage);
@@ -53,21 +52,22 @@ class LazyStorageTest {
   }
 
   @Test
-  public void set_storage_provider_property_not_matches_one_provides() throws Exception {
+  @SetSystemProperty(
-    System.setProperty(CONTEXT_STORAGE_PROVIDER_PROPERTY, "not.match.provider.class.name");
+      key = CONTEXT_STORAGE_PROVIDER_PROPERTY,
+      value = "not.match.provider.class.name")
+  void set_storage_provider_property_not_matches_one_providers() throws Exception {
     File serviceFile = createContextStorageProvider();
     try {
       assertThat(LazyStorage.createStorage(DEFERRED_STORAGE_FAILURE))
-          .isEqualTo(DefaultContext.threadLocalStorage());
+          .isEqualTo(ContextStorage.defaultStorage());
     } finally {
       serviceFile.delete();
     }
   }
 
   @Test
-  public void set_storage_provider_property_matches_one_provides() throws Exception {
+  @SetSystemProperty(key = CONTEXT_STORAGE_PROVIDER_PROPERTY, value = MOCK_CONTEXT_STORAGE_PROVIDER)
-    System.setProperty(
+  void set_storage_provider_property_matches_one_providers() throws Exception {
-        CONTEXT_STORAGE_PROVIDER_PROPERTY, MockContextStorageProvider.class.getName());
     File serviceFile = createContextStorageProvider();
     try {
       assertThat(LazyStorage.createStorage(DEFERRED_STORAGE_FAILURE)).isEqualTo(mockContextStorage);
@@ -76,6 +76,25 @@ class LazyStorageTest {
     }
   }
 
+  @Test
+  @SetSystemProperty(key = CONTEXT_STORAGE_PROVIDER_PROPERTY, value = "default")
+  void enforce_default_and_empty_providers() {
+    assertThat(LazyStorage.createStorage(DEFERRED_STORAGE_FAILURE))
+        .isEqualTo(ContextStorage.defaultStorage());
+  }
+
+  @Test
+  @SetSystemProperty(key = CONTEXT_STORAGE_PROVIDER_PROPERTY, value = "default")
+  void enforce_default_and_one_providers() throws IOException {
+    File serviceFile = createContextStorageProvider();
+    try {
+      assertThat(LazyStorage.createStorage(DEFERRED_STORAGE_FAILURE))
+          .isEqualTo(ContextStorage.defaultStorage());
+    } finally {
+      serviceFile.delete();
+    }
+  }
+
   private static File createContextStorageProvider() throws IOException {
     URL location =
         MockContextStorageProvider.class.getProtectionDomain().getCodeSource().getLocation();