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

Author: rstoyanchev

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 static Context extractThreadLocalValues(ThreadLocalAccessor accessor, Con
 			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<String, Object> 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);
 	}
 
 }

spring-graphql/src/test/java/org/springframework/graphql/TestThreadLocalAccessor.java

@@ -34,8 +34,15 @@ public class TestThreadLocalAccessor<T> implements ThreadLocalAccessor {
 	@Nullable
 	private Long threadId;
 
+	private boolean suppressThreadIdCheck;
+
 	public TestThreadLocalAccessor(ThreadLocal<T> threadLocal) {
+		this(threadLocal, false);
+	}
+
+	public TestThreadLocalAccessor(ThreadLocal<T> threadLocal, boolean suppressThreadIdCheck) {
 		this.threadLocal = threadLocal;
+		this.suppressThreadIdCheck = suppressThreadIdCheck;
 	}
 
 	@Override
@@ -61,10 +68,16 @@ public void resetValues(Map<String, Object> values) {
 	}
 
 	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.")

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<String> 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<String> 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");
+	}
+
+}