From a35a9d79dc12d57d3110218d826554d80ab3087a Mon Sep 17 00:00:00 2001
From: dsyer <dsyer@2e8e966d-6415-0410-94b9-dc4e65435d03>
Date: Wed, 27 Feb 2008 16:44:52 +0000
Subject: [PATCH] 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.
---
 .../batch/core/repository/JobRepository.java  |  26 +-
 .../batch/execution/job/AbstractJob.java      |   2 -
 .../repository/SimpleJobRepository.java       |  19 +-
 .../repository/dao/JdbcStepExecutionDao.java  |  42 +--
 .../execution/repository/dao/MapStepDao.java  |   7 +-
 .../repository/dao/StepExecutionDao.java      |  14 +-
 .../execution/step/ItemOrientedStep.java      | 134 ++++++--
 .../batch/execution/step/TaskletStep.java     |  26 +-
 .../execution/repository/MockStepDao.java     |  11 +-
 .../repository/SimpleJobRepositoryTests.java  |  25 +-
 .../repository/dao/AbstractStepDaoTests.java  |   8 +-
 .../execution/step/ItemOrientedStepTests.java | 293 +++++++++++++-----
 .../execution/step/TaskletStepTests.java      |  28 +-
 .../step/support/JobRepositorySupport.java    |   6 +
 .../item/exception/CommitFailedException.java |  46 +++
 15 files changed, 489 insertions(+), 198 deletions(-)
 create mode 100644 spring-batch-infrastructure/src/main/java/org/springframework/batch/item/exception/CommitFailedException.java

diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/repository/JobRepository.java b/spring-batch-core/src/main/java/org/springframework/batch/core/repository/JobRepository.java
index 986015051..af8e01484 100644
--- a/spring-batch-core/src/main/java/org/springframework/batch/core/repository/JobRepository.java
+++ b/spring-batch-core/src/main/java/org/springframework/batch/core/repository/JobRepository.java
@@ -81,18 +81,31 @@ public interface JobRepository {
 	public void saveOrUpdate(JobExecution jobExecution);
 
 	/**
-	 * Save or Update a StepExecution. If no ID is found a new instance will be
-	 * created. (saved). If an ID does exist it will be updated. It is not
-	 * advisable that an ID be assigned to a JobExecution before calling this
-	 * method. Instead, it should be left blank, to be assigned by a
-	 * JobRepository.
+	 * Save or update a {@link StepExecution}. If no ID is found a new instance
+	 * will be created (and saved). If an ID does exist it will be updated. It
+	 * is not advisable that an ID be assigned before calling this method.
+	 * Instead, it should be left blank, to be assigned by a
+	 * {@link JobRepository}. The {@link ExecutionContext} of the
+	 * {@link StepExecution} is <em>not</em> saved: see
+	 * {@link #saveExecutionContext(StepExecution)}.
 	 * 
-	 * Preconditions: StepExecution must have a valid StepId.
+	 * Preconditions: {@link StepExecution} must have a valid {@link Step}.
 	 * 
 	 * @param jobInstance
 	 */
 	public void saveOrUpdate(StepExecution stepExecution);
 
+	/**
+	 * Save the {@link ExecutionContext} of the given {@link StepExecution}.
+	 * Implementations are allowed to ensure that the {@link StepExecution} is
+	 * already saved by calling {@link #saveOrUpdate(StepExecution)} before
+	 * saving the {@link ExecutionContext}.
+	 * 
+	 * @param stepExecution the {@link StepExecution} containing the
+	 * {@link ExecutionContext} to be saved.
+	 */
+	void saveOrUpdateExecutionContext(StepExecution stepExecution);
+
 	/**
 	 * @return the last execution of step for the given job instance.
 	 */
@@ -103,5 +116,4 @@ public interface JobRepository {
 	 */
 	public int getStepExecutionCount(JobInstance jobInstance, Step step);
 
-
 }
diff --git a/spring-batch-execution/src/main/java/org/springframework/batch/execution/job/AbstractJob.java b/spring-batch-execution/src/main/java/org/springframework/batch/execution/job/AbstractJob.java
index 21ef40067..6de446cb5 100644
--- a/spring-batch-execution/src/main/java/org/springframework/batch/execution/job/AbstractJob.java
+++ b/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;
 
diff --git a/spring-batch-execution/src/main/java/org/springframework/batch/execution/repository/SimpleJobRepository.java b/spring-batch-execution/src/main/java/org/springframework/batch/execution/repository/SimpleJobRepository.java
index 9db70a515..27ad8ab33 100644
--- a/spring-batch-execution/src/main/java/org/springframework/batch/execution/repository/SimpleJobRepository.java
+++ b/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
 	 */
diff --git a/spring-batch-execution/src/main/java/org/springframework/batch/execution/repository/dao/JdbcStepExecutionDao.java b/spring-batch-execution/src/main/java/org/springframework/batch/execution/repository/dao/JdbcStepExecutionDao.java
index e1ff8bf10..77a28dab5 100644
--- a/spring-batch-execution/src/main/java/org/springframework/batch/execution/repository/dao/JdbcStepExecutionDao.java
+++ b/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);
diff --git a/spring-batch-execution/src/main/java/org/springframework/batch/execution/repository/dao/MapStepDao.java b/spring-batch-execution/src/main/java/org/springframework/batch/execution/repository/dao/MapStepDao.java
index 50f86f051..0ca393e54 100644
--- a/spring-batch-execution/src/main/java/org/springframework/batch/execution/repository/dao/MapStepDao.java
+++ b/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());
 	}
 
