From ab99dc2892db2e22866320779ede039b17f93a1e Mon Sep 17 00:00:00 2001 From: Rossen Stoyanchev Date: Thu, 24 Jun 2021 17:14:33 +0100 Subject: [PATCH] Improve ThreadLocal value handling ThreadLocal values from a Servlet container thread maybe end up being unnecessarily restored, e.g. if DataFetcher is invoked on the same thread and then also removed, which then impacts the filter chain. The ContextManager now saves the thread id when values are extracted and ignores restore or remove calls if still on the same thread. This should also be more optimal, avoiding ThreadLocal access if threads aren't switched. See gh-58 --- .../graphql/execution/ContextManager.java | 25 ++++--- .../graphql/TestThreadLocalAccessor.java | 13 ++++ .../execution/ContextManagerTests.java | 73 +++++++++++++++++++ 3 files changed, 100 insertions(+), 11 deletions(-) create mode 100644 spring-graphql/src/test/java/org/springframework/graphql/execution/ContextManagerTests.java 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"); + } + +}