Skip to content
This repository was archived by the owner on Dec 23, 2023. It is now read-only.

Commit 1a7feea

Browse files
sebright2sebright2
authored
Merge pull request #22 from dinooliva/grpc-context
Adds grpc-based thread-local support
2 parents e7e8752 + 35e0e5a commit 1a7feea

File tree

6 files changed

+242
-6
lines changed

6 files changed

+242
-6
lines changed

‎BUILD‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,13 +29,37 @@ java_library(
2929
],
3030
)
3131

32+
java_library(
33+
name = "census-grpc",
34+
srcs = glob(["grpc/java/com/google/census/*.java"]),
35+
deps = [
36+
":census-core",
37+
"@grpc_context//jar",
38+
"@jsr305//jar",
39+
],
40+
)
41+
42+
java_library(
43+
name = "census-grpc_native",
44+
srcs = glob(["grpc_native/java/com/google/census/*.java"]),
45+
deps = [
46+
":census-core",
47+
":census-grpc",
48+
"@grpc_context//jar",
49+
"@jsr305//jar",
50+
],
51+
)
52+
3253
java_binary(
3354
name = "CensusRunner",
3455
srcs = ["examples/java/com/google/census/CensusRunner.java"],
3556
main_class = "com.google.census.examples.CensusRunner",
3657
deps = [
3758
":census-core",
3859
":census-core_native",
60+
":census-grpc",
61+
":census-grpc_native",
62+
"@grpc_context//jar",
3963
"@guava//jar",
4064
"@jsr305//jar",
4165
],
@@ -68,6 +92,22 @@ java_test(
6892
],
6993
)
7094

