Allow root Context to be overridden (#3230)

* Allow root Context to be overridden

Adds ContextStorage.root() to provide a way to customize the
implementation of the root Context when you don't want the
root to be an ArrayBasedContext.

* Update context/src/main/java/io/opentelemetry/context/ContextStorage.java

Co-authored-by: Anuraag Agrawal <aanuraag@amazon.co.jp>
Co-authored-by: Anuraag Agrawal <anuraaga@gmail.com>

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

@@ -96,7 +96,7 @@ public interface Context {
    * is only a workaround hiding an underlying context propagation issue.
    */
   static Context root() {
-    return ArrayBasedContext.root();
+    return ContextStorage.get().root();
   }
 
   /**

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

@@ -102,4 +102,14 @@ public interface ContextStorage {
    */
   @Nullable
   Context current();
+
+  /**
+   * Returns the root {@link Context} which all other {@link Context} are derived from.
+   *
+   * <p>The default implementation returns the root {@code ArrayBasedContext}, but subclasses can
+   * override this method to return a root instance of a different {@link Context} implementation.
+   */
+  default Context root() {
+    return ArrayBasedContext.root();
+  }
 }

context/src/otelAsBraveTest/java/io/opentelemetry/context/BraveContextStorageProvider.java

@@ -43,7 +43,12 @@ public class BraveContextStorageProvider implements ContextStorageProvider {
       if (current != null) {
         return new BraveContextWrapper(current);
       }
-      return new BraveContextWrapper(TraceContext.newBuilder().traceId(1).spanId(1).build());
+      return root();
+    }
+
+    @Override
+    public Context root() {
+      return BraveContextWrapper.ROOT;
     }
   }
 
@@ -86,6 +91,9 @@ public class BraveContextStorageProvider implements ContextStorageProvider {
 
   private static class BraveContextWrapper implements Context {
 
+    private static final BraveContextWrapper ROOT =
+        new BraveContextWrapper(TraceContext.newBuilder().traceId(1).spanId(1).build());
+
     private final TraceContext braveContext;
 
     private BraveContextWrapper(TraceContext braveContext) {

context/src/otelInBraveTest/java/io/opentelemetry/context/BraveContextStorageProvider.java

@@ -50,10 +50,18 @@ public class BraveContextStorageProvider implements ContextStorageProvider {
     public Context current() {
       return new BraveContextWrapper(Tracing.current().currentTraceContext().get());
     }
+
+    @Override
+    public Context root() {
+      return BraveContextWrapper.ROOT;
+    }
   }
 
   // Need to wrap the Context because brave findExtra searches for perfect match of the class.
   static final class BraveContextWrapper implements Context {
+
+    static final BraveContextWrapper ROOT = new BraveContextWrapper(null, ArrayBasedContext.root());
+
     @Nullable private final TraceContext baseBraveContext;
     private final Context delegate;
 
@@ -93,7 +101,7 @@ public class BraveContextStorageProvider implements ContextStorageProvider {
 
   private static Context fromBraveContext(@Nullable TraceContext braveContext) {
     if (braveContext == null) {
-      return Context.root();
+      return BraveContextWrapper.ROOT.delegate;
     }
     List<Object> extra = braveContext.extra();
     for (int i = extra.size() - 1; i >= 0; i--) {

context/src/otelInGrpcTest/java/io/opentelemetry/context/GrpcContextStorageProvider.java

@@ -7,7 +7,7 @@ package io.opentelemetry.context;
 
 public class GrpcContextStorageProvider implements ContextStorageProvider {
   private static final io.grpc.Context.Key<Context> OTEL_CONTEXT =
-      io.grpc.Context.keyWithDefault("otel-context", Context.root());
+      io.grpc.Context.keyWithDefault("otel-context", GrpcContextWrapper.ROOT.context);
 
   @Override
   public ContextStorage get() {
@@ -44,9 +44,18 @@ public class GrpcContextStorageProvider implements ContextStorageProvider {
       io.grpc.Context grpcContext = io.grpc.Context.current();
       return GrpcContextWrapper.wrapperFromGrpc(grpcContext);
     }
+
+    @Override
+    public Context root() {
+      return GrpcContextWrapper.ROOT;
+    }
   }
 
   private static class GrpcContextWrapper implements Context {
+
+    private static final GrpcContextWrapper ROOT =
+        new GrpcContextWrapper(io.grpc.Context.ROOT, ArrayBasedContext.root());
+
     // If otel context changes the grpc Context may be out of sync.
     // There are 2 options here: 1. always update the grpc Context, 2. update only when needed.
     // Currently the second one is implemented.

docs/apidiffs/current_vs_latest/opentelemetry-context.txt

@@ -1,2 +1,4 @@
 Comparing source compatibility of  against 
-No changes.
+***! MODIFIED INTERFACE: PUBLIC ABSTRACT io.opentelemetry.context.ContextStorage  (not serializable)
+	===  CLASS FILE FORMAT VERSION: 52.0 <- 52.0
+	+++! NEW METHOD: PUBLIC(+) io.opentelemetry.context.Context root()