Skip to content

Commit 782a44a

Browse files
YifeiZhuangYifeiZhuang
authored
Implement ContextStorageOverride for opentelemetry context bridge (#11523)
1 parent 9b0c19e commit 782a44a

File tree

5 files changed

+292
-0
lines changed

5 files changed

+292
-0
lines changed

‎contextstorage/build.gradle

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
plugins {
2+
id "java-library"
3+
// until we are confident we like the name
4+
//id "maven-publish"
5+
6+
id "ru.vyarus.animalsniffer"
7+
}
8+
9+
description = 'gRPC: ContextStorageOverride'
10+
11+
dependencies {
12+
api project(':grpc-api')
13+
implementation libraries.opentelemetry.api
14+
15+
testImplementation libraries.junit,
16+
libraries.opentelemetry.sdk.testing,
17+
libraries.assertj.core
18+
testImplementation 'junit:junit:4.13.1'// opentelemetry.sdk.testing uses compileOnly for assertj
19+
20+
signature libraries.signature.java
21+
signature libraries.signature.android
22+
}
23+
24+
tasks.named("jar").configure {
25+
manifest {
26+
attributes('Automatic-Module-Name': 'io.grpc.override')
27+
}
28+
}

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

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
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 io.grpc.Context;
20+
21+
/**
22+
* Including this class in your dependencies will override the default gRPC context storage using
23+
* reflection. It is a bridge between {@link io.grpc.Context} and
24+
* {@link io.opentelemetry.context.Context}, i.e. propagating io.grpc.context.Context also
25+
* propagates io.opentelemetry.context, and propagating io.opentelemetry.context will also propagate
26+
* io.grpc.context.
27+
*/
28+
public final class ContextStorageOverride extends Context.Storage {
29+
30+
private final Context.Storage delegate = new OpenTelemetryContextStorage();
31+
32+
@Override
33+
public Context doAttach(Context toAttach) {
34+
return delegate.doAttach(toAttach);
35+
}
36+
37+
@Override
38+
public void detach(Context toDetach, Context toRestore) {
39+
delegate.detach(toDetach, toRestore);
40+
}
41+
42+
@Override
43+
public Context current() {
44+
return delegate.current();
45+
}
46+
}

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

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
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 io.grpc.Context;
20+
import io.opentelemetry.context.ContextKey;
21+
import io.opentelemetry.context.Scope;
22+
import java.util.logging.Level;
23+
import java.util.logging.Logger;
24+
25+
/**
26+
* A Context.Storage implementation that attaches io.grpc.context to OpenTelemetry's context and
27+
* io.opentelemetry.context is also saved in the io.grpc.context.
28+
* Bridge between {@link io.grpc.Context} and {@link io.opentelemetry.context.Context}.
29+
*/
30+
final class OpenTelemetryContextStorage extends Context.Storage {
31+
private static final Logger logger = Logger.getLogger(
32+
OpenTelemetryContextStorage.class.getName());
33+
34+
private static final io.grpc.Context.Key<io.opentelemetry.context.Context> OTEL_CONTEXT_OVER_GRPC
35+
= io.grpc.Context.key("otel-context-over-grpc");
36+
private static final Context.Key<Scope> OTEL_SCOPE = Context.key("otel-scope");
37+
private static final ContextKey<io.grpc.Context> GRPC_CONTEXT_OVER_OTEL =
38+
ContextKey.named("grpc-context-over-otel");
39+
40+
@Override
41+
@SuppressWarnings("MustBeClosedChecker")
42+
public Context doAttach(Context toAttach) {
43+
io.grpc.Context previous = current();
44+
io.opentelemetry.context.Context otelContext = OTEL_CONTEXT_OVER_GRPC.get(toAttach);
45+
if (otelContext == null) {
46+
otelContext = io.opentelemetry.context.Context.current();
47+
}
48+
Scope scope = otelContext.with(GRPC_CONTEXT_OVER_OTEL, toAttach).makeCurrent();
49+
return previous.withValue(OTEL_SCOPE, scope);
50+
}
51+
52+
@Override
53+
public void detach(Context toDetach, Context toRestore) {
54+
Scope scope = OTEL_SCOPE.get(toRestore);
55+
if (scope == null) {
56+
logger.log(
57+
Level.SEVERE, "Detaching context which was not attached.");
58+
} else {
59+
scope.close();
60+
}
61+
}
62+
63+
@Override
64+
public Context current() {
65+
io.opentelemetry.context.Context otelCurrent = io.opentelemetry.context.Context.current();
66+
io.grpc.Context grpcCurrent = otelCurrent.get(GRPC_CONTEXT_OVER_OTEL);
67+
if (grpcCurrent == null) {
68+
grpcCurrent = Context.ROOT;
69+
}
70+
return grpcCurrent.withValue(OTEL_CONTEXT_OVER_GRPC, otelCurrent);
71+
}
72+
}

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

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
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+
}

‎settings.gradle

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ include ":grpc-istio-interop-testing"
7777
include ":grpc-inprocess"
7878
include ":grpc-util"
7979
include ":grpc-opentelemetry"
80+
include ":grpc-opentelemetry-context-storage-override"
8081

8182
project(':grpc-api').projectDir = "$rootDir/api" as File
8283
project(':grpc-core').projectDir = "$rootDir/core" as File
@@ -113,6 +114,7 @@ project(':grpc-istio-interop-testing').projectDir = "$rootDir/istio-interop-test
113114
project(':grpc-inprocess').projectDir = "$rootDir/inprocess" as File
114115
project(':grpc-util').projectDir = "$rootDir/util" as File
115116
project(':grpc-opentelemetry').projectDir = "$rootDir/opentelemetry" as File
117+
project(':grpc-opentelemetry-context-storage-override').projectDir = "$rootDir/contextstorage" as File
116118

117119
if (settings.hasProperty('skipCodegen') && skipCodegen.toBoolean()) {
118120
println '*** Skipping the build of codegen and compilation of proto files because skipCodegen=true'

0 commit comments

Comments
 (0)