diff --git a/contextstorage/build.gradle b/contextstorage/build.gradle new file mode 100644 index 00000000000..10c7f3a3a86 --- /dev/null +++ b/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') + } +} diff --git a/contextstorage/src/main/java/io/grpc/override/ContextStorageOverride.java b/contextstorage/src/main/java/io/grpc/override/ContextStorageOverride.java new file mode 100644 index 00000000000..41b24765de0 --- /dev/null +++ b/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(); + } +} diff --git a/contextstorage/src/main/java/io/grpc/override/OpenTelemetryContextStorage.java b/contextstorage/src/main/java/io/grpc/override/OpenTelemetryContextStorage.java new file mode 100644 index 00000000000..01356e9f406 --- /dev/null +++ b/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 OTEL_CONTEXT_OVER_GRPC + = io.grpc.Context.key("otel-context-over-grpc"); + private static final Context.Key OTEL_SCOPE = Context.key("otel-scope"); + private static final ContextKey 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); + } +} diff --git a/contextstorage/src/test/java/io/grpc/override/OpenTelemetryContextStorageTest.java b/contextstorage/src/test/java/io/grpc/override/OpenTelemetryContextStorageTest.java new file mode 100644 index 00000000000..3c628964342 --- /dev/null +++ b/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 username = io.grpc.Context.key("username"); + private final ContextKey password = ContextKey.named("password"); + + @Test + public void grpcContextPropagation() throws Exception { + final Span parentSpan = tracerRule.spanBuilder("test-context").startSpan(); + final SettableFuture spanPropagated = SettableFuture.create(); + final SettableFuture grpcContextPropagated = SettableFuture.create(); + final SettableFuture spanDetached = SettableFuture.create(); + final SettableFuture 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 grpcPropagated = SettableFuture.create(); + final AtomicReference 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); + } + } +} diff --git a/settings.gradle b/settings.gradle index 61972e30b6c..bc0471734e2 100644 --- a/settings.gradle +++ b/settings.gradle @@ -76,6 +76,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 @@ -111,6 +112,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'