Implement ContextStorageOverride for opentelemetry context bridge (#1…

…1523)

contextstorage/build.gradle

@@ -0,0 +1,28 @@
+plugins {
+    id "java-library"
+    // until we are confident we like the name
+    //id "maven-publish"
+
+    id "ru.vyarus.animalsniffer"
+}
+
+description = 'gRPC: ContextStorageOverride'
+
+dependencies {
+    api project(':grpc-api')
+    implementation libraries.opentelemetry.api
+
+    testImplementation libraries.junit,
+            libraries.opentelemetry.sdk.testing,
+            libraries.assertj.core
+    testImplementation 'junit:junit:4.13.1'// opentelemetry.sdk.testing uses compileOnly for assertj
+
+    signature libraries.signature.java
+    signature libraries.signature.android
+}
+
+tasks.named("jar").configure {
+    manifest {
+        attributes('Automatic-Module-Name': 'io.grpc.override')
+    }
+}

contextstorage/src/main/java/io/grpc/override/ContextStorageOverride.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright 2024 The gRPC Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.grpc.override;
+
+import io.grpc.Context;
+
+/**
+ * Including this class in your dependencies will override the default gRPC context storage using
+ * reflection. It is a bridge between {@link io.grpc.Context} and
+ * {@link io.opentelemetry.context.Context}, i.e. propagating io.grpc.context.Context also
+ * propagates io.opentelemetry.context, and propagating io.opentelemetry.context will also propagate
+ * io.grpc.context.
+ */
+public final class ContextStorageOverride extends Context.Storage {
+
+  private final Context.Storage delegate = new OpenTelemetryContextStorage();
+
+  @Override
+  public Context doAttach(Context toAttach) {
+    return delegate.doAttach(toAttach);
+  }
+
+  @Override
+  public void detach(Context toDetach, Context toRestore) {
+    delegate.detach(toDetach, toRestore);
+  }
+
+  @Override
+  public Context current() {
+    return delegate.current();
+  }
+}

contextstorage/src/main/java/io/grpc/override/OpenTelemetryContextStorage.java

@@ -0,0 +1,72 @@
+/*
+ * Copyright 2024 The gRPC Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.grpc.override;
+
+import io.grpc.Context;
+import io.opentelemetry.context.ContextKey;
+import io.opentelemetry.context.Scope;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+/**
+ * A Context.Storage implementation that attaches io.grpc.context to OpenTelemetry's context and
+ * io.opentelemetry.context is also saved in the io.grpc.context.
+ * Bridge between {@link io.grpc.Context} and {@link io.opentelemetry.context.Context}.
+ */
+final class OpenTelemetryContextStorage extends Context.Storage {
+  private static final Logger logger = Logger.getLogger(
+      OpenTelemetryContextStorage.class.getName());
+
+  private static final io.grpc.Context.Key<io.opentelemetry.context.Context> OTEL_CONTEXT_OVER_GRPC
+      = io.grpc.Context.key("otel-context-over-grpc");
+  private static final Context.Key<Scope> OTEL_SCOPE = Context.key("otel-scope");
+  private static final ContextKey<io.grpc.Context> GRPC_CONTEXT_OVER_OTEL =
+      ContextKey.named("grpc-context-over-otel");
+
+  @Override
+  @SuppressWarnings("MustBeClosedChecker")
+  public Context doAttach(Context toAttach) {
+    io.grpc.Context previous = current();
+    io.opentelemetry.context.Context otelContext = OTEL_CONTEXT_OVER_GRPC.get(toAttach);
+    if (otelContext == null) {
+      otelContext = io.opentelemetry.context.Context.current();
+    }
+    Scope scope = otelContext.with(GRPC_CONTEXT_OVER_OTEL, toAttach).makeCurrent();
+    return previous.withValue(OTEL_SCOPE, scope);
+  }
+
+  @Override
+  public void detach(Context toDetach, Context toRestore) {
+    Scope scope = OTEL_SCOPE.get(toRestore);
+    if (scope == null) {
+      logger.log(
+          Level.SEVERE, "Detaching context which was not attached.");
+    } else {
+      scope.close();
+    }
+  }
+
+  @Override
+  public Context current() {
+    io.opentelemetry.context.Context otelCurrent = io.opentelemetry.context.Context.current();
+    io.grpc.Context grpcCurrent = otelCurrent.get(GRPC_CONTEXT_OVER_OTEL);
+    if (grpcCurrent == null) {
+      grpcCurrent = Context.ROOT;
+    }
+    return grpcCurrent.withValue(OTEL_CONTEXT_OVER_GRPC, otelCurrent);
+  }
+}

contextstorage/src/test/java/io/grpc/override/OpenTelemetryContextStorageTest.java