95+
java_test(
96+
name = "CensusGrpcContextTest",
97+
srcs = ["grpc/javatests/com/google/census/CensusGrpcContextTest.java"],
98+
deps = [
99+
":census-core",
100+
":census-core_native",
101+
":census-grpc",
102+
":census-grpc_native",
103+
"@grpc_context//jar",
104+
"@guava//jar",
105+
"@jsr305//jar",
106+
"@junit//jar",
107+
"@truth//jar",
108+
],
109+
)
110+
71111
java_test(
72112
name = "MetricMapTest",
73113
srcs = ["core/javatests/com/google/census/MetricMapTest.java"],

‎WORKSPACE‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,11 @@
99
# See the License for the specific language governing permissions and
1010
# limitations under the License.
1111

12+
maven_jar(
13+
name = "grpc_context",
14+
artifact = "io.grpc:grpc-context:1.0.1",
15+
)
16+
1217
maven_jar(
1318
name = "guava",
1419
artifact = "com.google.guava:guava:19.0",

‎examples/java/com/google/census/CensusRunner.java‎

Lines changed: 40 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,14 @@
1515

1616
import com.google.census.Census;
1717
import com.google.census.CensusContext;
18+
import com.google.census.CensusGrpcContext;
1819
import com.google.census.MetricMap;
1920
import com.google.census.MetricName;
2021
import com.google.census.TagKey;
2122
import com.google.census.TagValue;
2223

24+
import io.grpc.Context;
25+
2326
public class CensusRunner {
2427
private static final TagKey K1 = new TagKey("k1");
2528
private static final TagKey K2 = new TagKey("k2");
@@ -36,11 +39,42 @@ public class CensusRunner {
3639

3740
public static void main(String args[]) {
3841
System.out.println("Hello Census World");
39-
System.out.println("Default Tags: " + Census.getCensusContextFactory().getDefault());
40-
CensusContext context1 = Census.getCensusContextFactory().getDefault().with(K1, V1, K2, V2);
41-
System.out.println("Context1 Tags: " + context1);
42-
CensusContext context2 = context1.with(K3, V3, K4, V4);
43-
System.out.println("Context2 Tags: " + context2);
44-
context2.record(MetricMap.of(M1, 0.2, M2, 0.4));
42+
System.out.println("Default Tags: " + DEFAULT);
43+
System.out.println("Current Tags: " + getCurrentCensusContext());
44+
Context context1 = withCurrent(DEFAULT.with(K1, V1, K2, V2));
45+
Context original = context1.attach();
46+
try {
47+
System.out.println(" Current Tags: " + getCurrentCensusContext());
48+
System.out.println(" Current == Default + tags1: "
49+
+ getCurrentCensusContext().equals(getCensusContext(context1)));
50+
Context context2 = withCurrent(getCurrentCensusContext().with(K3, V3, K4, V4));
51+
context2.attach();
52+
try {
53+
System.out.println(" Current Tags: " + getCurrentCensusContext());
54+
System.out.println(" Current == Default + tags1 + tags2: "
55+
+ getCurrentCensusContext().equals(getCensusContext(context2)));
56+
getCurrentCensusContext().record(MetricMap.of(M1, 0.2, M2, 0.4));
57+
} finally {
58+
context2.detach(context1);
59+
}
60+
} finally {
61+
context1.detach(original);
62+
}
63+
System.out.println("Current == Default: "
64+
+ getCurrentCensusContext().equals(DEFAULT));
65+
}
66+
67+
private static final CensusContext DEFAULT = Census.getCensusContextFactory().getDefault();
68+
69+
private static final CensusContext getCurrentCensusContext() {
70+
return getCensusContext(Context.current());
71+
}
72+
73+
private static final CensusContext getCensusContext(Context context) {
74+
return CensusGrpcContext.getInstance().get(context);
75+
}
76+
77+
private static final Context withCurrent(CensusContext context) {
78+
return CensusGrpcContext.getInstance().withCensusContext(Context.current(), context);
4579
}
4680
}

‎grpc/java/com/google/census/CensusGrpcContext.java‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
/*
2+
* Copyright 2016, Google Inc.
3+
* Licensed under the Apache License, Version 2.0 (the "License");
4+
* you may not use this file except in compliance with the License.
5+
* You may obtain a copy of the License at
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
* Unless required by applicable law or agreed to in writing, software
8+
* distributed under the License is distributed on an "AS IS" BASIS,
9+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
* See the License for the specific language governing permissions and
11+
* limitations under the License.
12+
*/
13+
14+
package com.google.census;
15+
16+
import io.grpc.Context;
17+
18+
/**
19+
* Utils for embedding {@link CensusContext} in gRPC {@link Context} to support thread-local
20+
* implicit propagation of the {@link CensusContext}.
21+
*/
22+
public abstract class CensusGrpcContext {
23+
private static final CensusGrpcContext INSTANCE = new Provider<CensusGrpcContext>(
24+
"com.google.census.CensusGrpcContextImpl").newInstance();
25+
26+
/** Returns the default {@link CensusGrpContext} instance. */
27+
public static CensusGrpcContext getInstance() {
28+
return INSTANCE;
29+
}
30+
31+
/**
32+
* Adds this {@link CensusContext} to the specified {@link Context} and returns the resulting
33+
* {@link Context}.
34+
*/
35+
public abstract Context withCensusContext(Context context, CensusContext censusContext);
36+
37+
/**
38+
* Extracts the CensusContext from the given Context. If there is no associated CensusContext,
39+
* the default CensusContext is returned.
40+
*/
41+
public abstract CensusContext get(Context context);
42+
}

‎grpc/javatests/com/google/census/CensusGrpcContextTest.java‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
/*
2+
* Copyright 2016, Google Inc.
3+
* Licensed under the Apache License, Version 2.0 (the "License");
4+
* you may not use this file except in compliance with the License.
5+
* You may obtain a copy of the License at
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
* Unless required by applicable law or agreed to in writing, software
8+
* distributed under the License is distributed on an "AS IS" BASIS,
9+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
* See the License for the specific language governing permissions and
11+
* limitations under the License.
12+
*/
13+
14+
package com.google.census;
15+
16+
import static com.google.common.truth.Truth.assertThat;
17+
18+
import io.grpc.Context;
19+
import java.util.concurrent.ExecutorService;
20+
import java.util.concurrent.Executors;
21+
import java.util.concurrent.Future;
22+
import org.junit.Test;
23+
import org.junit.runner.RunWith;
24+
import org.junit.runners.JUnit4;
25+
26+
/**
27+
* Tests for {@link CensusGrpcContext}.
28+
*/
29+
@RunWith(JUnit4.class)
30+
public class CensusGrpcContextTest {
31+
private static final CensusContext DEFAULT = Census.getCensusContextFactory().getDefault();
32+
private static final CensusGrpcContext CENSUS_GRPC_CONTEXT = CensusGrpcContext.getInstance();
33+
34+
@Test
35+
public void withCensusContext() {
36+
Context original = Context.current();
37+
assertThat(CENSUS_GRPC_CONTEXT.get(original)).isEqualTo(DEFAULT);
38+
CensusContext censusContext = DEFAULT.with(new TagKey("k1"), new TagValue("v1"));
39+
Context context = CENSUS_GRPC_CONTEXT.withCensusContext(original, censusContext);
40+
assertThat(CENSUS_GRPC_CONTEXT.get(context)).isEqualTo(censusContext);
41+
}
42+
43+
@Test
44+
public void attachCensusContext() {
45+
assertThat(CENSUS_GRPC_CONTEXT.get(Context.current())).isEqualTo(DEFAULT);
46+
testAttachCensusContext(DEFAULT.with(K1, V1)).run();
47+
assertThat(CENSUS_GRPC_CONTEXT.get(Context.current())).isEqualTo(DEFAULT);
48+
}
49+
50+
@Test
51+
public void attachCensusContextInMultipleThreads() throws Exception {
52+
assertThat(CENSUS_GRPC_CONTEXT.get(Context.current())).isEqualTo(DEFAULT);
53+
ExecutorService executor = Executors.newFixedThreadPool(2);
54+
Future<?> future1 = executor.submit(testAttachCensusContext(DEFAULT.with(K1, V1)));
55+
Future<?> future2 = executor.submit(testAttachCensusContext(DEFAULT.with(K2, V2)));
56+
future1.get();
57+
future2.get();
58+
assertThat(CENSUS_GRPC_CONTEXT.get(Context.current())).isEqualTo(DEFAULT);
59+
}
60+
61+
private static final TagKey K1 = new TagKey("k1");
62+
private static final TagKey K2 = new TagKey("k2");
63+
64+
private static final TagValue V1 = new TagValue("v1");
65+
private static final TagValue V2 = new TagValue("v2");
66+
67+
private static final Runnable testAttachCensusContext(final CensusContext censusContext) {
68+
return new Runnable() {
69+
@Override
70+
public void run() {
71+
assertThat(CENSUS_GRPC_CONTEXT.get(Context.current())).isEqualTo(DEFAULT);
72+
Context current = CENSUS_GRPC_CONTEXT.withCensusContext(Context.current(), censusContext);
73+
Context original = current.attach();
74+
assertThat(CENSUS_GRPC_CONTEXT.get(Context.current())).isEqualTo(censusContext);
75+
current.detach(original);
76+
assertThat(CENSUS_GRPC_CONTEXT.get(Context.current())).isEqualTo(DEFAULT);
77+
}
78+
};
79+
}
80+
}

‎grpc_native/java/com/google/census/CensusGrpcContextImpl.java‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
/*
2+
* Copyright 2016, Google Inc.
3+
* Licensed under the Apache License, Version 2.0 (the "License");
4+
* you may not use this file except in compliance with the License.
5+
* You may obtain a copy of the License at
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
* Unless required by applicable law or agreed to in writing, software
8+
* distributed under the License is distributed on an "AS IS" BASIS,
9+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
* See the License for the specific language governing permissions and
11+
* limitations under the License.
12+
*/
13+
14+
package com.google.census;
15+
16+
import io.grpc.Context;
17+
18+
/**
19+
* Utils for embedding {@link CensusContext} in gRPC {@link Context} to support thread-local
20+
* implicit propagation of the {@link CensusContext}.
21+
*/
22+
public class CensusGrpcContextImpl extends CensusGrpcContext {
23+
private static final Context.Key<CensusContext> CENSUS_CONTEXT_KEY =
24+
Context.keyWithDefault("CensusContextKey", Census.getCensusContextFactory().getDefault());
25+
26+
@Override
27+
public Context withCensusContext(Context context, CensusContext censusContext) {
28+
return context.withValue(CENSUS_CONTEXT_KEY, censusContext);
29+
}
30+
31+
@Override
32+
public CensusContext get(Context context) {
33+
return CENSUS_CONTEXT_KEY.get(context);
34+
}
35+
}

0 commit comments

Comments
 (0)