A variant of RetryContext propagation from client to worker thread

src/main/java/org/springframework/retry/concurrent/AbstractDelegatingRetryContextSupport.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2002-2019 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.concurrent;
+
+import java.util.concurrent.Callable;
+
+import org.springframework.retry.RetryContext;
+
+/**
+ * An internal support class that wraps {@link Callable} with
+ * {@link DelegatingRetryContextCallable} and {@link Runnable} with
+ * {@link DelegatingRetryContextRunnable}
+ * Inspired by org.springframework.security.concurrent package.
+ *
+ * @author Rob Winch
+ * @author Aleksandr Shamukov
+ */
+abstract class AbstractDelegatingRetryContextSupport {
+
+	private final RetryContext retryContext;
+
+	/**
+	 * Creates a new {@link AbstractDelegatingRetryContextSupport} that uses the specified
+	 * {@link RetryContext}.
+	 * @param retryContext the {@link RetryContext} to use for each
+	 * {@link DelegatingRetryContextRunnable} and each
+	 * {@link DelegatingRetryContextCallable} or null to default to the current
+	 * {@link RetryContext}.
+	 */
+	AbstractDelegatingRetryContextSupport(RetryContext retryContext) {
+		this.retryContext = retryContext;
+	}
+
+	protected final Runnable wrap(Runnable delegate) {
+		return DelegatingRetryContextRunnable.create(delegate, retryContext);
+	}
+
+	protected final <T> Callable<T> wrap(Callable<T> delegate) {
+		return DelegatingRetryContextCallable.create(delegate, retryContext);
+	}
+
+}

src/main/java/org/springframework/retry/concurrent/DelegatingRetryContextCallable.java

@@ -0,0 +1,115 @@
+/*
+ * Copyright 2002-2019 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.concurrent;
+
+import org.springframework.retry.RetryContext;
+import org.springframework.retry.support.RetrySynchronizationManager;
+import org.springframework.util.Assert;
+
+import java.util.concurrent.Callable;
+
+/**
+ * <p>
+ * Wraps a delegate {@link Callable} with logic for setting up a {@link RetryContext}
+ * before invoking the delegate {@link Callable} and then removing the
+ * {@link RetryContext} after the delegate has completed.
+ * </p>
+ * <p>
+ * If there is a {@link RetryContext} that already exists, it will be restored after the
+ * {@link #call()} method is invoked.
+ * </p>
+ * Inspired by org.springframework.security.concurrent package.
+ *
+ * @author Rob Winch
+ * @author Aleksandr Shamukov
+ */
+public final class DelegatingRetryContextCallable<V> implements Callable<V> {
+
+	private final Callable<V> delegate;
+
+	/**
+	 * The {@link RetryContext} that the delegate {@link Callable} will be ran as.
+	 */
+	private final RetryContext delegateRetryContext;
+
+	/**
+	 * The {@link RetryContext} that was on the {@link RetrySynchronizationManager} prior
+	 * to being set to the delegateRetryContext.
+	 */
+	private RetryContext originalRetryContext;
+
+	/**
+	 * Creates a new {@link DelegatingRetryContextCallable} with a specific
+	 * {@link RetryContext}.
+	 * @param delegate the delegate {@link DelegatingRetryContextCallable} to run with the
+	 * specified {@link RetryContext}. Cannot be null.
+	 * @param retryContext the {@link RetryContext} to establish for the delegate
+	 * {@link Callable}. Cannot be null.
+	 */
+	public DelegatingRetryContextCallable(Callable<V> delegate, RetryContext retryContext) {
+		Assert.notNull(delegate, "delegate cannot be null");
+		Assert.notNull(retryContext, "retryContext cannot be null");
+		this.delegate = delegate;
+		this.delegateRetryContext = retryContext;
+	}
+
+	/**
+	 * Creates a new {@link DelegatingRetryContextCallable} with the {@link RetryContext}
+	 * from the {@link RetrySynchronizationManager}.
+	 * @param delegate the delegate {@link Callable} to run under the current
+	 * {@link RetryContext}. Cannot be null.
+	 */
+	public DelegatingRetryContextCallable(Callable<V> delegate) {
+		this(delegate, RetrySynchronizationManager.getContext());
+	}
+
+	@Override
+	public V call() throws Exception {
+		this.originalRetryContext = RetrySynchronizationManager.getContext();
+
+		try {
+			RetrySynchronizationManager.register(delegateRetryContext);
+			return delegate.call();
+		}
+		finally {
+			RetrySynchronizationManager.register(originalRetryContext);
+			this.originalRetryContext = null;
+		}
+	}
+
+	@Override
+	public String toString() {
+		return delegate.toString();
+	}
+
+	/**
+	 * Creates a {@link DelegatingRetryContextCallable} and with the given
+	 * {@link Callable} and {@link RetryContext}, but if the retryContext is null will
+	 * defaults to the current {@link RetryContext} on the
+	 * {@link RetrySynchronizationManager}
+	 * @param delegate the delegate {@link DelegatingRetryContextCallable} to run with the
+	 * specified {@link RetryContext}. Cannot be null.
+	 * @param retryContext the {@link RetryContext} to establish for the delegate
+	 * {@link Callable}. If null, defaults to
+	 * {@link RetrySynchronizationManager#getContext()}
+	 * @return
+	 */
+	public static <V> Callable<V> create(Callable<V> delegate, RetryContext retryContext) {
+		return retryContext == null ? new DelegatingRetryContextCallable<>(delegate)
+				: new DelegatingRetryContextCallable<>(delegate, retryContext);
+	}
+
+}