diff --git a/spring-batch-execution/src/main/java/org/springframework/batch/execution/repository/dao/StepExecutionDao.java b/spring-batch-execution/src/main/java/org/springframework/batch/execution/repository/dao/StepExecutionDao.java
index 7b88a0346..8fecda0e5 100644
--- a/spring-batch-execution/src/main/java/org/springframework/batch/execution/repository/dao/StepExecutionDao.java
+++ b/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);
 
 }
diff --git a/spring-batch-execution/src/main/java/org/springframework/batch/execution/step/ItemOrientedStep.java b/spring-batch-execution/src/main/java/org/springframework/batch/execution/step/ItemOrientedStep.java
index b96236d13..4d72ada7b 100644
--- a/spring-batch-execution/src/main/java/org/springframework/batch/execution/step/ItemOrientedStep.java
+++ b/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;
+		}
+
+	}
+
 }
diff --git a/spring-batch-execution/src/main/java/org/springframework/batch/execution/step/TaskletStep.java b/spring-batch-execution/src/main/java/org/springframework/batch/execution/step/TaskletStep.java
index 052a15e7a..30ea7cdc3 100644
--- a/spring-batch-execution/src/main/java/org/springframework/batch/execution/step/TaskletStep.java
+++ b/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);
 	}
 
 }
diff --git a/spring-batch-execution/src/test/java/org/springframework/batch/execution/repository/MockStepDao.java b/spring-batch-execution/src/test/java/org/springframework/batch/execution/repository/MockStepDao.java
index 05be6ea71..0b165b26b 100644
--- a/spring-batch-execution/src/test/java/org/springframework/batch/execution/repository/MockStepDao.java
+++ b/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) {
diff --git a/spring-batch-execution/src/test/java/org/springframework/batch/execution/repository/SimpleJobRepositoryTests.java b/spring-batch-execution/src/test/java/org/springframework/batch/execution/repository/SimpleJobRepositoryTests.java
index a2013424a..29f602c9e 100644
--- a/spring-batch-execution/src/test/java/org/springframework/batch/execution/repository/SimpleJobRepositoryTests.java
+++ b/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);
diff --git a/spring-batch-execution/src/test/java/org/springframework/batch/execution/repository/dao/AbstractStepDaoTests.java b/spring-batch-execution/src/test/java/org/springframework/batch/execution/repository/dao/AbstractStepDaoTests.java
index aeea1538b..9da7b78aa 100644
--- a/spring-batch-execution/src/test/java/org/springframework/batch/execution/repository/dao/AbstractStepDaoTests.java
+++ b/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);
 	}
diff --git a/spring-batch-execution/src/test/java/org/springframework/batch/execution/step/ItemOrientedStepTests.java b/spring-batch-execution/src/test/java/org/springframework/batch/execution/step/ItemOrientedStepTests.java
index 491e17c3b..c5fcba801 100644
--- a/spring-batch-execution/src/test/java/org/springframework/batch/execution/step/ItemOrientedStepTests.java
+++ b/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 {
+		}
+
+	}
+
 }
diff --git a/spring-batch-execution/src/test/java/org/springframework/batch/execution/step/TaskletStepTests.java b/spring-batch-execution/src/test/java/org/springframework/batch/execution/step/TaskletStepTests.java
index f15fd66d6..d87120427 100644
--- a/spring-batch-execution/src/test/java/org/springframework/batch/execution/step/TaskletStepTests.java
+++ b/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());
diff --git a/spring-batch-execution/src/test/java/org/springframework/batch/execution/step/support/JobRepositorySupport.java b/spring-batch-execution/src/test/java/org/springframework/batch/execution/step/support/JobRepositorySupport.java
index 0831242e7..c456d28ea 100644
--- a/spring-batch-execution/src/test/java/org/springframework/batch/execution/step/support/JobRepositorySupport.java
+++ b/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)
diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/exception/CommitFailedException.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/exception/CommitFailedException.java
new file mode 100644
index 000000000..9fbadf6d9
--- /dev/null
+++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/exception/CommitFailedException.java
@@ -0,0 +1,46 @@
+/*
+ * Copyright 2006-2007 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
+ *
+ *      http://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.batch.item.exception;
+
+/**
+ * @author Dave Syer
+ *
+ */
+public class CommitFailedException extends StreamException {
+	/**
+	 * @param message
+	 */
+	public CommitFailedException(String message) {
+		super(message);
+	}
+
+	/**
+	 * @param msg
+	 * @param nested
+	 */
+	public CommitFailedException(String msg, Throwable nested) {
+		super(msg, nested);
+	}
+
+	/**
+	 * @param msg
+	 * @param nested
+	 */
+	public CommitFailedException(Throwable nested) {
+		super(nested);
+	}	
+
+}