@@ -0,0 +1,144 @@
+/*
+ * Copyright 2024 The gRPC Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.grpc.override;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+
+import com.google.common.util.concurrent.SettableFuture;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.ContextKey;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.sdk.testing.junit4.OpenTelemetryRule;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+import org.junit.Assert;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class OpenTelemetryContextStorageTest {
+  @Rule
+  public final OpenTelemetryRule openTelemetryRule = OpenTelemetryRule.create();
+  private Tracer tracerRule = openTelemetryRule.getOpenTelemetry().getTracer(
+      "context-storage-test");
+  private final io.grpc.Context.Key<String> username = io.grpc.Context.key("username");
+  private final ContextKey<String> password = ContextKey.named("password");
+
+  @Test
+  public void grpcContextPropagation() throws Exception {
+    final Span parentSpan = tracerRule.spanBuilder("test-context").startSpan();
+    final SettableFuture<Span> spanPropagated = SettableFuture.create();
+    final SettableFuture<String> grpcContextPropagated = SettableFuture.create();
+    final SettableFuture<Span> spanDetached = SettableFuture.create();
+    final SettableFuture<String> grpcContextDetached = SettableFuture.create();
+
+    io.grpc.Context grpcContext;
+    try (Scope scope = Context.current().with(parentSpan).makeCurrent()) {
+      grpcContext = io.grpc.Context.current().withValue(username, "jeff");
+    }
+    new Thread(new Runnable() {
+      @Override
+      public void run() {
+        io.grpc.Context previous = grpcContext.attach();
+        try {
+          grpcContextPropagated.set(username.get(io.grpc.Context.current()));
+          spanPropagated.set(Span.fromContext(io.opentelemetry.context.Context.current()));
+        } finally {
+          grpcContext.detach(previous);
+          spanDetached.set(Span.fromContext(io.opentelemetry.context.Context.current()));
+          grpcContextDetached.set(username.get(io.grpc.Context.current()));
+        }
+      }
+    }).start();
+    Assert.assertEquals(spanPropagated.get(5, TimeUnit.SECONDS), parentSpan);
+    Assert.assertEquals(grpcContextPropagated.get(5, TimeUnit.SECONDS), "jeff");
+    Assert.assertEquals(spanDetached.get(5, TimeUnit.SECONDS), Span.getInvalid());
+    Assert.assertNull(grpcContextDetached.get(5, TimeUnit.SECONDS));
+  }
+
+  @Test
+  public void otelContextPropagation() throws Exception {
+    final SettableFuture<String> grpcPropagated = SettableFuture.create();
+    final AtomicReference<String> otelPropagation = new AtomicReference<>();
+
+    io.grpc.Context grpcContext = io.grpc.Context.current().withValue(username, "jeff");
+    io.grpc.Context previous = grpcContext.attach();
+    Context original = Context.current().with(password, "valentine");
+    try {
+      new Thread(
+          () -> {
+            try (Scope scope = original.makeCurrent()) {
+              otelPropagation.set(Context.current().get(password));
+              grpcPropagated.set(username.get(io.grpc.Context.current()));
+            }
+          }
+      ).start();
+    } finally {
+      grpcContext.detach(previous);
+    }
+    Assert.assertEquals(grpcPropagated.get(5, TimeUnit.SECONDS), "jeff");
+    Assert.assertEquals(otelPropagation.get(), "valentine");
+  }
+
+  @Test
+  public void grpcOtelMix() {
+    io.grpc.Context grpcContext = io.grpc.Context.current().withValue(username, "jeff");
+    Context otelContext = Context.current().with(password, "valentine");
+    Assert.assertNull(username.get(io.grpc.Context.current()));
+    Assert.assertNull(Context.current().get(password));
+    io.grpc.Context previous = grpcContext.attach();
+    try {
+      assertEquals(username.get(io.grpc.Context.current()), "jeff");
+      try (Scope scope = otelContext.makeCurrent()) {
+        Assert.assertEquals(Context.current().get(password), "valentine");
+        assertNull(username.get(io.grpc.Context.current()));
+
+        io.grpc.Context grpcContext2 = io.grpc.Context.current().withValue(username, "frank");
+        io.grpc.Context previous2 = grpcContext2.attach();
+        try {
+          assertEquals(username.get(io.grpc.Context.current()), "frank");
+          Assert.assertEquals(Context.current().get(password), "valentine");
+        } finally {
+          grpcContext2.detach(previous2);
+        }
+        assertNull(username.get(io.grpc.Context.current()));
+        Assert.assertEquals(Context.current().get(password), "valentine");
+      }
+    } finally {
+      grpcContext.detach(previous);
+    }
+    Assert.assertNull(username.get(io.grpc.Context.current()));
+    Assert.assertNull(Context.current().get(password));
+  }
+
+  @Test
+  public void grpcContextDetachError() {
+    io.grpc.Context grpcContext = io.grpc.Context.current().withValue(username, "jeff");
+    io.grpc.Context previous = grpcContext.attach();
+    try {
+      previous.detach(grpcContext);
+      assertEquals(username.get(io.grpc.Context.current()), "jeff");
+    } finally {
+      grpcContext.detach(previous);
+    }
+  }
+}

settings.gradle

@@ -77,6 +77,7 @@ include ":grpc-istio-interop-testing"
 include ":grpc-inprocess"
 include ":grpc-util"
 include ":grpc-opentelemetry"
+include ":grpc-opentelemetry-context-storage-override"
 
 project(':grpc-api').projectDir = "$rootDir/api" as File
 project(':grpc-core').projectDir = "$rootDir/core" as File
@@ -113,6 +114,7 @@ project(':grpc-istio-interop-testing').projectDir = "$rootDir/istio-interop-test
 project(':grpc-inprocess').projectDir = "$rootDir/inprocess" as File
 project(':grpc-util').projectDir = "$rootDir/util" as File
 project(':grpc-opentelemetry').projectDir = "$rootDir/opentelemetry" as File
+project(':grpc-opentelemetry-context-storage-override').projectDir = "$rootDir/contextstorage" as File
 
 if (settings.hasProperty('skipCodegen') && skipCodegen.toBoolean()) {
     println '*** Skipping the build of codegen and compilation of proto files because skipCodegen=true'