diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextManager.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextManager.java index cc36df21..52bdd856 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextManager.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextManager.java @@ -38,6 +38,8 @@ public abstract class ContextManager { private static final String CONTEXT_VIEW_KEY = ContextManager.class.getName() + ".CONTEXT_VIEW"; + private static final String THREAD_ID = ContextManager.class.getName() + ".THREAD_ID"; + private static final String THREAD_LOCAL_VALUES_KEY = ContextManager.class.getName() + ".THREAD_VALUES_ACCESSOR"; private static final String THREAD_LOCAL_ACCESSOR_KEY = ContextManager.class.getName() + ".THREAD_LOCAL_ACCESSOR"; @@ -80,38 +82,39 @@ public abstract class ContextManager { return context; } return context.putAll((ContextView) Context.of( - THREAD_LOCAL_VALUES_KEY, valuesMap, THREAD_LOCAL_ACCESSOR_KEY, accessor)); + THREAD_LOCAL_VALUES_KEY, valuesMap, + THREAD_LOCAL_ACCESSOR_KEY, accessor, + THREAD_ID, Thread.currentThread().getId())); } /** - * Look up saved ThreadLocal values and use them to re-establish ThreadLocal context. + * Look up saved ThreadLocal values and restore them if any are found. + * This is a no-op if invoked on the thread that values were extracted on. * @param contextView the reactor {@link ContextView} */ static void restoreThreadLocalValues(ContextView contextView) { ThreadLocalAccessor accessor = getThreadLocalAccessor(contextView); if (accessor != null) { - accessor.restoreValues(getThreadLocalValues(contextView)); + accessor.restoreValues(contextView.get(THREAD_LOCAL_VALUES_KEY)); } } /** - * Look up saved ThreadLocal values and remove associated ThreadLocal context. + * Look up saved ThreadLocal values and remove the ThreadLocal values. + * This is a no-op if invoked on the thread that values were extracted on. * @param contextView the reactor {@link ContextView} */ static void resetThreadLocalValues(ContextView contextView) { ThreadLocalAccessor accessor = getThreadLocalAccessor(contextView); if (accessor != null) { - accessor.resetValues(getThreadLocalValues(contextView)); + accessor.resetValues(contextView.get(THREAD_LOCAL_VALUES_KEY)); } } @Nullable - private static ThreadLocalAccessor getThreadLocalAccessor(ContextView contextView) { - return (contextView.hasKey(THREAD_LOCAL_ACCESSOR_KEY) ? contextView.get(THREAD_LOCAL_ACCESSOR_KEY) : null); - } - - private static Map getThreadLocalValues(ContextView contextView) { - return contextView.get(THREAD_LOCAL_VALUES_KEY); + private static ThreadLocalAccessor getThreadLocalAccessor(ContextView view) { + Long id = view.getOrDefault(THREAD_ID, null); + return (id != null && id != Thread.currentThread().getId() ? view.get(THREAD_LOCAL_ACCESSOR_KEY) : null); } } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/TestThreadLocalAccessor.java b/spring-graphql/src/test/java/org/springframework/graphql/TestThreadLocalAccessor.java index 66b780d7..326617ca 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/TestThreadLocalAccessor.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/TestThreadLocalAccessor.java @@ -34,8 +34,15 @@ public class TestThreadLocalAccessor implements ThreadLocalAccessor { @Nullable private Long threadId; + private boolean suppressThreadIdCheck; + public TestThreadLocalAccessor(ThreadLocal threadLocal) { + this(threadLocal, false); + } + + public TestThreadLocalAccessor(ThreadLocal threadLocal, boolean suppressThreadIdCheck) { this.threadLocal = threadLocal; + this.suppressThreadIdCheck = suppressThreadIdCheck; } @Override @@ -61,10 +68,16 @@ public class TestThreadLocalAccessor implements ThreadLocalAccessor { } private void saveThreadId() { + if (this.suppressThreadIdCheck) { + return; + } this.threadId = Thread.currentThread().getId(); } private void checkThreadId() { + if (this.suppressThreadIdCheck) { + return; + } assertThat(this.threadId).as("No threadId to check. Was extractValues not called?").isNotNull(); assertThat(Thread.currentThread().getId() != this.threadId) .as("ThreadLocal value extracted and restored on the same thread. Propagation not tested effectively.") diff --git a/spring-graphql/src/test/java/org/springframework/graphql/execution/ContextManagerTests.java b/spring-graphql/src/test/java/org/springframework/graphql/execution/ContextManagerTests.java new file mode 100644 index 00000000..d3035730 --- /dev/null +++ b/spring-graphql/src/test/java/org/springframework/graphql/execution/ContextManagerTests.java @@ -0,0 +1,73 @@ +/* + * Copyright 2002-2021 the original author or 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 + * + * https://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 org.springframework.graphql.execution; + +import java.time.Duration; + +import org.junit.jupiter.api.Test; +import reactor.core.publisher.Mono; +import reactor.util.context.Context; + +import org.springframework.graphql.TestThreadLocalAccessor; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Unit tests for {@link ContextManager}. + * @author Rossen Stoyanchev + */ +public class ContextManagerTests { + + @Test + void restoreThreadLocaValues() { + ThreadLocal threadLocal = new ThreadLocal<>(); + threadLocal.set("myValue"); + + Context context = ContextManager.extractThreadLocalValues( + new TestThreadLocalAccessor<>(threadLocal), Context.empty()); + try { + Mono.delay(Duration.ofMillis(10)) + .doOnNext(aLong -> { + assertThat(threadLocal.get()).isNull(); + ContextManager.restoreThreadLocalValues(context); + assertThat(threadLocal.get()).isEqualTo("myValue"); + ContextManager.resetThreadLocalValues(context); + }) + .block(); + } + finally { + threadLocal.remove(); + } + } + + @Test + void restoreThreadLocaValuesOnSameThreadIsNoOp() { + ThreadLocal threadLocal = new ThreadLocal<>(); + threadLocal.set("myValue"); + + Context context = ContextManager.extractThreadLocalValues( + new TestThreadLocalAccessor<>(threadLocal, true), Context.empty()); + + threadLocal.remove(); + ContextManager.restoreThreadLocalValues(context); + assertThat(threadLocal.get()).isNull(); + + threadLocal.set("anotherValue"); + ContextManager.resetThreadLocalValues(context); + assertThat(threadLocal.get()).isEqualTo("anotherValue"); + } + +}