Add Option to Avoid Use of ThreadLocal

* Add an option to store contexts in a map instead of a `ThreadLocal`
* Subclass existing tests to verify functionality

Author: garyrussell

README.md

@@ -166,13 +166,18 @@ RetryTemplate.builder()
 
 ### Using `RetryContext`
 
-The method parameter for the `RetryCallback` is a `RetryContext`. Many callbacks ignore
-the context. However, if necessary, you can use it as an attribute bag to store data for
-the duration of the iteration.
-
-A `RetryContext` has a parent context if there is a nested retry in progress in the same
-thread. The parent context is occasionally useful for storing data that needs to be shared
-between calls to execute.
+The method parameter for the `RetryCallback` is a `RetryContext`.
+Many callbacks ignore the context.
+However, if necessary, you can use it as an attribute bag to store data for the duration of the iteration.
+It also has some useful properties, such as `retryCount`.
+
+A `RetryContext` has a parent context if there is a nested retry in progress in the same thread.
+The parent context is occasionally useful for storing data that needs to be shared between calls to execute.
+
+If you dont have access to the context directly, you can obtain the current context within the scope of the retries by calling `RetrySynchronizationManager.getContext()`.
+By default, the context is stored in a `ThreadLocal`.
+JEP 444 recommends that `ThreadLocal` should be avoided when using virtual threads, available in Java 21 and beyond.
+To store the contexts in a `Map` instead of a `ThreadLocal`, call `RetrySynchronizationManager.setUseThreadLocal(false)`.
 
 ### Using `RecoveryCallback`
 

src/main/java/org/springframework/retry/support/RetrySynchronizationManager.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2006-2022 the original author or authors.
+ * Copyright 2006-2023 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.
@@ -16,6 +16,10 @@
 
 package org.springframework.retry.support;
 
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.springframework.lang.Nullable;
 import org.springframework.retry.RetryCallback;
 import org.springframework.retry.RetryContext;
 import org.springframework.retry.RetryOperations;
@@ -30,6 +34,7 @@
  * {@link RetryOperations} implementations.
  *
  * @author Dave Syer
+ * @author Gary Russell
  *
  */
 public final class RetrySynchronizationManager {
@@ -39,13 +44,41 @@ private RetrySynchronizationManager() {
 
 	private static final ThreadLocal<RetryContext> context = new ThreadLocal<>();
 
+	private static final Map<Thread, RetryContext> contexts = new ConcurrentHashMap<>();
+
+	private static boolean useThreadLocal = true;
+
+	/**
+	 * Set to false to store the context in a map (keyed by the current thread) instead of
+	 * in a {@link ThreadLocal}. Recommended when using virtual threads.
+	 * @param use true to use a {@link ThreadLocal} (default true).
+	 * @since 2.0.3
+	 */
+	public static void setUseThreadLocal(boolean use) {
+		useThreadLocal = use;
+	}
+
+	/**
+	 * Return true if contexts are held in a ThreadLocal (default) rather than a Map.
+	 * @return the useThreadLocal
+	 * @since 2.0.3
+	 */
+	public static boolean isUseThreadLocal() {
+		return useThreadLocal;
+	}
+
 	/**
 	 * Public accessor for the locally enclosing {@link RetryContext}.
 	 * @return the current retry context, or null if there isn't one
 	 */
+	@Nullable
 	public static RetryContext getContext() {
-		RetryContext result = context.get();
-		return result;
+		if (useThreadLocal) {
+			return context.get();
+		}
+		else {
+			return contexts.get(Thread.currentThread());
+		}
 	}
 
 	/**
@@ -55,21 +88,40 @@ public static RetryContext getContext() {
 	 * @param context the new context to register
 	 * @return the old context if there was one
 	 */
+	@Nullable
 	public static RetryContext register(RetryContext context) {
-		RetryContext oldContext = getContext();
-		RetrySynchronizationManager.context.set(context);
-		return oldContext;
+		if (useThreadLocal) {
+			RetryContext oldContext = getContext();
+			RetrySynchronizationManager.context.set(context);
+			return oldContext;
+		}
+		else {
+			RetryContext oldContext = contexts.get(Thread.currentThread());
+			contexts.put(Thread.currentThread(), context);
+			return oldContext;
+		}
 	}
 
 	/**
 	 * Clear the current context at the end of a batch - should only be used by
 	 * {@link RetryOperations} implementations.
 	 * @return the old value if there was one.
 	 */
+	@Nullable
 	public static RetryContext clear() {
 		RetryContext value = getContext();
 		RetryContext parent = value == null ? null : value.getParent();
-		RetrySynchronizationManager.context.set(parent);
+		if (useThreadLocal) {
+			RetrySynchronizationManager.context.set(parent);
+		}
+		else {
+			if (parent != null) {
+				contexts.put(Thread.currentThread(), parent);
+			}
+			else {
+				contexts.remove(Thread.currentThread());
+			}
+		}
 		return value;
 	}
 

src/test/java/org/springframework/retry/annotation/EnableRetryNoThreadLocalTests.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright 2023 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.retry.annotation;
+
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+
+import org.springframework.retry.support.RetrySynchronizationManager;
+
+public class EnableRetryNoThreadLocalTests extends EnableRetryTests {
+
+	@BeforeAll
+	static void before() {
+		RetrySynchronizationManager.setUseThreadLocal(false);
+	}
+
+	@AfterAll
+	static void after() {
+		RetrySynchronizationManager.setUseThreadLocal(true);
+	}
+
+}

src/test/java/org/springframework/retry/annotation/EnableRetryWithBackoffNoThreadLocalTests.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2023 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.retry.annotation;
+
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+
+import org.springframework.retry.support.RetrySynchronizationManager;
+
+/**
+ * @author Gary Russell
+ */
+public class EnableRetryWithBackoffNoThreadLocalTests extends EnableRetryWithBackoffTests {
+
+	@BeforeAll
+	static void before() {
+		RetrySynchronizationManager.setUseThreadLocal(false);
+	}
+
+	@AfterAll
+	static void after() {
+		RetrySynchronizationManager.setUseThreadLocal(true);
+	}
+
+}

src/test/java/org/springframework/retry/support/RetrySynchronizationManagerNoThreadLocalTests.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2023 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.retry.support;
+
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+
+import org.springframework.retry.RetryContext;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class RetrySynchronizationManagerNoThreadLocalTests extends RetrySynchronizationManagerTests {
+
+	@BeforeAll
+	static void before() {
+		RetrySynchronizationManager.setUseThreadLocal(false);
+	}
+
+	@AfterAll
+	static void after() {
+		RetrySynchronizationManager.setUseThreadLocal(true);
+	}
+
+	@Override
+	@BeforeEach
+	public void setUp() {
+		RetrySynchronizationManagerTests.clearAll();
+		RetryContext status = RetrySynchronizationManager.getContext();
+		assertThat(status).isNull();
+	}
+
+}

src/test/java/org/springframework/retry/support/RetryTemplateNoThreadLocalTests.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2023 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.retry.support;
+
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+
+/**
+ * @author Gary Russell
+ */
+public class RetryTemplateNoThreadLocalTests extends RetryTemplateTests {
+
+	@BeforeAll
+	static void before() {
+		RetrySynchronizationManager.setUseThreadLocal(false);
+	}
+
+	@AfterAll
+	static void after() {
+		RetrySynchronizationManager.setUseThreadLocal(true);
+	}
+
+}