src/main/java/org/springframework/retry/concurrent/DelegatingRetryContextExecutor.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright 2002-2019 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.concurrent;
+
+import org.springframework.retry.RetryContext;
+import org.springframework.retry.support.RetrySynchronizationManager;
+import org.springframework.util.Assert;
+
+import java.util.concurrent.Executor;
+
+/**
+ * An {@link Executor} which wraps each {@link Runnable} in a
+ * {@link DelegatingRetryContextRunnable}.
+ * Inspired by org.springframework.security.concurrent package.
+ *
+ * @author Rob Winch
+ * @author Aleksandr Shamukov
+ */
+public class DelegatingRetryContextExecutor extends AbstractDelegatingRetryContextSupport implements Executor {
+
+	private final Executor delegate;
+
+	/**
+	 * Creates a new {@link DelegatingRetryContextExecutor} that uses the specified
+	 * {@link RetryContext}.
+	 * @param delegateExecutor the {@link Executor} to delegate to. Cannot be null.
+	 * @param retryContext the {@link RetryContext} to use for each
+	 * {@link DelegatingRetryContextRunnable} or null to default to the current
+	 * {@link RetryContext}
+	 */
+	public DelegatingRetryContextExecutor(Executor delegateExecutor, RetryContext retryContext) {
+		super(retryContext);
+		Assert.notNull(delegateExecutor, "delegateExecutor cannot be null");
+		this.delegate = delegateExecutor;
+	}
+
+	/**
+	 * Creates a new {@link DelegatingRetryContextExecutor} that uses the current
+	 * {@link RetryContext} from the {@link RetrySynchronizationManager} at the time the
+	 * task is submitted.
+	 * @param delegate the {@link Executor} to delegate to. Cannot be null.
+	 */
+	public DelegatingRetryContextExecutor(Executor delegate) {
+		this(delegate, null);
+	}
+
+	public final void execute(Runnable task) {
+		task = wrap(task);
+		delegate.execute(task);
+	}
+
+	protected final Executor getDelegateExecutor() {
+		return delegate;
+	}
+
+}

src/main/java/org/springframework/retry/concurrent/DelegatingRetryContextExecutorService.java

