diff --git a/spring-batch-core/.classpath b/spring-batch-core/.classpath index 06b858ff6..6886233c2 100644 --- a/spring-batch-core/.classpath +++ b/spring-batch-core/.classpath @@ -1,10 +1,11 @@ + - - - - - - - - + + + + + + + + diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/scope/context/StepContextRepeatCallback.java b/spring-batch-core/src/main/java/org/springframework/batch/core/scope/context/StepContextRepeatCallback.java index 3581fc0c3..a8d2adbbe 100644 --- a/spring-batch-core/src/main/java/org/springframework/batch/core/scope/context/StepContextRepeatCallback.java +++ b/spring-batch-core/src/main/java/org/springframework/batch/core/scope/context/StepContextRepeatCallback.java @@ -26,6 +26,7 @@ import org.springframework.batch.core.StepExecution; import org.springframework.batch.repeat.RepeatCallback; import org.springframework.batch.repeat.RepeatContext; import org.springframework.batch.repeat.RepeatStatus; +import org.springframework.util.ObjectUtils; /** * Convenient base class for clients who need to do something in a repeat @@ -68,6 +69,7 @@ public abstract class StepContextRepeatCallback implements RepeatCallback { // The StepContext has to be the same for all chunks, // otherwise step-scoped beans will be re-initialised for each chunk. StepContext stepContext = StepSynchronizationManager.register(stepExecution); + logger.debug("Preparing chunk execution for StepContext: "+ObjectUtils.identityToString(stepContext)); ChunkContext chunkContext = attributeQueue.poll(); if (chunkContext == null) { diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/scope/context/StepSynchronizationManager.java b/spring-batch-core/src/main/java/org/springframework/batch/core/scope/context/StepSynchronizationManager.java index 7f3a4a59e..07eecb08d 100644 --- a/spring-batch-core/src/main/java/org/springframework/batch/core/scope/context/StepSynchronizationManager.java +++ b/spring-batch-core/src/main/java/org/springframework/batch/core/scope/context/StepSynchronizationManager.java @@ -15,7 +15,10 @@ */ package org.springframework.batch.core.scope.context; +import java.util.HashMap; +import java.util.Map; import java.util.Stack; +import java.util.concurrent.atomic.AtomicInteger; import org.springframework.batch.core.Step; import org.springframework.batch.core.StepExecution; @@ -32,11 +35,32 @@ import org.springframework.batch.core.StepExecution; */ public class StepSynchronizationManager { - /** - * Don't use InheritableThreadLocal because there are side effects if a step - * is trying to run multiple child steps (e.g. with partitioning). + /* + * We have to deal with single and multi-threaded execution, with a single + * and with multiple step execution instances. That's 2x2 = 4 scenarios. */ - private static final ThreadLocal> contextHolder = new ThreadLocal>(); + + /** + * Storage for the current step execution; has to be ThreadLocal because it + * is needed to locate a StepContext in components that are not part of a + * Step (like when re-hydrating a scoped proxy). Doesn't use + * InheritableThreadLocal because there are side effects if a step is trying + * to run multiple child steps (e.g. with partitioning). The Stack is used + * to cover the single threaded case, so that the API is the same as + * multi-threaded. + */ + private static final ThreadLocal> executionHolder = new ThreadLocal>(); + + /** + * Reference counter for each step execution: how many threads are using the + * same one? + */ + private static final Map counts = new HashMap(); + + /** + * Simple map from a running step execution to the associated context. + */ + private static final Map contexts = new HashMap(); /** * Getter for the current context if there is one, otherwise returns null. @@ -45,17 +69,16 @@ public class StepSynchronizationManager { * has not been registered for this thread). */ public static StepContext getContext() { - Stack current = getCurrent(); - if (current.isEmpty()) { + if (getCurrent().isEmpty()) { return null; } - return current.peek(); + return contexts.get(getCurrent().peek()); } /** - * Method for registering a context with the current thread - always put a - * matching {@link #close()} call in a finally block to ensure that the - * correct context is available in the enclosing block. + * Register a context with the current thread - always put a matching + * {@link #close()} call in a finally block to ensure that the correct + * context is available in the enclosing block. * * @param stepExecution the step context to register * @return a new {@link StepContext} or the current one if it has the same @@ -65,19 +88,13 @@ public class StepSynchronizationManager { if (stepExecution == null) { return null; } - StepContext current = getContext(); - StepContext context; - if (current != null && current.getStepExecution().equals(stepExecution)) { - /* - * If the new context has the same step execution we don't want a - * new set of attributes, otherwise auto proxied beans get created - * twice for the same execution. - */ - context = current; - } else { + getCurrent().push(stepExecution); + StepContext context = contexts.get(stepExecution); + if (context == null) { context = new StepContext(stepExecution); + contexts.put(stepExecution, context); } - getCurrent().push(context); + increment(); return context; } @@ -94,21 +111,43 @@ public class StepSynchronizationManager { if (oldSession == null) { return; } - getCurrent().pop(); + decrement(); } - private static Stack getCurrent() { - if (contextHolder.get() == null) { - contextHolder.set(new Stack()); + private static void decrement() { + StepExecution current = getCurrent().pop(); + if (current != null) { + int remaining = counts.get(current).decrementAndGet(); + if (remaining <= 0) { + contexts.remove(current); + } } - return contextHolder.get(); + } + + private static void increment() { + StepExecution current = getCurrent().peek(); + if (current != null) { + AtomicInteger count = counts.get(current); + if (count == null) { + count = new AtomicInteger(); + counts.put(current, count); + } + count.incrementAndGet(); + } + } + + private static Stack getCurrent() { + if (executionHolder.get() == null) { + executionHolder.set(new Stack()); + } + return executionHolder.get(); } /** - * A "deep" close operation. Call this instead of {@link #close()} if the - * step execution for the current context is ending. Delegates to - * {@link StepContext#close()} and then ensures that {@link #close()} is - * also called in a finally block. + * A convenient "deep" close operation. Call this instead of + * {@link #close()} if the step execution for the current context is ending. + * Delegates to {@link StepContext#close()} and then ensures that + * {@link #close()} is also called in a finally block. */ public static void release() { StepContext context = getContext(); diff --git a/spring-batch-core/src/test/java/org/springframework/batch/core/scope/AsyncStepScopeIntegrationTests.java b/spring-batch-core/src/test/java/org/springframework/batch/core/scope/AsyncStepScopeIntegrationTests.java index 73ec96ca0..40ea70560 100644 --- a/spring-batch-core/src/test/java/org/springframework/batch/core/scope/AsyncStepScopeIntegrationTests.java +++ b/spring-batch-core/src/test/java/org/springframework/batch/core/scope/AsyncStepScopeIntegrationTests.java @@ -51,6 +51,7 @@ public class AsyncStepScopeIntegrationTests implements BeanFactoryAware { @Before public void countBeans() { + StepSynchronizationManager.release(); beanCount = beanFactory.getBeanDefinitionCount(); } @@ -71,7 +72,7 @@ public class AsyncStepScopeIntegrationTests implements BeanFactoryAware { } @Test - public void testGetMultiple() throws Exception { + public void testGetMultipleInMultipleThreads() throws Exception { List> tasks = new ArrayList>(); @@ -105,4 +106,44 @@ public class AsyncStepScopeIntegrationTests implements BeanFactoryAware { } + @Test + public void testGetSameInMultipleThreads() throws Exception { + + List> tasks = new ArrayList>(); + final StepExecution stepExecution = new StepExecution("foo", new JobExecution(0L), 123L); + ExecutionContext executionContext = stepExecution.getExecutionContext(); + executionContext.put("foo", "foo"); + StepSynchronizationManager.register(stepExecution); + assertEquals("foo", simple.getName()); + + for (int i = 0; i < 12; i++) { + final String value = "foo"+i; + FutureTask task = new FutureTask(new Callable() { + public String call() throws Exception { + ExecutionContext executionContext = stepExecution.getExecutionContext(); + executionContext.put("foo", value); + StepContext context = StepSynchronizationManager.register(stepExecution); + logger.debug("Registered: "+context.getStepExecutionContext()); + try { + return simple.getName(); + } + finally { + StepSynchronizationManager.close(); + } + } + }); + tasks.add(task); + taskExecutor.execute(task); + } + + StepSynchronizationManager.close(); + + int i = 0; + for (FutureTask task : tasks) { + assertEquals("foo", task.get()); + i++; + } + + } + } diff --git a/spring-batch-core/src/test/java/org/springframework/batch/core/scope/context/StepSynchronizationManagerTests.java b/spring-batch-core/src/test/java/org/springframework/batch/core/scope/context/StepSynchronizationManagerTests.java index 7ca3d639e..fb55942dd 100644 --- a/spring-batch-core/src/test/java/org/springframework/batch/core/scope/context/StepSynchronizationManagerTests.java +++ b/spring-batch-core/src/test/java/org/springframework/batch/core/scope/context/StepSynchronizationManagerTests.java @@ -6,14 +6,17 @@ import static org.junit.Assert.assertNull; import java.util.ArrayList; import java.util.List; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.FutureTask; +import java.util.concurrent.TimeUnit; import org.junit.After; import org.junit.Before; import org.junit.Test; import org.springframework.batch.core.JobExecution; import org.springframework.batch.core.StepExecution; -import org.springframework.batch.core.scope.context.StepContext; -import org.springframework.batch.core.scope.context.StepSynchronizationManager; public class StepSynchronizationManagerTests { @@ -48,6 +51,30 @@ public class StepSynchronizationManagerTests { assertEquals(0, list.size()); } + @Test + public void testMultithreaded() throws Exception { + StepContext context = StepSynchronizationManager.register(stepExecution); + ExecutorService executorService = Executors.newFixedThreadPool(2); + FutureTask task = new FutureTask(new Callable() { + public StepContext call() throws Exception { + try { + StepSynchronizationManager.register(stepExecution); + StepContext context = StepSynchronizationManager.getContext(); + context.setAttribute("foo", "bar"); + return context; + } + finally { + StepSynchronizationManager.close(); + } + } + }); + executorService.execute(task); + executorService.awaitTermination(1, TimeUnit.SECONDS); + assertEquals(context.attributeNames().length, task.get().attributeNames().length); + StepSynchronizationManager.close(); + assertNull(StepSynchronizationManager.getContext()); + } + @Test public void testRelease() { StepContext context = StepSynchronizationManager.register(stepExecution);