Implement ContextStorageOverride for opentelemetry context bridge (#1… · grpc/grpc-java@782a44a
1+/*
2+ * Copyright 2024 The gRPC Authors
3+ *
4+ * Licensed under the Apache License, Version 2.0 (the "License");
5+ * you may not use this file except in compliance with the License.
6+ * You may obtain a copy of the License at
7+ *
8+ * http://www.apache.org/licenses/LICENSE-2.0
9+ *
10+ * Unless required by applicable law or agreed to in writing, software
11+ * distributed under the License is distributed on an "AS IS" BASIS,
12+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+ * See the License for the specific language governing permissions and
14+ * limitations under the License.
15+ */
16+17+package io.grpc.override;
18+19+import static org.junit.Assert.assertEquals;
20+import static org.junit.Assert.assertNull;
21+22+import com.google.common.util.concurrent.SettableFuture;
23+import io.opentelemetry.api.trace.Span;
24+import io.opentelemetry.api.trace.Tracer;
25+import io.opentelemetry.context.Context;
26+import io.opentelemetry.context.ContextKey;
27+import io.opentelemetry.context.Scope;
28+import io.opentelemetry.sdk.testing.junit4.OpenTelemetryRule;
29+import java.util.concurrent.TimeUnit;
30+import java.util.concurrent.atomic.AtomicReference;
31+import org.junit.Assert;
32+import org.junit.Rule;
33+import org.junit.Test;
34+import org.junit.runner.RunWith;
35+import org.junit.runners.JUnit4;
36+37+@RunWith(JUnit4.class)
38+public class OpenTelemetryContextStorageTest {
39+@Rule
40+public final OpenTelemetryRule openTelemetryRule = OpenTelemetryRule.create();
41+private Tracer tracerRule = openTelemetryRule.getOpenTelemetry().getTracer(
42+"context-storage-test");
43+private final io.grpc.Context.Key<String> username = io.grpc.Context.key("username");
44+private final ContextKey<String> password = ContextKey.named("password");
45+46+@Test
47+public void grpcContextPropagation() throws Exception {
48+final Span parentSpan = tracerRule.spanBuilder("test-context").startSpan();
49+final SettableFuture<Span> spanPropagated = SettableFuture.create();
50+final SettableFuture<String> grpcContextPropagated = SettableFuture.create();
51+final SettableFuture<Span> spanDetached = SettableFuture.create();
52+final SettableFuture<String> grpcContextDetached = SettableFuture.create();
53+54+io.grpc.Context grpcContext;
55+try (Scope scope = Context.current().with(parentSpan).makeCurrent()) {
56+grpcContext = io.grpc.Context.current().withValue(username, "jeff");
57+ }
58+new Thread(new Runnable() {
59+@Override
60+public void run() {
61+io.grpc.Context previous = grpcContext.attach();
62+try {
63+grpcContextPropagated.set(username.get(io.grpc.Context.current()));
64+spanPropagated.set(Span.fromContext(io.opentelemetry.context.Context.current()));
65+ } finally {
66+grpcContext.detach(previous);
67+spanDetached.set(Span.fromContext(io.opentelemetry.context.Context.current()));
68+grpcContextDetached.set(username.get(io.grpc.Context.current()));
69+ }
70+ }
71+ }).start();
72+Assert.assertEquals(spanPropagated.get(5, TimeUnit.SECONDS), parentSpan);
73+Assert.assertEquals(grpcContextPropagated.get(5, TimeUnit.SECONDS), "jeff");
74+Assert.assertEquals(spanDetached.get(5, TimeUnit.SECONDS), Span.getInvalid());
75+Assert.assertNull(grpcContextDetached.get(5, TimeUnit.SECONDS));
76+ }
77+78+@Test
79+public void otelContextPropagation() throws Exception {
80+final SettableFuture<String> grpcPropagated = SettableFuture.create();
81+final AtomicReference<String> otelPropagation = new AtomicReference<>();
82+83+io.grpc.Context grpcContext = io.grpc.Context.current().withValue(username, "jeff");
84+io.grpc.Context previous = grpcContext.attach();
85+Context original = Context.current().with(password, "valentine");
86+try {
87+new Thread(
88+ () -> {
89+try (Scope scope = original.makeCurrent()) {
90+otelPropagation.set(Context.current().get(password));
91+grpcPropagated.set(username.get(io.grpc.Context.current()));
92+ }
93+ }
94+ ).start();
95+ } finally {
96+grpcContext.detach(previous);
97+ }
98+Assert.assertEquals(grpcPropagated.get(5, TimeUnit.SECONDS), "jeff");
99+Assert.assertEquals(otelPropagation.get(), "valentine");
100+ }
101+102+@Test
103+public void grpcOtelMix() {
104+io.grpc.Context grpcContext = io.grpc.Context.current().withValue(username, "jeff");
105+Context otelContext = Context.current().with(password, "valentine");
106+Assert.assertNull(username.get(io.grpc.Context.current()));
107+Assert.assertNull(Context.current().get(password));
108+io.grpc.Context previous = grpcContext.attach();
109+try {
110+assertEquals(username.get(io.grpc.Context.current()), "jeff");
111+try (Scope scope = otelContext.makeCurrent()) {
112+Assert.assertEquals(Context.current().get(password), "valentine");
113+assertNull(username.get(io.grpc.Context.current()));
114+115+io.grpc.Context grpcContext2 = io.grpc.Context.current().withValue(username, "frank");
116+io.grpc.Context previous2 = grpcContext2.attach();
117+try {
118+assertEquals(username.get(io.grpc.Context.current()), "frank");
119+Assert.assertEquals(Context.current().get(password), "valentine");
120+ } finally {
121+grpcContext2.detach(previous2);
122+ }
123+assertNull(username.get(io.grpc.Context.current()));
124+Assert.assertEquals(Context.current().get(password), "valentine");
125+ }
126+ } finally {
127+grpcContext.detach(previous);
128+ }
129+Assert.assertNull(username.get(io.grpc.Context.current()));
130+Assert.assertNull(Context.current().get(password));
131+ }
132+133+@Test
134+public void grpcContextDetachError() {
135+io.grpc.Context grpcContext = io.grpc.Context.current().withValue(username, "jeff");
136+io.grpc.Context previous = grpcContext.attach();
137+try {
138+previous.detach(grpcContext);
139+assertEquals(username.get(io.grpc.Context.current()), "jeff");
140+ } finally {
141+grpcContext.detach(previous);
142+ }
143+ }
144+}