@@ -0,0 +1,139 @@
+/*
+ * Copyright 2002-2019 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.concurrent;
+
+import org.springframework.retry.RetryContext;
+import org.springframework.retry.support.RetrySynchronizationManager;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+/**
+ * An {@link ExecutorService} which wraps each {@link Runnable} in a
+ * {@link DelegatingRetryContextRunnable} and each {@link Callable} in a
+ * {@link DelegatingRetryContextCallable}.
+ * Inspired by org.springframework.security.concurrent package.
+ *
+ * @author Rob Winch
+ * @author Aleksandr Shamukov
+ */
+public class DelegatingRetryContextExecutorService extends DelegatingRetryContextExecutor implements ExecutorService {
+
+	/**
+	 * Creates a new {@link DelegatingRetryContextExecutorService} that uses the specified
+	 * {@link RetryContext}.
+	 * @param delegateExecutorService the {@link ExecutorService} to delegate to. Cannot
+	 * be null.
+	 * @param retryContext the {@link RetryContext} to use for each
+	 * {@link DelegatingRetryContextRunnable} and each
+	 * {@link DelegatingRetryContextCallable}.
+	 */
+	public DelegatingRetryContextExecutorService(ExecutorService delegateExecutorService, RetryContext retryContext) {
+		super(delegateExecutorService, retryContext);
+	}
+
+	/**
+	 * Creates a new {@link DelegatingRetryContextExecutorService} that uses the current
+	 * {@link RetryContext} from the {@link RetrySynchronizationManager}.
+	 * @param delegate the {@link ExecutorService} to delegate to. Cannot be null.
+	 */
+	public DelegatingRetryContextExecutorService(ExecutorService delegate) {
+		this(delegate, null);
+	}
+
+	public final void shutdown() {
+		getDelegate().shutdown();
+	}
+
+	public final List<Runnable> shutdownNow() {
+		return getDelegate().shutdownNow();
+	}
+
+	public final boolean isShutdown() {
+		return getDelegate().isShutdown();
+	}
+
+	public final boolean isTerminated() {
+		return getDelegate().isTerminated();
+	}
+
+	public final boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
+		return getDelegate().awaitTermination(timeout, unit);
+	}
+
+	public final <T> Future<T> submit(Callable<T> task) {
+		task = wrap(task);
+		return getDelegate().submit(task);
+	}
+
+	public final <T> Future<T> submit(Runnable task, T result) {
+		task = wrap(task);
+		return getDelegate().submit(task, result);
+	}
+
+	public final Future<?> submit(Runnable task) {
+		task = wrap(task);
+		return getDelegate().submit(task);
+	}
+
+	@SuppressWarnings({ "rawtypes", "unchecked" })
+	public final List invokeAll(Collection tasks) throws InterruptedException {
+		tasks = createTasks(tasks);
+		return getDelegate().invokeAll(tasks);
+	}
+
+	@SuppressWarnings({ "rawtypes", "unchecked" })
+	public final List invokeAll(Collection tasks, long timeout, TimeUnit unit) throws InterruptedException {
+		tasks = createTasks(tasks);
+		return getDelegate().invokeAll(tasks, timeout, unit);
+	}
+
+	@SuppressWarnings({ "rawtypes", "unchecked" })
+	public final Object invokeAny(Collection tasks) throws InterruptedException, ExecutionException {
+		tasks = createTasks(tasks);
+		return getDelegate().invokeAny(tasks);
+	}
+
+	@SuppressWarnings({ "rawtypes", "unchecked" })
+	public final Object invokeAny(Collection tasks, long timeout, TimeUnit unit)
+			throws InterruptedException, ExecutionException, TimeoutException {
+		tasks = createTasks(tasks);
+		return getDelegate().invokeAny(tasks, timeout, unit);
+	}
+
+	private <T> Collection<Callable<T>> createTasks(Collection<Callable<T>> tasks) {
+		if (tasks == null) {
+			return null;
+		}
+		List<Callable<T>> results = new ArrayList<>(tasks.size());
+		for (Callable<T> task : tasks) {
+			results.add(wrap(task));
+		}
+		return results;
+	}
+
+	private ExecutorService getDelegate() {
+		return (ExecutorService) getDelegateExecutor();
+	}
+
+}