Fixed StepScope not being handled correctly - creating too many auto-proxies for the same stepExecution.

spring-batch-core/src/main/java/org/springframework/batch/core/scope/StepContextRepeatCallback.java

@@ -35,13 +35,13 @@ public abstract class StepContextRepeatCallback implements RepeatCallback {
 
 	private final Queue<ChunkContext> attributeQueue = new LinkedBlockingQueue<ChunkContext>();
 
-	private final StepContext stepContext;
+	private final StepExecution stepExecution;
 
 	/**
 	 * @param stepExecution
 	 */
 	public StepContextRepeatCallback(StepExecution stepExecution) {
-		this.stepContext = new StepContext(stepExecution);
+		this.stepExecution = stepExecution;
 	}
 
 	/**
@@ -54,15 +54,16 @@ public abstract class StepContextRepeatCallback implements RepeatCallback {
 	 */
 	public RepeatStatus doInIteration(RepeatContext context) throws Exception {
 
+		// 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);
+
 		ChunkContext chunkContext = attributeQueue.poll();
 		if (chunkContext == null) {
 			chunkContext = new ChunkContext();
 		}
-		int start = stepContext.attributeNames().length;
+		int start = chunkContext.attributeNames().length;
 
-		// The StepContext has to be the same for all chunks,
-		// otherwise step-scoped beans will be re-initialised for each chunk.
-		StepSynchronizationManager.register(stepContext);
 		try {
 			return doInStepContext(context, stepContext);
 		}

spring-batch-core/src/main/java/org/springframework/batch/core/scope/StepScopeManager.java

@@ -20,7 +20,7 @@ public class StepScopeManager {
 
 	@Around("execution(void org.springframework.batch.core.Step+.execute(*)) && target(step) && args(stepExecution)")
 	public void execute(Step step, StepExecution stepExecution) throws JobInterruptedException {
-		StepSynchronizationManager.register(new StepContext(stepExecution));
+		StepSynchronizationManager.register(stepExecution);
 		try {
 			step.execute(stepExecution);
 		}

spring-batch-core/src/main/java/org/springframework/batch/core/scope/StepSynchronizationManager.java

@@ -18,6 +18,7 @@ package org.springframework.batch.core.scope;
 import java.util.Stack;
 
 import org.springframework.batch.core.Step;
+import org.springframework.batch.core.StepExecution;
 
 /**
  * Central convenience class for framework use in managing the step scope
@@ -56,20 +57,35 @@ public class StepSynchronizationManager {
 	 * matching {@link #close()} call in a finally block to ensure that the
 	 * correct context is available in the enclosing block.
 	 * 
-	 * @param context the step context to register
+	 * @param stepExecution the step context to register
+	 * @return a new {@link StepContext} or the current one if it has the same
+	 * {@link StepExecution}
 	 */
-	public static void register(StepContext context) {
-		if (context == null) {
-			return;
+	public static StepContext register(StepExecution stepExecution) {
+		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 {
+			context = new StepContext(stepExecution);
 		}
 		getCurrent().push(context);
+		return context;
 	}
 
 	/**
 	 * Method for de-registering the current context - should always and only be
-	 * used by in conjunction with a matching {@link #register(StepContext)} to
-	 * ensure that {@link #getContext()} always returns the correct value. Does
-	 * not call {@link StepContext#close()} - that is left up to the caller
+	 * used by in conjunction with a matching {@link #register(StepExecution)}
+	 * to ensure that {@link #getContext()} always returns the correct value.
+	 * Does not call {@link StepContext#close()} - that is left up to the caller
 	 * because he has a reference to the context (having registered it) and only
 	 * he has knowledge of when the step actually ended.
 	 */
@@ -97,10 +113,11 @@ public class StepSynchronizationManager {
 	public static void release() {
 		StepContext context = getContext();
 		try {
-			if (context!=null) {
+			if (context != null) {
 				context.close();
 			}
-		} finally {
+		}
+		finally {
 			close();
 		}
 	}

spring-batch-core/src/main/java/org/springframework/batch/core/step/AbstractStep.java

@@ -32,7 +32,6 @@ import org.springframework.batch.core.launch.NoSuchJobException;
 import org.springframework.batch.core.launch.support.ExitCodeMapper;
 import org.springframework.batch.core.listener.CompositeStepExecutionListener;
 import org.springframework.batch.core.repository.JobRepository;
-import org.springframework.batch.core.scope.StepContext;
 import org.springframework.batch.core.scope.StepSynchronizationManager;
 import org.springframework.batch.item.ExecutionContext;
 import org.springframework.beans.factory.BeanNameAware;
@@ -47,8 +46,7 @@ import org.springframework.util.Assert;
  * @author Ben Hale
  * @author Robert Kasanicky
  */
-public abstract class AbstractStep implements Step, InitializingBean,
-		BeanNameAware {
+public abstract class AbstractStep implements Step, InitializingBean, BeanNameAware {
 
 	private static final Log logger = LogFactory.getLog(AbstractStep.class);
 
@@ -109,8 +107,7 @@ public abstract class AbstractStep implements Step, InitializingBean,
 	/**
 	 * Public setter for the startLimit.
 	 * 
-	 * @param startLimit
-	 *            the startLimit to set
+	 * @param startLimit the startLimit to set
 	 */
 	public void setStartLimit(int startLimit) {
 		this.startLimit = startLimit;
@@ -124,8 +121,7 @@ public abstract class AbstractStep implements Step, InitializingBean,
 	 * Public setter for flag that determines whether the step should start
 	 * again if it is already complete. Defaults to false.
 	 * 
-	 * @param allowStartIfComplete
-	 *            the value of the flag to set
+	 * @param allowStartIfComplete the value of the flag to set
 	 */
 	public void setAllowStartIfComplete(boolean allowStartIfComplete) {
 		this.allowStartIfComplete = allowStartIfComplete;
@@ -145,20 +141,17 @@ public abstract class AbstractStep implements Step, InitializingBean,
 	 * should set the {@link ExitStatus} on the {@link StepExecution} before
 	 * returning.
 	 * 
-	 * @param stepExecution
-	 *            the current step context
+	 * @param stepExecution the current step context
 	 * @throws Exception
 	 */
-	protected abstract void doExecute(StepExecution stepExecution)
-			throws Exception;
+	protected abstract void doExecute(StepExecution stepExecution) throws Exception;
 
 	/**
 	 * Extension point for subclasses to provide callbacks to their
 	 * collaborators at the beginning of a step, to open or acquire resources.
 	 * Does nothing by default.
 	 * 
-	 * @param ctx
-	 *            the {@link ExecutionContext} to use
+	 * @param ctx the {@link ExecutionContext} to use
 	 * @throws Exception
 	 */
 	protected void open(ExecutionContext ctx) throws Exception {
@@ -169,8 +162,7 @@ public abstract class AbstractStep implements Step, InitializingBean,
 	 * collaborators at the end of a step (right at the end of the finally
 	 * block), to close or release resources. Does nothing by default.
 	 * 
-	 * @param ctx
-	 *            the {@link ExecutionContext} to use
+	 * @param ctx the {@link ExecutionContext} to use
 	 * @throws Exception
 	 */
 	protected void close(ExecutionContext ctx) throws Exception {
@@ -182,8 +174,8 @@ public abstract class AbstractStep implements Step, InitializingBean,
 	 * logic ({@link #doExecute(StepExecution)}) and resource closing (
 	 * {@link #close(ExecutionContext)}).
 	 */
-	public final void execute(StepExecution stepExecution)
-			throws JobInterruptedException, UnexpectedJobExecutionException {
+	public final void execute(StepExecution stepExecution) throws JobInterruptedException,
+			UnexpectedJobExecutionException {
 		stepExecution.setStartTime(new Date());
 		stepExecution.setStatus(BatchStatus.STARTED);
 		getJobRepository().update(stepExecution);
@@ -191,8 +183,7 @@ public abstract class AbstractStep implements Step, InitializingBean,
 		ExitStatus exitStatus = ExitStatus.FAILED;
 		Exception commitException = null;
 
-		StepContext stepContext = new StepContext(stepExecution);
-		StepSynchronizationManager.register(stepContext);
+		StepSynchronizationManager.register(stepExecution);
 
 		try {
 			getCompositeListener().beforeStep(stepExecution);
@@ -212,32 +203,34 @@ public abstract class AbstractStep implements Step, InitializingBean,
 			try {
 				getJobRepository().update(stepExecution);
 				getJobRepository().updateExecutionContext(stepExecution);
-			} catch (Exception e) {
+			}
+			catch (Exception e) {
 				commitException = e;
 				exitStatus = exitStatus.and(ExitStatus.UNKNOWN);
 			}
 
-		} catch (Throwable e) {
+		}
+		catch (Throwable e) {
 
-			logger.error("Encountered an error executing the step: "
-					+ e.getClass() + ": " + e.getMessage(), e);
+			logger.error("Encountered an error executing the step: " + e.getClass() + ": " + e.getMessage(), e);
 			stepExecution.setStatus(determineBatchStatus(e));
 			exitStatus = getDefaultExitStatusForFailure(e);
 			stepExecution.addFailureException(e);
 
 			try {
 				getJobRepository().updateExecutionContext(stepExecution);
-			} catch (Exception ex) {
-				logger.error(
-						"Encountered an error on listener error callback.", ex);
+			}
+			catch (Exception ex) {
+				logger.error("Encountered an error on listener error callback.", ex);
 				stepExecution.addFailureException(ex);
 			}
-		} finally {
+		}
+		finally {
 
 			try {
-				exitStatus = exitStatus.and(getCompositeListener().afterStep(
-						stepExecution));
-			} catch (Exception e) {
+				exitStatus = exitStatus.and(getCompositeListener().afterStep(stepExecution));
+			}
+			catch (Exception e) {
 				logger.error("Exception in afterStep callback", e);
 			}
 
@@ -246,22 +239,21 @@ public abstract class AbstractStep implements Step, InitializingBean,
 
 			try {
 				getJobRepository().update(stepExecution);
-			} catch (Exception e) {
+			}
+			catch (Exception e) {
 				if (commitException == null) {
 					commitException = e;
-				} else {
-					logger
-							.error(
-									"Exception while updating step execution after commit exception",
-									e);
+				}
+				else {
+					logger.error("Exception while updating step execution after commit exception", e);
 				}
 			}
 
 			try {
 				close(stepExecution.getExecutionContext());
-			} catch (Exception e) {
-				logger.error(
-						"Exception while closing step execution resources", e);
+			}
+			catch (Exception e) {
+				logger.error("Exception while closing step execution resources", e);
 				stepExecution.addFailureException(e);
 			}
 
@@ -269,11 +261,8 @@ public abstract class AbstractStep implements Step, InitializingBean,
 
 			if (commitException != null) {
 				stepExecution.setStatus(BatchStatus.UNKNOWN);
-				logger
-						.error(
-								"Encountered an error saving batch meta data."
-										+ "This job is now in an unknown state and should not be restarted.",
-								commitException);
+				logger.error("Encountered an error saving batch meta data."
+						+ "This job is now in an unknown state and should not be restarted.", commitException);
 				stepExecution.addFailureException(commitException);
 			}
 		}
@@ -285,10 +274,11 @@ public abstract class AbstractStep implements Step, InitializingBean,
 	private static BatchStatus determineBatchStatus(Throwable e) {
 		if (e instanceof FatalException) {
 			return BatchStatus.UNKNOWN;
-		} else if (e instanceof JobInterruptedException
-				|| e.getCause() instanceof JobInterruptedException) {
+		}
+		else if (e instanceof JobInterruptedException || e.getCause() instanceof JobInterruptedException) {
 			return BatchStatus.STOPPED;
-		} else {
+		}
+		else {
 			return BatchStatus.FAILED;
 		}
 	}
@@ -297,8 +287,7 @@ public abstract class AbstractStep implements Step, InitializingBean,
 	 * Register a step listener for callbacks at the appropriate stages in a
 	 * step execution.
 	 * 
-	 * @param listener
-	 *            a {@link StepExecutionListener}
+	 * @param listener a {@link StepExecutionListener}
 	 */
 	public void registerStepExecutionListener(StepExecutionListener listener) {
 		this.listener.register(listener);
@@ -307,8 +296,7 @@ public abstract class AbstractStep implements Step, InitializingBean,
 	/**
 	 * Register each of the objects as listeners.
 	 * 
-	 * @param listeners
-	 *            an array of listener objects of known types.
+	 * @param listeners an array of listener objects of known types.
 	 */
 	public void setStepExecutionListeners(StepExecutionListener[] listeners) {
 		for (int i = 0; i < listeners.length; i++) {
@@ -326,8 +314,7 @@ public abstract class AbstractStep implements Step, InitializingBean,
 	/**
 	 * Public setter for {@link JobRepository}.
 	 * 
-	 * @param jobRepository
-	 *            is a mandatory dependence (no default).
+	 * @param jobRepository is a mandatory dependence (no default).
 	 */
 	public void setJobRepository(JobRepository jobRepository) {
 		this.jobRepository = jobRepository;
@@ -341,21 +328,18 @@ public abstract class AbstractStep implements Step, InitializingBean,
 	 * Default mapping from throwable to {@link ExitStatus}. Clients can modify
 	 * the exit code using a {@link StepExecutionListener}.
 	 * 
-	 * @param ex
-	 *            the cause of the failure
+	 * @param ex the cause of the failure
 	 * @return an {@link ExitStatus}
 	 */
 	private ExitStatus getDefaultExitStatusForFailure(Throwable ex) {
 		ExitStatus exitStatus;
-		if (ex instanceof JobInterruptedException
-				|| ex.getCause() instanceof JobInterruptedException) {
-			exitStatus = ExitStatus.INTERRUPTED
-					.addExitDescription(JobInterruptedException.class.getName());
-		} else if (ex instanceof NoSuchJobException
-				|| ex.getCause() instanceof NoSuchJobException) {
-			exitStatus = new ExitStatus(ExitCodeMapper.NO_SUCH_JOB, ex
-					.getClass().getName());
-		} else {
+		if (ex instanceof JobInterruptedException || ex.getCause() instanceof JobInterruptedException) {
+			exitStatus = ExitStatus.INTERRUPTED.addExitDescription(JobInterruptedException.class.getName());
+		}
+		else if (ex instanceof NoSuchJobException || ex.getCause() instanceof NoSuchJobException) {
+			exitStatus = new ExitStatus(ExitCodeMapper.NO_SUCH_JOB, ex.getClass().getName());
+		}
+		else {
 			StringWriter writer = new StringWriter();
 			ex.printStackTrace(new PrintWriter(writer));
 			String message = writer.toString();

spring-batch-core/src/test/java/org/springframework/batch/core/partition/RestartIntegrationTests.java

@@ -87,7 +87,6 @@ public class RestartIntegrationTests {
 		// Two attempts
 		assertEquals(2, afterMaster-beforeMaster);
 		// One failure and two successes
-		System.out.println(afterPartition);
 		assertEquals(3, afterPartition-beforePartition);
 
 	}

spring-batch-core/src/test/java/org/springframework/batch/core/scope/StepContextRepeatCallbackTests.java

@@ -19,6 +19,7 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
+import org.junit.After;
 import org.junit.Test;
 import org.springframework.batch.core.ExitStatus;
 import org.springframework.batch.core.JobExecution;
@@ -35,6 +36,11 @@ public class StepContextRepeatCallbackTests {
 	private StepExecution stepExecution = new StepExecution("foo", new JobExecution(0L), 123L);
 	private boolean addedAttribute = false;
 	private boolean removedAttribute = false;
+	
+	@After
+	public void cleanUpStepContext() {
+		StepSynchronizationManager.close();
+	}
 
 	@Test
 	public void testDoInIteration() throws Exception {
@@ -51,6 +57,7 @@ public class StepContextRepeatCallbackTests {
 
 	@Test
 	public void testUnfinishedWork() throws Exception {
+		StepSynchronizationManager.register(stepExecution);
 		StepContextRepeatCallback callback = new StepContextRepeatCallback(stepExecution) {
 			@Override
 			public RepeatStatus doInStepContext(RepeatContext context, StepContext stepContext) throws Exception {

spring-batch-core/src/test/java/org/springframework/batch/core/scope/StepScopeIntegrationTests.java

@@ -38,6 +38,7 @@ public class StepScopeIntegrationTests {
 	@Before
 	@After
 	public void start() {
+		StepSynchronizationManager.close();
 		TestStep.reset();
 	}
 		

spring-batch-core/src/test/java/org/springframework/batch/core/scope/StepScopeTests.java

@@ -18,6 +18,7 @@ package org.springframework.batch.core.scope;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -41,11 +42,13 @@ public class StepScopeTests {
 
 	private StepScope scope = new StepScope();
 
-	private StepContext context = new StepContext(new StepExecution("foo", new JobExecution(0L)));
+	private StepExecution stepExecution = new StepExecution("foo", new JobExecution(0L));
+
+	private StepContext context;
 
 	@Before
 	public void setUp() throws Exception {
-		StepSynchronizationManager.register(context);
+		context = StepSynchronizationManager.register(stepExecution);
 	}
 
 	@After
@@ -98,8 +101,7 @@ public class StepScopeTests {
 	@Test
 	public void testGetWithSomethingAlreadyInParentContext() {
 		context.setAttribute("foo", "bar");
-		StepContext context = new StepContext(new StepExecution("bar", new JobExecution(0L)));
-		StepSynchronizationManager.register(context);
+		StepContext context = StepSynchronizationManager.register(new StepExecution("bar", new JobExecution(0L)));
 		Object value = scope.get("foo", new ObjectFactory() {
 			public Object getObject() throws BeansException {
 				return "spam";
@@ -111,6 +113,13 @@ public class StepScopeTests {
 		assertEquals("bar", scope.get("foo", null));
 	}
 
+	@Test
+	public void testParentContextWithSameStepExecution() {
+		context.setAttribute("foo", "bar");
+		StepContext other = StepSynchronizationManager.register(stepExecution);
+		assertSame(other, context);
+	}
+
 	@Test
 	public void testGetConversationId() {
 		String id = scope.getConversationId();

spring-batch-core/src/test/java/org/springframework/batch/core/scope/StepSynchronizationManagerTests.java

@@ -28,20 +28,19 @@ public class StepSynchronizationManagerTests {
 	@Test
 	public void testGetContext() {
 		assertNull(StepSynchronizationManager.getContext());
-		StepSynchronizationManager.register(new StepContext(stepExecution));
+		StepSynchronizationManager.register(stepExecution);
 		assertNotNull(StepSynchronizationManager.getContext());
 	}
 
 	@Test
 	public void testClose() {
-		StepContext context = new StepContext(stepExecution);
 		final List<String> list = new ArrayList<String>();
+		StepContext context = StepSynchronizationManager.register(stepExecution);
 		context.registerDestructionCallback("foo", new Runnable() {
 			public void run() {
 				list.add("foo");
 			}
 		});
-		StepSynchronizationManager.register(context);
 		StepSynchronizationManager.close();
 		assertNull(StepSynchronizationManager.getContext());
 		assertEquals(0, list.size());
@@ -49,14 +48,13 @@ public class StepSynchronizationManagerTests {
 
 	@Test
 	public void testRelease() {
-		StepContext context = new StepContext(stepExecution);
+		StepContext context = StepSynchronizationManager.register(stepExecution);
 		final List<String> list = new ArrayList<String>();
 		context.registerDestructionCallback("foo", new Runnable() {
 			public void run() {
 				list.add("foo");
 			}
 		});
-		StepSynchronizationManager.register(context);
 		// On release we expect the destruction callbacks to be called
 		StepSynchronizationManager.release();
 		assertNull(StepSynchronizationManager.getContext());
@@ -72,9 +70,8 @@ public class StepSynchronizationManagerTests {
 
 	@Test
 	public void testRegisterTwice() {
-		StepContext context = new StepContext(stepExecution);
-		StepSynchronizationManager.register(context);
-		StepSynchronizationManager.register(context);
+		StepSynchronizationManager.register(stepExecution);
+		StepSynchronizationManager.register(stepExecution);
 		StepSynchronizationManager.close();
 		// if someone registers you have to assume they are going to close, so
 		// the last thing you want is for the close to remove another context

spring-batch-core/src/test/java/org/springframework/batch/core/scope/TestStep.java

@@ -7,13 +7,13 @@ import org.springframework.batch.core.StepExecution;
 public class TestStep implements Step {
 
 	private static StepContext context;
-	
+
 	private Collaborator collaborator;
-	
+
 	public void setCollaborator(Collaborator collaborator) {
 		this.collaborator = collaborator;
 	}
-	
+
 	public static StepContext getContext() {
 		return context;
 	}
@@ -24,7 +24,9 @@ public class TestStep implements Step {
 
 	public void execute(StepExecution stepExecution) throws JobInterruptedException {
 		context = StepSynchronizationManager.getContext();
-		context.setAttribute("collaborator", collaborator.getName());
+		if (context != null) {
+			context.setAttribute("collaborator", collaborator.getName());
+		}
 	}
 
 	public String getName() {

spring-batch-core/src/test/resources/foo.sql

@@ -1,3 +1,5 @@
+DROP TABLE T_FOOS;
+
 CREATE TABLE T_FOOS (
 	ID BIGINT NOT NULL,
 	NAME VARCHAR(45),