Skip to content

Commit 98b2c45

Browse files
sjoerdtalsmasjoerdtalsma
authored
Task: Simplify GrpcContextManager (#656)
* Simplify gRPC manager. Signed-off-by: Sjoerd Talsma <sjoerdtalsma@users.noreply.github.com> * Simplify gRPC context manager and add trace logging for debugging. Signed-off-by: Sjoerd Talsma <sjoerdtalsma@users.noreply.github.com> * Fix too-early string formatting in new log-trace statements. Signed-off-by: Sjoerd Talsma <sjoerdtalsma@users.noreply.github.com> * Fix non-final parameters in lambda's Signed-off-by: Sjoerd Talsma <sjoerdtalsma@users.noreply.github.com> --------- Signed-off-by: Sjoerd Talsma <sjoerdtalsma@users.noreply.github.com>
1 parent 5d8f9ba commit 98b2c45

File tree

4 files changed

+176
-124
lines changed

4 files changed

+176
-124
lines changed

‎managers/context-manager-grpc/src/main/java/io/grpc/override/ContextStorageOverride.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
* @since 2.0.0
2626
*/
2727
public class ContextStorageOverride extends Context.Storage {
28+
2829
/**
2930
* Create a new context storage override that delegates to the {@link GrpcContextManager}.
3031
*/

‎managers/context-manager-grpc/src/main/java/nl/talsmasoftware/context/managers/grpc/GrpcContextManager.java‎

Lines changed: 96 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,9 @@
1919
import nl.talsmasoftware.context.api.Context;
2020
import nl.talsmasoftware.context.api.ContextManager;
2121
import nl.talsmasoftware.context.api.ContextSnapshot;
22+
import nl.talsmasoftware.context.api.ContextSnapshot.Reactivation;
2223

24+
import java.util.logging.Level;
2325
import java.util.logging.Logger;
2426

2527
import static io.grpc.Context.ROOT;
@@ -38,27 +40,14 @@
3840
*/
3941
public class GrpcContextManager extends io.grpc.Context.Storage implements ContextManager<io.grpc.Context> {
4042
private static final GrpcContextManager INSTANCE = new GrpcContextManager();
41-
private static final Key<ContextSnapshot> GRPC_SNAPSHOT_KEY = key("contextsnapshot-over-grpc");
42-
private static final Key<ContextSnapshot.Reactivation> GRPC_REACTIVATION_KEY = key("contextsnapshot-reactivation");
43+
4344
private static final Logger LOGGER = Logger.getLogger(GrpcContextManager.class.getName());
44-
private static final ThreadLocal<io.grpc.Context> STORAGE = new ThreadLocal<io.grpc.Context>() {
45-
@Override
46-
public void set(io.grpc.Context value) {
47-
if (ROOT.equals(value)) { // prevent ROOT from being set
48-
value = null;
49-
}
50-
if (value != null) {
51-
// do not store ContextSnapshot or its reactivation in the threadlocal!
52-
if (GRPC_SNAPSHOT_KEY.get(value) != null) {
53-
value = value.withValue(GRPC_SNAPSHOT_KEY, null);
54-
}
55-
if (GRPC_REACTIVATION_KEY.get(value) != null) {
56-
value = value.withValue(GRPC_REACTIVATION_KEY, null);
57-
}
58-
}
59-
super.set(value);
60-
}
61-
};
45+
private static final Key<ContextSnapshot> GRPC_SNAPSHOT_KEY = key("contextsnapshot-over-grpc");
46+
private static final Key<Reactivation> GRPC_REACTIVATION_KEY = key("contextsnapshot-reactivation");
47+
private static final ThreadLocal<GrpcContextManager> CAPTURING = new ThreadLocal<>();
48+
private static final io.grpc.Context DUMMY = io.grpc.Context.ROOT.withValue(key("dummy"), "dummy");
49+
50+
private static final ThreadLocal<io.grpc.Context> STORAGE = new ThreadLocal<>();
6251

6352
/**
6453
* Default constructor for java8 ServiceLoader.
@@ -76,22 +65,6 @@ public static GrpcContextManager provider() {
7665
return INSTANCE;
7766
}
7867

79-
/**
80-
* The current gRPC context.
81-
*
82-
* <p>
83-
* This is the managed gRPC context, which does <em>not</em> contain a {@link ContextSnapshot}.
84-
* This method is accessed by the {@link ContextSnapshot} implementation itself to include the gRPC context in the
85-
* captured {@link ContextSnapshot}.
86-
*
87-
* @return The current gRPC context if available, or {@code null} otherwise.
88-
* @see #current()
89-
*/
90-
@Override
91-
public io.grpc.Context getActiveContextValue() {
92-
return STORAGE.get();
93-
}
94-
9568
/**
9669
* The current gRPC context.
9770
*
@@ -105,26 +78,20 @@ public io.grpc.Context getActiveContextValue() {
10578
*/
10679
@Override
10780
public io.grpc.Context current() {
108-
io.grpc.Context current = getActiveContextValue();
109-
return (current == null ? ROOT : current).withValue(GRPC_SNAPSHOT_KEY, ContextSnapshot.capture());
110-
}
111-
112-
/**
113-
* Activate the given gRPC context.
114-
*
115-
* <p>
116-
* This stores the given gRPC context as the current gRPC context. The previous gRPC context will be restored
117-
* when the returned {@link Context} is closed.
118-
*
119-
* @param value The gRPC context to activate.
120-
* @return A context that restores the previous gRPC context when closed.
121-
* @see #doAttach(io.grpc.Context)
122-
*/
123-
@Override
124-
public Context activate(io.grpc.Context value) {
125-
io.grpc.Context previous = STORAGE.get();
126-
STORAGE.set(value);
127-
return () -> STORAGE.set(previous);
81+
if (CAPTURING.get() != null) {
82+
LOGGER.finest("gRPC current(): Returning ROOT context because already capturing.");
83+
return ROOT;
84+
}
85+
try {
86+
LOGGER.finest("--> gRPC current(): Capturing gRPC context.");
87+
CAPTURING.set(this);
88+
ContextSnapshot snapshot = ContextSnapshot.capture();
89+
io.grpc.Context current = nullToRoot(STORAGE.get()).withValue(GRPC_SNAPSHOT_KEY, snapshot);
90+
LOGGER.finest(() -> "<-- gRPC current(): Returning current gRPC context " + current + " with captured " + snapshot + ".");
91+
return current;
92+
} finally {
93+
CAPTURING.remove();
94+
}
12895
}
12996

13097
/**
@@ -137,24 +104,19 @@ public Context activate(io.grpc.Context value) {
137104
* It <em>must</em> be restored in the same thread, which then closes the reactivated ContextSnapshot.
138105
*
139106
* @param toAttach the context to be attached
140-
* @return The previous gRPC context to be restored in the corresponding detach call.
107+
* @return The previous gRPC context to be restored in the corresponding 'detach' call.
141108
* @see #activate(io.grpc.Context)
142109
* @see #detach(io.grpc.Context, io.grpc.Context)
143110
*/
144111
@Override
145112
public io.grpc.Context doAttach(io.grpc.Context toAttach) {
146-
io.grpc.Context toRestore = STORAGE.get();
147-
if (toRestore == null) {
148-
toRestore = ROOT;
149-
}
150-
STORAGE.set(toAttach);
151-
113+
io.grpc.Context toRestore = nullToRoot(STORAGE.get());
152114
ContextSnapshot snapshot = toAttach == null ? null : GRPC_SNAPSHOT_KEY.get(toAttach);
153115
if (snapshot != null) {
116+
toAttach = toAttach.withValue(GRPC_SNAPSHOT_KEY, null);
154117
toRestore = toRestore.withValue(GRPC_REACTIVATION_KEY, snapshot.reactivate());
155-
} else {
156-
LOGGER.warning(() -> "No context snapshot found in gRPC context to be attached: " + toAttach);
157118
}
119+
STORAGE.set(rootOrDummyToNull(toAttach));
158120
return toRestore;
159121
}
160122

@@ -172,14 +134,16 @@ public io.grpc.Context doAttach(io.grpc.Context toAttach) {
172134
*/
173135
@Override
174136
public void detach(io.grpc.Context toDetach, io.grpc.Context toRestore) {
175-
ContextSnapshot.Reactivation reactivation = GRPC_REACTIVATION_KEY.get(toRestore);
176-
if (reactivation == null) {
177-
// should this ever happen?
178-
LOGGER.warning(() -> "No snapshot reactivation present in gRPC context to be restored: " + toRestore);
179-
} else {
137+
LOGGER.log(Level.FINEST, "--> gRPC detach(): Detaching gRPC context by restoring {0}.", toRestore);
138+
Reactivation reactivation = toRestore == null ? null : GRPC_REACTIVATION_KEY.get(toRestore);
139+
if (reactivation != null) {
140+
LOGGER.finest(() -> "--> grpc detach(): Closing reactivation from toRestore: " + reactivation + ".");
180141
reactivation.close();
142+
toRestore = toRestore.withValue(GRPC_REACTIVATION_KEY, null);
143+
LOGGER.finest(() -> "--- gRPC detach(): Reactivation from toRestore closed: " + reactivation + ".");
181144
}
182-
STORAGE.set(toRestore);
145+
STORAGE.set(rootOrDummyToNull(toRestore));
146+
LOGGER.log(Level.FINEST, "<-- gRPC detach(): Detached gRPC context by restoring {0}.", toRestore);
183147
}
184148

185149
/**
@@ -188,5 +152,66 @@ public void detach(io.grpc.Context toDetach, io.grpc.Context toRestore) {
188152
@Override
189153
public void clear() {
190154
STORAGE.remove();
155+
LOGGER.finest("<-> clear(): Cleared current gRPC context.");
156+
}
157+
158+
/**
159+
* The current gRPC context.
160+
*
161+
* <p>
162+
* This is the managed gRPC context, which does <em>not</em> contain a {@link ContextSnapshot}.
163+
* This method is accessed by the {@link ContextSnapshot} implementation itself to include the gRPC context in the
164+
* captured {@link ContextSnapshot}.
165+
*
166+
* @return The current gRPC context if available, or {@code null} otherwise.
167+
* @see #current()
168+
*/
169+
@Override
170+
public io.grpc.Context getActiveContextValue() {
171+
if (CAPTURING.get() != null) {
172+
LOGGER.finest("<-> getActiveContextValue(): Returning DUMMY context because already capturing.");
173+
return DUMMY;
174+
}
175+
try {
176+
LOGGER.finest("--> getActiveContextValue(): Capturing active context value.");
177+
CAPTURING.set(this);
178+
io.grpc.Context current = STORAGE.get();
179+
LOGGER.finest(() -> "<-- getActiveContextValue() Returning current gRPC context " + current + " from storage.");
180+
return current;
181+
} finally {
182+
CAPTURING.remove();
183+
}
184+
}
185+
186+
/**
187+
* Activate the given gRPC context.
188+
*
189+
* <p>
190+
* This stores the given gRPC context as the current gRPC context. The previous gRPC context will be restored
191+
* when the returned {@link Context} is closed.
192+
*
193+
* @param value The gRPC context to activate.
194+
* @return A context that restores the previous gRPC context when closed.
195+
* @see #doAttach(io.grpc.Context)
196+
*/
197+
@Override
198+
public Context activate(io.grpc.Context value) {
199+
if (value == DUMMY) {
200+
LOGGER.finest("<-> activate(DUMMY): Returning no-op context.");
201+
return () -> {
202+
};
203+
}
204+
LOGGER.finest(() -> "--> activate(" + value + "): Attaching gRPC context " + value + " as current context.");
205+
final io.grpc.Context toRestore = doAttach(value);
206+
LOGGER.finest(() -> "<-- activate(" + value + "): Returning context that restores " + toRestore + ".");
207+
return () -> detach(value, toRestore);
208+
}
209+
210+
private static io.grpc.Context nullToRoot(io.grpc.Context context) {
211+
return context == null ? ROOT : context;
212+
}
213+
214+
private static io.grpc.Context rootOrDummyToNull(io.grpc.Context context) {
215+
return context == ROOT || context == DUMMY ? null : context;
191216
}
192217
}

‎managers/context-manager-grpc/src/test/java/nl/talsmasoftware/context/managers/grpc/GrpcContextManagerTest.java‎

Lines changed: 72 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -17,84 +17,103 @@
1717

1818
import io.grpc.Context;
1919
import nl.talsmasoftware.context.api.ContextManager;
20-
import nl.talsmasoftware.context.api.ContextSnapshot;
2120
import nl.talsmasoftware.context.managers.locale.CurrentLocaleHolder;
2221
import org.junit.jupiter.api.AfterEach;
22+
import org.junit.jupiter.api.BeforeAll;
2323
import org.junit.jupiter.api.BeforeEach;
24-
import org.junit.jupiter.api.Disabled;
2524
import org.junit.jupiter.api.Test;
2625

26+
import java.io.IOException;
2727
import java.util.Locale;
2828
import java.util.UUID;
29-
import java.util.concurrent.ExecutorService;
30-
import java.util.concurrent.Executors;
31-
import java.util.concurrent.TimeUnit;
29+
import java.util.logging.LogManager;
3230

3331
import static org.assertj.core.api.Assertions.assertThat;
34-
import static org.assertj.core.api.InstanceOfAssertFactories.OPTIONAL;
3532

36-
@Disabled("gRPC context manager must be refactored to properly support handling null.")
33+
//@Disabled("gRPC context manager must be refactored to properly support handling null.")
3734
class GrpcContextManagerTest {
3835
static final Context.Key<String> TEST_KEY = Context.key("test-key");
39-
static final Locale DUTCH = Locale.of("nl", "NL");
40-
static final ExecutorService THREAD_POOL = Executors.newCachedThreadPool();
36+
37+
GrpcContextManager subject = GrpcContextManager.provider();
38+
39+
@BeforeAll
40+
static void setUpLogging() throws IOException {
41+
LogManager.getLogManager().readConfiguration(GrpcContextManagerTest.class.getResourceAsStream("/logging.properties"));
42+
}
4143

4244
@BeforeEach
4345
@AfterEach
4446
void clearAllContexts() {
47+
subject.clear();
4548
ContextManager.clearAll();
4649
}
4750

4851
@Test
49-
void grpcContextIsPropagatedToBackgroundThread() {
52+
void testNullContext() {
53+
try (var ignored = subject.activate(Context.ROOT.withValue(TEST_KEY, "test1"))) {
54+
assertThat(subject.getActiveContextValue()).isNotNull();
55+
assertThat(TEST_KEY.get()).isEqualTo("test1");
56+
57+
try (var ignored2 = subject.activate(null)) {
58+
assertThat(subject.getActiveContextValue()).isNull();
59+
assertThat(TEST_KEY.get()).isNull();
60+
61+
try (var ignored3 = subject.activate(Context.ROOT.withValue(TEST_KEY, "test2"))) {
62+
assertThat(subject.getActiveContextValue()).isNotNull();
63+
assertThat(TEST_KEY.get()).isEqualTo("test2");
64+
}
65+
66+
assertThat(subject.getActiveContextValue()).isNull();
67+
assertThat(TEST_KEY.get()).isNull();
68+
}
69+
70+
assertThat(subject.getActiveContextValue()).isNotNull();
71+
assertThat(TEST_KEY.get()).isEqualTo("test1");
72+
}
73+
}
74+
75+
@Test
76+
void testSetGrpcContextValue() {
77+
String testValue1 = "test-" + UUID.randomUUID();
78+
String testValue2 = "test2-" + UUID.randomUUID();
79+
assertThat(TEST_KEY.get()).isNull();
80+
Context.current().withValue(TEST_KEY, testValue1).run(() -> {
81+
assertThat(TEST_KEY.get()).isEqualTo(testValue1);
82+
Context.current().withValue(TEST_KEY, testValue2).run(() -> {
83+
assertThat(TEST_KEY.get()).isEqualTo(testValue2);
84+
});
85+
assertThat(TEST_KEY.get()).isEqualTo(testValue1);
86+
});
87+
assertThat(TEST_KEY.get()).isNull();
88+
}
89+
90+
@Test
91+
void testActivateGrpcContext() {
92+
// given
5093
String testValue = "test-" + UUID.randomUUID();
51-
Context context = Context.current().withValue(TEST_KEY, testValue).attach();
52-
try {
53-
54-
assertThat(THREAD_POOL.submit(() -> TEST_KEY.get()))
55-
.as("Test value from gRPC context in plain thread")
56-
.succeedsWithin(1, TimeUnit.SECONDS)
57-
.isNull();
58-
59-
ContextSnapshot snapshot = ContextSnapshot.capture();
60-
assertThat(THREAD_POOL.submit(snapshot.wrap(() -> TEST_KEY.get())))
61-
.as("Test value from gRPC context in thread with snapshot")
62-
.succeedsWithin(1, TimeUnit.SECONDS)
63-
.isEqualTo(testValue);
64-
65-
} finally {
66-
context.detach(Context.current());
94+
Context randomContext = Context.current().withValue(TEST_KEY, testValue);
95+
assertThat(TEST_KEY.get()).isNull();
96+
97+
// when
98+
try (var ignored = subject.activate(randomContext)) {
99+
// then
100+
assertThat(TEST_KEY.get()).isEqualTo(testValue);
67101
}
102+
103+
assertThat(TEST_KEY.get()).isNull();
68104
}
69105

70106
@Test
71-
void contextManagersArePropagatedWithGrpcContext() {
72-
CurrentLocaleHolder.set(null);
73-
Context contextWithNoLocale = Context.current();
74-
75-
CurrentLocaleHolder.set(DUTCH);
76-
77-
assertThat(THREAD_POOL.submit(CurrentLocaleHolder::get))
78-
.as("Current Locale in plain thread")
79-
.succeedsWithin(1, TimeUnit.SECONDS)
80-
.asInstanceOf(OPTIONAL)
81-
.isEmpty();
82-
83-
assertThat(THREAD_POOL.submit(Context.current().wrap(CurrentLocaleHolder::get)))
84-
.as("Current Locale in thread with gRPC context")
85-
.succeedsWithin(1, TimeUnit.SECONDS)
86-
.asInstanceOf(OPTIONAL)
87-
.contains(DUTCH);
88-
89-
assertThat(THREAD_POOL.submit(contextWithNoLocale.wrap(CurrentLocaleHolder::get)))
90-
.as("Current Locale in thread with gRPC context with no locale set")
91-
.succeedsWithin(1, TimeUnit.SECONDS)
92-
.asInstanceOf(OPTIONAL)
93-
.isEmpty();
94-
95-
assertThat(THREAD_POOL.submit(Context.ROOT.wrap(CurrentLocaleHolder::get)))
96-
.as("Current Locale in plain thread")
97-
.succeedsWithin(1, TimeUnit.SECONDS);
107+
void grpcContextAutomaticallyPropagatesLocaleContext() {
108+
// given
109+
CurrentLocaleHolder.set(Locale.GERMANY);
110+
Context grpcContext = Context.current();
111+
CurrentLocaleHolder.set(Locale.FRANCE);
112+
113+
assertThat(CurrentLocaleHolder.getOrDefault()).isEqualTo(Locale.FRANCE);
114+
grpcContext.run(() -> assertThat(CurrentLocaleHolder.getOrDefault()).isEqualTo(Locale.GERMANY));
115+
116+
assertThat(CurrentLocaleHolder.getOrDefault()).isEqualTo(Locale.FRANCE);
98117
}
99118

100119
}

0 commit comments

Comments
 (0)