OPEN - issue BATCH-368: StepExecution attributes can overflow and cause spurious OptimisticLockingException

http://jira.springframework.org/browse/BATCH-368

Split the save of a StepExecution into two parts at the JobRepository level, that way we can see if the failure was in saving context or just the step execution.
Lots of exception handling and tests in step implementations.

spring-batch-execution/src/main/java/org/springframework/batch/execution/job/AbstractJob.java

@@ -20,9 +20,7 @@ import java.util.ArrayList;
 import java.util.List;
 
 import org.springframework.batch.core.domain.Job;
-import org.springframework.batch.core.domain.JobExecution;
 import org.springframework.batch.core.domain.Step;
-import org.springframework.batch.io.exception.BatchCriticalException;
 import org.springframework.beans.factory.BeanNameAware;
 import org.springframework.util.ClassUtils;
 

spring-batch-execution/src/main/java/org/springframework/batch/execution/repository/SimpleJobRepository.java

@@ -108,8 +108,8 @@ public class SimpleJobRepository implements JobRepository {
 	 * <ul>
 	 * <li>If there are none, or the {@link Job} is marked restartable, then we
 	 * create a new {@link JobInstance}</li>
-	 * <li>If there is more than one and the {@link Job} is not marked
-	 * as restartable, it is an error. This could be caused by a job whose
+	 * <li>If there is more than one and the {@link Job} is not marked as
+	 * restartable, it is an error. This could be caused by a job whose
 	 * restartable flag has changed to be more strict (true not false)
 	 * <em>after</em> it has been executed at least once.</li>
 	 * <li>If there is precisely one existing {@link JobInstance} then we check
@@ -136,7 +136,7 @@ public class SimpleJobRepository implements JobRepository {
 	 * JobInstance.getJobExecutionCount() is greater than Job.getStartLimit()
 	 * @throws JobExecutionAlreadyRunningException if a job execution is found
 	 * for the given {@link JobIdentifier} that is already running
-	 * @throws CannotRestartJobInstanceException 
+	 * @throws CannotRestartJobInstanceException
 	 * 
 	 */
 	public JobExecution createJobExecution(Job job, JobParameters jobParameters)
@@ -163,7 +163,7 @@ public class SimpleJobRepository implements JobRepository {
 			if (!job.isRestartable()) {
 				throw new BatchRestartException("JobInstance already exists and is not restartable");
 			}
-			
+
 			jobInstance.setJobExecutionCount(jobExecutionDao.getJobExecutionCount(jobInstance));
 			if (jobInstance.getJobExecutionCount() > job.getStartLimit()) {
 				throw new BatchRestartException("Restart Max exceeded for Job: " + jobInstance.toString());
@@ -252,15 +252,22 @@ public class SimpleJobRepository implements JobRepository {
 				jobExecutionDao.saveJobExecution(jobExecution);
 			}
 			stepExecutionDao.saveStepExecution(stepExecution);
-			stepExecutionDao.saveExecutionContext(stepExecution);
 		}
 		else {
 			// existing execution, update
 			stepExecutionDao.updateStepExecution(stepExecution);
-			stepExecutionDao.updateExecutionContext(stepExecution);
 		}
 	}
 
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.batch.core.repository.JobRepository#saveOrUpdateExecutionContext(org.springframework.batch.core.domain.StepExecution)
+	 */
+	public void saveOrUpdateExecutionContext(StepExecution stepExecution) {
+		saveOrUpdate(stepExecution);
+		stepExecutionDao.saveOrUpdateExecutionContext(stepExecution);
+	}
+
 	/**
 	 * @return the last execution of the step within given job instance
 	 */

spring-batch-execution/src/main/java/org/springframework/batch/execution/repository/dao/JdbcStepExecutionDao.java

@@ -120,38 +120,6 @@ public class JdbcStepExecutionDao extends AbstractJdbcBatchMetadataDao implement
 		return executionContext;
 	}
 
-	/**
-	 * Insert execution attributes. A {@link LobHandler} must be provided, since
-	 * any attributes that don't match a provided type are serialized into a
-	 * blob.
-	 */
-	public void saveExecutionContext(final StepExecution stepExecution) {
-
-		final Long executionId = stepExecution.getId();
-		final ExecutionContext executionContext = stepExecution.getExecutionContext();
-		Assert.notNull(executionId, "ExecutionId must not be null.");
-		Assert.notNull(executionContext, "The ExecutionContext must not be null.");
-
-		for (Iterator it = executionContext.entrySet().iterator(); it.hasNext();) {
-			Entry entry = (Entry) it.next();
-			final String key = entry.getKey().toString();
-			final Object value = entry.getValue();
-
-			if (value instanceof String) {
-				insertExecutionAttribute(executionId, key, value, AttributeType.STRING);
-			}
-			else if (value instanceof Double) {
-				insertExecutionAttribute(executionId, key, value, AttributeType.DOUBLE);
-			}
-			else if (value instanceof Long) {
-				insertExecutionAttribute(executionId, key, value, AttributeType.LONG);
-			}
-			else {
-				insertExecutionAttribute(executionId, key, value, AttributeType.OBJECT);
-			}
-		}
-	}
-
 	private void insertExecutionAttribute(final Long executionId, final String key, final Object value,
 			final AttributeType type) {
 		PreparedStatementCallback callback = new AbstractLobCreatingPreparedStatementCallback(lobHandler) {
@@ -237,13 +205,13 @@ public class JdbcStepExecutionDao extends AbstractJdbcBatchMetadataDao implement
 	}
 
 	/**
-	 * update execution attributes. A lob creator must be used, since any
-	 * attributes that don't match a provided type must be serialized into a
+	 * Save or update execution attributes. A lob creator must be used, since
+	 * any attributes that don't match a provided type must be serialized into a
 	 * blob.
 	 * 
 	 * @see {@link LobCreator}
 	 */
-	public void updateExecutionContext(final StepExecution stepExecution) {
+	public void saveOrUpdateExecutionContext(final StepExecution stepExecution) {
 
 		Long executionId = stepExecution.getId();
 		ExecutionContext executionContext = stepExecution.getExecutionContext();
@@ -320,6 +288,10 @@ public class JdbcStepExecutionDao extends AbstractJdbcBatchMetadataDao implement
 		}
 	}
 
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.batch.execution.repository.dao.StepExecutionDao#updateStepExecution(org.springframework.batch.core.domain.StepExecution)
+	 */
 	public void updateStepExecution(StepExecution stepExecution) {
 
 		validateStepExecution(stepExecution);

spring-batch-execution/src/main/java/org/springframework/batch/execution/repository/dao/MapStepDao.java

@@ -46,12 +46,7 @@ public class MapStepDao implements StepExecutionDao {
 		return (ExecutionContext) contextsByStepExecutionId.get(stepExecution.getId());
 	}
 
-	public void saveExecutionContext(StepExecution stepExecution) {
-		contextsByStepExecutionId.put(stepExecution.getId(), stepExecution.getExecutionContext());
-	}
-
-	public void updateExecutionContext(StepExecution stepExecution) {
-		Assert.notNull(contextsByStepExecutionId.get(stepExecution.getId()), "execution context should already be saved");
+	public void saveOrUpdateExecutionContext(StepExecution stepExecution) {
 		contextsByStepExecutionId.put(stepExecution.getId(), stepExecution.getExecutionContext());
 	}
 

spring-batch-execution/src/main/java/org/springframework/batch/execution/repository/dao/StepExecutionDao.java

@@ -37,18 +37,12 @@ public interface StepExecutionDao {
 	/**
 	 * Save the {@link ExecutionContext} of the given {@link StepExecution}.
 	 * 
-	 * @param executionId to be saved
-	 * @param executionContext to be saved.
-	 * @throws IllegalArgumentException if the executionId or attributes are
-	 * null.
+	 * @param stepExecution the {@link StepExecution} containing the
+	 * {@link ExecutionContext} to be saved.
+	 * @throws IllegalArgumentException if the attributes are null.
 	 */
-	void saveExecutionContext(StepExecution stepExecution);
+	void saveOrUpdateExecutionContext(StepExecution stepExecution);
 
-	/**
-	 * Update the ExecutionContext of given {@link StepExecution}.
-	 */
-	void updateExecutionContext(StepExecution stepExecution);
-	
 	StepExecution getStepExecution(JobExecution jobExecution, Step step);
 
 }

spring-batch-execution/src/main/java/org/springframework/batch/execution/step/ItemOrientedStep.java

@@ -41,7 +41,7 @@ import org.springframework.batch.item.ItemRecoverer;
 import org.springframework.batch.item.ItemStream;
 import org.springframework.batch.item.ItemWriter;
 import org.springframework.batch.item.KeyedItemReader;
-import org.springframework.batch.item.exception.ResetFailedException;
+import org.springframework.batch.item.exception.CommitFailedException;
 import org.springframework.batch.item.stream.SimpleStreamManager;
 import org.springframework.batch.repeat.ExitStatus;
 import org.springframework.batch.repeat.RepeatCallback;
@@ -223,6 +223,7 @@ public class ItemOrientedStep extends AbstractStep implements InitializingBean {
 		boolean isRestart = jobRepository.getStepExecutionCount(jobInstance, this) > 0 ? true : false;
 
 		ExitStatus status = ExitStatus.FAILED;
+		final ExceptionHolder fatalException = new ExceptionHolder();
 
 		try {
 
@@ -230,7 +231,7 @@ public class ItemOrientedStep extends AbstractStep implements InitializingBean {
 			// We need to save the step execution right away, before we start
 			// using its ID. It would be better to make the creation atomic in
 			// the caller.
-			updateStatus(stepExecution, BatchStatus.STARTED);
+			fatalException.setException(updateStatus(stepExecution, BatchStatus.STARTED));
 
 			StepContext parentStepContext = StepSynchronizationManager.getContext();
 			final StepContext stepContext = new SimpleStepContext(stepExecution, parentStepContext);
@@ -280,15 +281,33 @@ public class ItemOrientedStep extends AbstractStep implements InitializingBean {
 							stepExecution.apply(contribution);
 
 							streamManager.update(stepExecution.getExecutionContext());
-							jobRepository.saveOrUpdate(stepExecution);
+							try {
+								stepExecution.setStatus(BatchStatus.COMPLETED);
+								jobRepository.saveOrUpdateExecutionContext(stepExecution);
+							}
+							catch (Exception e) {
+								fatalException.setException(e);
+								stepExecution.setStatus(BatchStatus.UNKNOWN);
+								throw new CommitFailedException("Fatal error detected during commit", e);
+							}
 
 						}
 
-						itemReader.mark();
-						itemWriter.flush();
-						streamManager.commit(transaction);
+						try {
+							itemReader.mark();
+							itemWriter.flush();
+							streamManager.commit(transaction);
+						}
+						catch (Exception e) {
+							fatalException.setException(e);
+							stepExecution.setStatus(BatchStatus.UNKNOWN);
+							throw new CommitFailedException("Fatal error detected during commit", e);
+						}
 
 					}
+					catch (CommitFailedException e) {
+						throw e;
+					}
 					catch (Throwable t) {
 						/*
 						 * Any exception thrown within the transaction template
@@ -305,16 +324,9 @@ public class ItemOrientedStep extends AbstractStep implements InitializingBean {
 							itemWriter.clear();
 							streamManager.rollback(transaction);
 						}
-						catch (ResetFailedException e) {
-							// The original Throwable cause is in danger of
-							// being lost here, so we log the reset
-							// failure and re-throw with cause of the rollback.
-							logger.error("Encountered reset error on rollback: "
-									+ "one of the streams may be in an inconsistent state, "
-									+ "so this step should not proceed", e);
-							throw new ResetFailedException("Encountered reset error on rollback.  "
-									+ "Consult logs for the cause of the reet failure.  "
-									+ "The cause of the original rollback is incuded here.", t);
+						catch (Exception e) {
+							fatalException.setException(e);
+							stepExecution.setStatus(BatchStatus.UNKNOWN);
 						}
 						if (t instanceof RuntimeException) {
 							throw (RuntimeException) t;
@@ -334,7 +346,11 @@ public class ItemOrientedStep extends AbstractStep implements InitializingBean {
 				}
 			});
 
-			updateStatus(stepExecution, BatchStatus.COMPLETED);
+			fatalException.setException(updateStatus(stepExecution, BatchStatus.COMPLETED));
+		}
+		catch (CommitFailedException e) {
+			logger.error("Fatal error detected during commit.");
+			throw e;
 		}
 		catch (RuntimeException e) {
 
@@ -344,12 +360,12 @@ public class ItemOrientedStep extends AbstractStep implements InitializingBean {
 				updateStatus(stepExecution, BatchStatus.STOPPED);
 				throw (JobInterruptedException) e.getCause();
 			}
-			else if (e instanceof ResetFailedException) {
-				updateStatus(stepExecution, BatchStatus.UNKNOWN);
-				throw (ResetFailedException) e;
+			else if (!fatalException.hasException()) {
+				updateStatus(stepExecution, BatchStatus.FAILED);
+				throw e;
 			}
 			else {
-				updateStatus(stepExecution, BatchStatus.FAILED);
+				logger.error("Fatal error detected during rollback caused by underlying exception: ", e);
 				throw e;
 			}
 
@@ -357,21 +373,45 @@ public class ItemOrientedStep extends AbstractStep implements InitializingBean {
 		finally {
 			stepExecution.setExitStatus(status);
 			stepExecution.setEndTime(new Date(System.currentTimeMillis()));
+
 			try {
-				jobRepository.saveOrUpdate(stepExecution);
-				streamManager.close(stepExecution.getExecutionContext());
-			}
-			catch (Exception e) {
-				logger
-						.error(
-								"Failed to update step execution: probably fatal, so there is already an exception on the stack.",
-								e);
+
+				try {
+					jobRepository.saveOrUpdate(stepExecution);
+				}
+				catch (RuntimeException e) {
+					String msg = "Fatal error detected during final save of meta data";
+					logger.error(msg, e);
+					if (!fatalException.hasException()) {
+						fatalException.setException(e);
+					}
+					throw new BatchCriticalException(msg, fatalException.getException());
+				}
+
+				try {
+					streamManager.close(stepExecution.getExecutionContext());
+				}
+				catch (RuntimeException e) {
+					String msg = "Fatal error detected during close of streams. "
+							+ "The job execution completed (possibly unsuccessfully but with consistent meta-data).";
+					logger.error(msg, e);
+					if (!fatalException.hasException()) {
+						fatalException.setException(e);
+					}
+					throw new BatchCriticalException(msg, fatalException.getException());
+				}
+
 			}
 			finally {
 				// clear any registered synchronizations
-
 				StepSynchronizationManager.close();
 			}
+
+			if (fatalException.hasException()) {
+				throw new BatchCriticalException("Encountered an error saving batch meta data.", fatalException
+						.getException());
+			}
+
 		}
 
 	}
@@ -527,14 +567,44 @@ public class ItemOrientedStep extends AbstractStep implements InitializingBean {
 	 * @param stepExecution the current stepExecution
 	 * @param status the status to set
 	 */
-	private void updateStatus(StepExecution stepExecution, BatchStatus status) {
+	private Exception updateStatus(StepExecution stepExecution, BatchStatus status) {
 		stepExecution.setStatus(status);
 		try {
 			jobRepository.saveOrUpdate(stepExecution);
+			return null;
 		}
 		catch (Exception e) {
-			logger.error("Failed to update step execution with status: probably fatal.", e);
+			return e;
 		}
 
 	}
+
+	/**
+	 * @author Dave Syer
+	 * 
+	 */
+	private static class ExceptionHolder {
+
+		private Exception exception;
+
+		public boolean hasException() {
+			return exception != null;
+		}
+
+		/**
+		 * @param exception
+		 */
+		public void setException(Exception exception) {
+			this.exception = exception;
+		}
+
+		/**
+		 * @return
+		 */
+		public Exception getException() {
+			return this.exception;
+		}
+
+	}
+
 }

spring-batch-execution/src/main/java/org/springframework/batch/execution/step/TaskletStep.java

@@ -173,6 +173,7 @@ public class TaskletStep implements Step, InitializingBean, BeanNameAware {
 		updateStatus(stepExecution, BatchStatus.STARTED);
 
 		ExitStatus exitStatus = ExitStatus.FAILED;
+		Exception fatalException = null;
 		try {
 
 			StepContext parentStepContext = StepSynchronizationManager.getContext();
@@ -191,9 +192,20 @@ public class TaskletStep implements Step, InitializingBean, BeanNameAware {
 				}
 			});
 
-			updateStatus(stepExecution, BatchStatus.COMPLETED);
+			try {
+				jobRepository.saveOrUpdateExecutionContext(stepExecution);
+				updateStatus(stepExecution, BatchStatus.COMPLETED);
+			} catch (Exception e) {
+				fatalException = e;
+				updateStatus(stepExecution, BatchStatus.UNKNOWN);
+			}
 
 		}
+		catch (RuntimeException e) {
+			logger.error("Encountered an error running the tasklet");
+			updateStatus(stepExecution, BatchStatus.FAILED);
+			throw e;
+		}
 		catch (Exception e) {
 			logger.error("Encountered an error running the tasklet");
 			updateStatus(stepExecution, BatchStatus.FAILED);
@@ -206,18 +218,22 @@ public class TaskletStep implements Step, InitializingBean, BeanNameAware {
 				jobRepository.saveOrUpdate(stepExecution);
 			}
 			catch (Exception e) {
-				logger.error("Encountered error saving batch meta data.  "
-						+ "This job is now in an unknown state and should not be restarted.", e);
+				fatalException = e;
 			}
 			finally {
 				StepSynchronizationManager.close();
+				if (fatalException!=null) {
+					logger.error("Encountered an error saving batch meta data."
+							+ "This job is now in an unknown state and should not be restarted.", fatalException);
+					throw new BatchCriticalException("Encountered an error saving batch meta data.", fatalException);
+				}
 			}
-		}
+		}		
+	
 	}
 
 	private void updateStatus(StepExecution stepExecution, BatchStatus status) {
 		stepExecution.setStatus(status);
-		jobRepository.saveOrUpdate(stepExecution);
 	}
 
 }

spring-batch-execution/src/test/java/org/springframework/batch/execution/repository/MockStepDao.java

@@ -29,8 +29,6 @@ public class MockStepDao implements StepExecutionDao {
 
 	private List newSteps;
 
-	private int currentNewStep = 0;
-
 	public List findStepInstances(JobInstance job) {
 		return newSteps;
 	}
@@ -45,18 +43,11 @@ public class MockStepDao implements StepExecutionDao {
 		this.newSteps = steps;
 	}
 
-	public void resetCurrentNewStep() {
-		currentNewStep = 0;
-	}
-
 	public ExecutionContext findExecutionContext(StepExecution stepExecution) {
 		return null;
 	}
 
-	public void saveExecutionContext(StepExecution stepExecution) {
-	}
-
-	public void updateExecutionContext(StepExecution stepExecution) {
+	public void saveOrUpdateExecutionContext(StepExecution stepExecution) {
 	}
 
 	public StepExecution getStepExecution(JobExecution jobExecution, Step step) {

spring-batch-execution/src/test/java/org/springframework/batch/execution/repository/SimpleJobRepositoryTests.java

@@ -324,23 +324,44 @@ public class SimpleJobRepositoryTests extends TestCase {
 		ExecutionContext executionContext = new ExecutionContext();
 		stepExecution.setExecutionContext(executionContext);
 		stepExecutionDao.updateStepExecution(stepExecution);
-		stepExecutionDao.updateExecutionContext(stepExecution);
 		stepExecutionDaoControl.replay();
 		jobRepository.saveOrUpdate(stepExecution);
 		stepExecutionDaoControl.verify();
 	}
 
+	public void testUpdateExecutionContext() {
+		StepExecution stepExecution = new StepExecution(new StepSupport("stepName"), null, new Long(1));
+		stepExecution.setId(new Long(11));
+		ExecutionContext executionContext = new ExecutionContext();
+		stepExecution.setExecutionContext(executionContext);
+		stepExecutionDao.updateStepExecution(stepExecution);
+		stepExecutionDao.saveOrUpdateExecutionContext(stepExecution);
+		stepExecutionDaoControl.replay();
+		jobRepository.saveOrUpdateExecutionContext(stepExecution);
+		stepExecutionDaoControl.verify();
+	}
+
 	public void testSaveExistingStepExecution() {
 		StepExecution stepExecution = new StepExecution(new StepSupport("stepName"), new JobExecution(null), null);
 		ExecutionContext executionContext = new ExecutionContext();
 		stepExecution.setExecutionContext(executionContext);
 		stepExecutionDao.saveStepExecution(stepExecution);
-		stepExecutionDao.saveExecutionContext(stepExecution);
 		stepExecutionDaoControl.replay();
 		jobRepository.saveOrUpdate(stepExecution);
 		stepExecutionDaoControl.verify();
 	}
 
+	public void testSaveExistingExecutionContext() {
+		StepExecution stepExecution = new StepExecution(new StepSupport("stepName"), new JobExecution(null), null);
+		ExecutionContext executionContext = new ExecutionContext();
+		stepExecution.setExecutionContext(executionContext);
+		stepExecutionDao.saveStepExecution(stepExecution);
+		stepExecutionDao.saveOrUpdateExecutionContext(stepExecution);
+		stepExecutionDaoControl.replay();
+		jobRepository.saveOrUpdateExecutionContext(stepExecution);
+		stepExecutionDaoControl.verify();
+	}
+
 	public void testSaveOrUpdateStepExecutionException() {
 
 		StepExecution stepExecution = new StepExecution(new StepSupport("stepName"), null, null);

spring-batch-execution/src/test/java/org/springframework/batch/execution/repository/dao/AbstractStepDaoTests.java

@@ -117,7 +117,7 @@ public abstract class AbstractStepDaoTests extends AbstractTransactionalDataSour
 
 	public void testUpdateStepWithExecutionContext() {
 		stepExecution.setExecutionContext(executionContext);
-		stepExecutionDao.saveExecutionContext(stepExecution);
+		stepExecutionDao.saveOrUpdateExecutionContext(stepExecution);
 		ExecutionContext tempAttributes = stepExecutionDao.findExecutionContext(stepExecution);
 		assertEquals(executionContext, tempAttributes);
 	}
@@ -143,7 +143,7 @@ public abstract class AbstractStepDaoTests extends AbstractTransactionalDataSour
 		execution.setExitStatus(new ExitStatus(false, ExitStatusExceptionClassifier.FATAL_EXCEPTION,
 				"java.lang.Exception"));
 		stepExecutionDao.saveStepExecution(execution);
-		stepExecutionDao.saveExecutionContext(execution);
+		stepExecutionDao.saveOrUpdateExecutionContext(execution);
 		StepExecution retrievedExecution = stepExecutionDao.getStepExecution(jobExecution, step2);
 		assertNotNull(retrievedExecution);
 		assertEquals(execution, retrievedExecution);
@@ -205,14 +205,14 @@ public abstract class AbstractStepDaoTests extends AbstractTransactionalDataSour
 	public void testSaveExecutionContext(){
 		
 		stepExecution.setExecutionContext(executionContext);
-		stepExecutionDao.saveExecutionContext(stepExecution);
+		stepExecutionDao.saveOrUpdateExecutionContext(stepExecution);
 		ExecutionContext attributes = stepExecutionDao.findExecutionContext(stepExecution);
 		assertEquals(executionContext, attributes);
 		executionContext.putString("newString", "newString");
 		executionContext.putLong("newLong", 1);
 		executionContext.putDouble("newDouble", 2.5);
 		executionContext.put("newSerializable", "serializableValue");
-		stepExecutionDao.updateExecutionContext(stepExecution);
+		stepExecutionDao.saveOrUpdateExecutionContext(stepExecution);
 		attributes = stepExecutionDao.findExecutionContext(stepExecution);
 		assertEquals(executionContext, attributes);
 	}

spring-batch-execution/src/test/java/org/springframework/batch/execution/step/ItemOrientedStepTests.java

@@ -27,7 +27,6 @@ import org.springframework.batch.core.domain.JobExecution;
 import org.springframework.batch.core.domain.JobInstance;
 import org.springframework.batch.core.domain.JobInterruptedException;
 import org.springframework.batch.core.domain.JobParameters;
-import org.springframework.batch.core.domain.Step;
 import org.springframework.batch.core.domain.StepContribution;
 import org.springframework.batch.core.domain.StepExecution;
 import org.springframework.batch.execution.job.JobSupport;
@@ -38,6 +37,7 @@ import org.springframework.batch.execution.repository.dao.MapStepDao;
 import org.springframework.batch.execution.scope.StepSynchronizationManager;
 import org.springframework.batch.execution.step.support.JobRepositorySupport;
 import org.springframework.batch.execution.step.support.StepInterruptionPolicy;
+import org.springframework.batch.io.exception.BatchCriticalException;
 import org.springframework.batch.item.ExecutionContext;
 import org.springframework.batch.item.ItemReader;
 import org.springframework.batch.item.ItemWriter;
@@ -64,6 +64,8 @@ public class ItemOrientedStepTests extends TestCase {
 
 	ArrayList processed = new ArrayList();
 
+	private List list = new ArrayList();
+
 	ItemWriter processor = new AbstractItemWriter() {
 		public void write(Object data) throws Exception {
 			processed.add((String) data);
@@ -84,6 +86,7 @@ public class ItemOrientedStepTests extends TestCase {
 
 	private AbstractStep getStep(String[] strings) throws Exception {
 		ItemOrientedStep step = new ItemOrientedStep();
+		step.setName("stepName");
 		step.setItemWriter(processor);
 		step.setItemReader(getReader(strings));
 		step.setJobRepository(new JobRepositorySupport());
@@ -113,9 +116,8 @@ public class ItemOrientedStepTests extends TestCase {
 
 	public void testStepExecutor() throws Exception {
 
-		Step step = new StepSupport("stepName");
 		JobExecution jobExecutionContext = new JobExecution(jobInstance);
-		StepExecution stepExecution = new StepExecution(step, jobExecutionContext);
+		StepExecution stepExecution = new StepExecution(itemOrientedStep, jobExecutionContext);
 
 		itemOrientedStep.execute(stepExecution);
 		assertEquals(1, processed.size());
@@ -130,10 +132,9 @@ public class ItemOrientedStepTests extends TestCase {
 		template.setCompletionPolicy(new SimpleCompletionPolicy(1));
 		itemOrientedStep.setChunkOperations(template);
 
-		Step step = new StepSupport("stepName");
 		JobExecution jobExecution = new JobExecution(jobInstance);
 
-		StepExecution stepExecution = new StepExecution(step, jobExecution);
+		StepExecution stepExecution = new StepExecution(itemOrientedStep, jobExecution);
 		StepContribution contribution = stepExecution.createStepContribution();
 		itemOrientedStep.processChunk(contribution);
 		assertEquals(1, processed.size());
@@ -150,13 +151,12 @@ public class ItemOrientedStepTests extends TestCase {
 		template.setCompletionPolicy(new SimpleCompletionPolicy(1));
 		itemOrientedStep.setChunkOperations(template);
 
-		final Step step = new StepSupport("stepName");
 		final JobExecution jobExecution = new JobExecution(jobInstance);
-		final StepExecution stepExecution = new StepExecution(step, jobExecution);
+		final StepExecution stepExecution = new StepExecution(itemOrientedStep, jobExecution);
 
 		itemOrientedStep.setItemReader(new AbstractItemReader() {
 			public Object read() throws Exception {
-				assertEquals(step.getName(), stepExecution.getStepName());
+				assertEquals(itemOrientedStep.getName(), stepExecution.getStepName());
 				assertNotNull(StepSynchronizationManager.getContext().getStepExecution());
 				return "foo";
 			}
@@ -175,10 +175,9 @@ public class ItemOrientedStepTests extends TestCase {
 		template.setCompletionPolicy(new SimpleCompletionPolicy(1));
 		itemOrientedStep.setStepOperations(template);
 
-		final Step step = new StepSupport("stepName");
 		final JobExecution jobExecution = new JobExecution(jobInstance);
 		jobExecution.setId(new Long(1));
-		final StepExecution stepExecution = new StepExecution(step, jobExecution);
+		final StepExecution stepExecution = new StepExecution(itemOrientedStep, jobExecution);
 
 		template.setListener(new RepeatListenerSupport() {
 			public void open(RepeatContext context) {
@@ -199,9 +198,8 @@ public class ItemOrientedStepTests extends TestCase {
 		SimpleJobRepository repository = new SimpleJobRepository(new MapJobInstanceDao(), new MapJobExecutionDao(), new MapStepDao());
 		itemOrientedStep.setJobRepository(repository);
 
-		Step step = new StepSupport("stepName");
 		JobExecution jobExecutionContext = new JobExecution(jobInstance);
-		StepExecution stepExecution = new StepExecution(step, jobExecutionContext);
+		StepExecution stepExecution = new StepExecution(itemOrientedStep, jobExecutionContext);
 
 		itemOrientedStep.execute(stepExecution);
 		assertEquals(1, processed.size());
@@ -224,10 +222,9 @@ public class ItemOrientedStepTests extends TestCase {
 
 		};
 
-		Step step = new StepSupport("stepName");
 		itemOrientedStep.setItemReader(itemReader);
 		JobExecution jobExecutionContext = new JobExecution(jobInstance);
-		StepExecution stepExecution = new StepExecution(step, jobExecutionContext);
+		StepExecution stepExecution = new StepExecution(itemOrientedStep, jobExecutionContext);
 
 		try {
 			itemOrientedStep.execute(stepExecution);
@@ -255,10 +252,9 @@ public class ItemOrientedStepTests extends TestCase {
 
 		};
 
-		Step step = new StepSupport("stepName");
 		itemOrientedStep.setItemReader(itemReader);
 		JobExecution jobExecutionContext = new JobExecution(jobInstance);
-		StepExecution stepExecution = new StepExecution(step, jobExecutionContext);
+		StepExecution stepExecution = new StepExecution(itemOrientedStep, jobExecutionContext);
 
 		try {
 			itemOrientedStep.execute(stepExecution);
@@ -274,11 +270,10 @@ public class ItemOrientedStepTests extends TestCase {
 	 * saveExecutionAttributes = true, doesn't have restoreFrom called on it.
 	 */
 	public void testNonRestartedJob() throws Exception {
-		Step step = new StepSupport("stepName");
 		MockRestartableItemReader tasklet = new MockRestartableItemReader();
 		itemOrientedStep.setItemReader(tasklet);
 		JobExecution jobExecutionContext = new JobExecution(jobInstance);
-		StepExecution stepExecution = new StepExecution(step, jobExecutionContext);
+		StepExecution stepExecution = new StepExecution(itemOrientedStep, jobExecutionContext);
 
 		itemOrientedStep.execute(stepExecution);
 
@@ -286,6 +281,36 @@ public class ItemOrientedStepTests extends TestCase {
 		assertTrue(tasklet.isGetExecutionAttributesCalled());
 	}
 
+	public void testSuccessfulExecutionWithExecutionContext() throws Exception {
+		final JobExecution jobExecution = new JobExecution(jobInstance);
+		final StepExecution stepExecution = new StepExecution(itemOrientedStep, jobExecution);
+		itemOrientedStep.setJobRepository(new JobRepositorySupport() {
+			public void saveOrUpdateExecutionContext(StepExecution stepExecution) {
+				list.add(stepExecution);
+			}
+		});
+		itemOrientedStep.execute(stepExecution);
+		assertEquals(1, list.size());
+	}
+
+	public void testSuccessfulExecutionWithFailureOnSaveOfExecutionContext() throws Exception {
+		final JobExecution jobExecution = new JobExecution(jobInstance);
+		final StepExecution stepExecution = new StepExecution(itemOrientedStep, jobExecution);
+		itemOrientedStep.setJobRepository(new JobRepositorySupport() {
+			public void saveOrUpdateExecutionContext(StepExecution stepExecution) {
+				throw new RuntimeException("foo");
+			}
+		});
+		try {
+			itemOrientedStep.execute(stepExecution);
+			fail("Expected BatchCriticalException");
+		}
+		catch (BatchCriticalException e) {
+			assertEquals("foo", e.getCause().getMessage());
+		}
+		assertEquals(BatchStatus.UNKNOWN, stepExecution.getStatus());
+	}
+
 	/*
 	 * make sure a job that has been executed before, and is therefore being
 	 * restarted, is restored.
@@ -315,11 +340,10 @@ public class ItemOrientedStepTests extends TestCase {
 	 * it.
 	 */
 	public void testNoSaveExecutionAttributesRestartableJob() {
-		Step step = new StepSupport("stepName");
 		MockRestartableItemReader tasklet = new MockRestartableItemReader();
 		itemOrientedStep.setItemReader(tasklet);
 		JobExecution jobExecutionContext = new JobExecution(jobInstance);
-		StepExecution stepExecution = new StepExecution(step, jobExecutionContext);
+		StepExecution stepExecution = new StepExecution(itemOrientedStep, jobExecutionContext);
 
 		try {
 			itemOrientedStep.execute(stepExecution);
@@ -337,14 +361,13 @@ public class ItemOrientedStepTests extends TestCase {
 	 * Restartable.
 	 */
 	public void testRestartJobOnNonRestartableTasklet() throws Exception {
-		Step step = new StepSupport("stepName");
 		itemOrientedStep.setItemReader(new AbstractItemReader() {
 			public Object read() throws Exception {
 				return "foo";
 			}
 		});
 		JobExecution jobExecution = new JobExecution(jobInstance);
-		StepExecution stepExecution = new StepExecution(step, jobExecution);
+		StepExecution stepExecution = new StepExecution(itemOrientedStep, jobExecution);
 
 		itemOrientedStep.execute(stepExecution);
 	}
@@ -362,14 +385,13 @@ public class ItemOrientedStepTests extends TestCase {
 	}
 
 	public void testStreamManager() throws Exception {
-		Step step = new StepSupport("stepName");
 		itemOrientedStep.setItemReader(new AbstractItemReader() {
 			public Object read() throws Exception {
 				return "foo";
 			}
 		});
 		JobExecution jobExecution = new JobExecution(jobInstance);
-		StepExecution stepExecution = new StepExecution(step, jobExecution);
+		StepExecution stepExecution = new StepExecution(itemOrientedStep, jobExecution);
 
 		assertEquals(false, stepExecution.getExecutionContext().containsKey("foo"));
 
@@ -389,53 +411,6 @@ public class ItemOrientedStepTests extends TestCase {
 		assertEquals("bar", stepExecution.getExecutionContext().getString("foo"));
 	}
 
-	private class MockRestartableItemReader extends ItemStreamSupport implements ItemReader {
-
-		private boolean getExecutionAttributesCalled = false;
-
-		private boolean restoreFromCalled = false;
-
-		private boolean restoreFromCalledWithSomeContext = false;
-
-		private ExecutionContext executionContext;
-
-		public Object read() throws Exception {
-			StepSynchronizationManager.getContext().setAttribute("TASKLET_TEST", this);
-			return "item";
-		}
-
-		public boolean isRestoreFromCalledWithSomeContext() {
-			return restoreFromCalledWithSomeContext;
-		}
-
-		public void update(ExecutionContext executionContext) {
-			getExecutionAttributesCalled = true;
-			executionContext.putString("spam", "bucket");
-		}
-
-		public boolean isGetExecutionAttributesCalled() {
-			return getExecutionAttributesCalled;
-		}
-
-		public boolean isRestoreFromCalled() {
-			return restoreFromCalled;
-		}
-
-		public void open(ExecutionContext executionContext) throws StreamException {
-			this.executionContext = executionContext;
-		}
-
-		public void close(ExecutionContext executionContext) throws StreamException {
-		}
-
-		public void mark() throws MarkFailedException {
-		}
-
-		public void reset() throws ResetFailedException {
-		}
-
-	}
-
 	public void testStatusForInterruptedException() {
 
 		StepInterruptionPolicy interruptionPolicy = new StepInterruptionPolicy() {
@@ -464,9 +439,8 @@ public class ItemOrientedStepTests extends TestCase {
 
 		itemOrientedStep.setItemReader(itemReader);
 
-		Step step = new StepSupport("stepName");
 		JobExecution jobExecutionContext = new JobExecution(jobInstance);
-		StepExecution stepExecution = new StepExecution(step, jobExecutionContext);
+		StepExecution stepExecution = new StepExecution(itemOrientedStep, jobExecutionContext);
 
 		stepExecution.setExecutionContext(new ExecutionContext(PropertiesConverter.stringToProperties("foo=bar")));
 		// step.setLastExecution(stepExecution);
@@ -483,6 +457,33 @@ public class ItemOrientedStepTests extends TestCase {
 		}
 	}
 
+	public void testStatusForNormalFailure() throws Exception {
+
+		ItemReader itemReader = new AbstractItemReader() {
+			public Object read() throws Exception {
+				// Trigger a rollback
+				throw new RuntimeException("Foo");
+			}
+		};
+		itemOrientedStep.setItemReader(itemReader);
+
+		JobExecution jobExecutionContext = jobInstance.createJobExecution();
+		StepExecution stepExecution = new StepExecution(itemOrientedStep, jobExecutionContext);
+
+		stepExecution.setExecutionContext(new ExecutionContext(PropertiesConverter.stringToProperties("foo=bar")));
+		// step.setLastExecution(stepExecution);
+
+		try {
+			itemOrientedStep.execute(stepExecution);
+			fail("Expected RuntimeException");
+		}
+		catch (RuntimeException ex) {
+			assertEquals(BatchStatus.FAILED, stepExecution.getStatus());
+			// The original rollback was caused by this one:
+			assertEquals("Foo", ex.getMessage());
+		}
+	}
+
 	public void testStatusForResetFailedException() throws Exception {
 
 		ItemReader itemReader = new AbstractItemReader() {
@@ -500,24 +501,160 @@ public class ItemOrientedStepTests extends TestCase {
 			}
 		});
 
-		Step step = new StepSupport("stepName");
 		JobExecution jobExecutionContext = jobInstance.createJobExecution();
-		StepExecution stepExecution = new StepExecution(step, jobExecutionContext);
+		StepExecution stepExecution = new StepExecution(itemOrientedStep, jobExecutionContext);
 
 		stepExecution.setExecutionContext(new ExecutionContext(PropertiesConverter.stringToProperties("foo=bar")));
 		// step.setLastExecution(stepExecution);
 
 		try {
 			itemOrientedStep.execute(stepExecution);
-			fail("Expected ResetFailedException");
+			fail("Expected BatchCriticalException");
 		}
-		catch (ResetFailedException ex) {
+		catch (BatchCriticalException ex) {
 			assertEquals(BatchStatus.UNKNOWN, stepExecution.getStatus());
 			String msg = stepExecution.getExitStatus().getExitDescription();
 			assertTrue("Message does not contain ResetFailedException: " + msg, msg.contains("ResetFailedException"));
 			// The original rollback was caused by this one:
-			assertEquals("Foo", ex.getCause().getMessage());
+			assertEquals("Bar", ex.getCause().getMessage());
 		}
 	}
 
+	public void testStatusForCommitFailedException() throws Exception {
+
+		itemOrientedStep.setStreamManager(new SimpleStreamManager(transactionManager) {
+			public void commit(TransactionStatus status) {
+				super.commit(status);
+				// Simulate failure on rollback when stream resets
+				throw new RuntimeException("Bar");
+			}
+		});
+
+		JobExecution jobExecutionContext = jobInstance.createJobExecution();
+		StepExecution stepExecution = new StepExecution(itemOrientedStep, jobExecutionContext);
+
+		stepExecution.setExecutionContext(new ExecutionContext(PropertiesConverter.stringToProperties("foo=bar")));
+		// step.setLastExecution(stepExecution);
+
+		try {
+			itemOrientedStep.execute(stepExecution);
+			fail("Expected BatchCriticalException");
+		}
+		catch (BatchCriticalException ex) {
+			assertEquals(BatchStatus.UNKNOWN, stepExecution.getStatus());
+			String msg = stepExecution.getExitStatus().getExitDescription();
+			assertEquals("", msg);
+			msg = ex.getMessage();
+			assertTrue("Message does not contain 'saving': " + msg, msg.contains("saving"));
+			// The original rollback was caused by this one:
+			assertEquals("Bar", ex.getCause().getMessage());
+		}
+	}
+
+	public void testStatusForFinalUpdateFailedException() throws Exception {
+
+		itemOrientedStep.setJobRepository(new JobRepositorySupport() {
+			public void saveOrUpdate(StepExecution stepExecution) {
+				if (stepExecution.getEndTime()!=null) {
+					throw new RuntimeException("Bar");
+				}
+				super.saveOrUpdate(stepExecution);
+			}
+		});
+
+		JobExecution jobExecutionContext = jobInstance.createJobExecution();
+		StepExecution stepExecution = new StepExecution(itemOrientedStep, jobExecutionContext);
+
+		try {
+			itemOrientedStep.execute(stepExecution);
+			fail("Expected RuntimeException");
+		}
+		catch (RuntimeException ex) {
+			// The job actually completeed, but teh streams couldn't be closed.
+			assertEquals(BatchStatus.COMPLETED, stepExecution.getStatus());
+			String msg = stepExecution.getExitStatus().getExitDescription();
+			assertEquals("", msg);
+			msg = ex.getMessage();
+			assertTrue("Message does not contain 'final': " + msg, msg.contains("final"));
+			// The original rollback was caused by this one:
+			assertEquals("Bar", ex.getCause().getMessage());
+		}
+	}
+
+	public void testStatusForCloseFailedException() throws Exception {
+
+		itemOrientedStep.setStreamManager(new SimpleStreamManager(transactionManager) {
+			public void close(ExecutionContext executionContext) throws StreamException {
+				super.close(executionContext);
+				// Simulate failure on rollback when stream resets
+				throw new RuntimeException("Bar");
+			}
+		});
+
+		JobExecution jobExecutionContext = jobInstance.createJobExecution();
+		StepExecution stepExecution = new StepExecution(itemOrientedStep, jobExecutionContext);
+
+		stepExecution.setExecutionContext(new ExecutionContext(PropertiesConverter.stringToProperties("foo=bar")));
+		// step.setLastExecution(stepExecution);
+
+		try {
+			itemOrientedStep.execute(stepExecution);
+			fail("Expected BatchCriticalException");
+		}
+		catch (BatchCriticalException ex) {
+			// The job actually completeed, but teh streams couldn't be closed.
+			assertEquals(BatchStatus.COMPLETED, stepExecution.getStatus());
+			String msg = stepExecution.getExitStatus().getExitDescription();
+			assertEquals("", msg);
+			msg = ex.getMessage();
+			assertTrue("Message does not contain 'close': " + msg, msg.contains("close"));
+			// The original rollback was caused by this one:
+			assertEquals("Bar", ex.getCause().getMessage());
+		}
+	}
+
+	private class MockRestartableItemReader extends ItemStreamSupport implements ItemReader {
+
+		private boolean getExecutionAttributesCalled = false;
+
+		private boolean restoreFromCalled = false;
+
+		private boolean restoreFromCalledWithSomeContext = false;
+
+		public Object read() throws Exception {
+			StepSynchronizationManager.getContext().setAttribute("TASKLET_TEST", this);
+			return "item";
+		}
+
+		public boolean isRestoreFromCalledWithSomeContext() {
+			return restoreFromCalledWithSomeContext;
+		}
+
+		public void update(ExecutionContext executionContext) {
+			getExecutionAttributesCalled = true;
+			executionContext.putString("spam", "bucket");
+		}
+
+		public boolean isGetExecutionAttributesCalled() {
+			return getExecutionAttributesCalled;
+		}
+
+		public boolean isRestoreFromCalled() {
+			return restoreFromCalled;
+		}
+
+		public void open(ExecutionContext executionContext) throws StreamException {
+		}
+
+		public void close(ExecutionContext executionContext) throws StreamException {
+		}
+
+		public void mark() throws MarkFailedException {
+		}
+
+		public void reset() throws ResetFailedException {
+		}
+
+	}
+
 }

spring-batch-execution/src/test/java/org/springframework/batch/execution/step/TaskletStepTests.java

@@ -5,6 +5,7 @@ import java.util.List;
 
 import junit.framework.TestCase;
 
+import org.springframework.batch.core.domain.BatchStatus;
 import org.springframework.batch.core.domain.JobExecution;
 import org.springframework.batch.core.domain.JobInstance;
 import org.springframework.batch.core.domain.JobInterruptedException;
@@ -69,6 +70,31 @@ public class TaskletStepTests extends TestCase {
 		assertNotNull(stepExecution.getEndTime());
 	}
 
+	public void testSuccessfulExecutionWithExecutionContext() throws Exception {
+		TaskletStep step = new TaskletStep(new StubTasklet(false, false), new JobRepositorySupport() {
+			public void saveOrUpdateExecutionContext(StepExecution stepExecution) {
+				list.add(stepExecution);
+			}
+		});
+		step.execute(stepExecution);
+		assertEquals(1, list.size());
+	}
+
+	public void testSuccessfulExecutionWithFailureOnSaveOfExecutionContext() throws Exception {
+		TaskletStep step = new TaskletStep(new StubTasklet(false, false, true), new JobRepositorySupport() {
+			public void saveOrUpdateExecutionContext(StepExecution stepExecution) {
+				throw new RuntimeException("foo");
+			}
+		});
+		try {
+			step.execute(stepExecution);
+			fail("Expected BatchCriticalException");
+		} catch (BatchCriticalException e){
+			assertEquals("foo", e.getCause().getMessage());
+		}
+		assertEquals(BatchStatus.UNKNOWN, stepExecution.getStatus());
+	}
+
 	public void testFailureExecution() throws Exception {
 		TaskletStep step = new TaskletStep(new StubTasklet(true, false), new JobRepositorySupport());
 		step.execute(stepExecution);
@@ -97,7 +123,7 @@ public class TaskletStepTests extends TestCase {
 			step.execute(stepExecution);
 			fail();
 		}
-		catch (BatchCriticalException e) {
+		catch (RuntimeException e) {
 			assertNotNull(stepExecution.getStartTime());
 			assertEquals(ExitStatus.FAILED, stepExecution.getExitStatus());
 			assertNotNull(stepExecution.getEndTime());

spring-batch-execution/src/test/java/org/springframework/batch/execution/step/support/JobRepositorySupport.java

@@ -47,6 +47,12 @@ public class JobRepositorySupport implements JobRepository {
 	 */
 	public void saveOrUpdate(StepExecution stepExecution) {
 	}
+	
+	/* (non-Javadoc)
+	 * @see org.springframework.batch.core.repository.JobRepository#saveOrUpdateExecutionContext(org.springframework.batch.core.domain.StepExecution)
+	 */
+	public void saveOrUpdateExecutionContext(StepExecution stepExecution) {
+	}
 
 	/* (non-Javadoc)
 	 * @see org.springframework.batch.container.common.repository.JobRepository#update(org.springframework.batch.container.common.domain.Job)