OPEN - issue BATCH-1437: Support for CallbackPreferringPlatformTransactionManager (and for native TX in WAS)

Refactored TaskletStep so that it uses a TransactionCallback.  It's actually much nicer that way because commit failures can be detected and accounted for with a proper rollback, instead of just panicking and using BatchStatus.UNKNOWN.

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

@@ -213,7 +213,7 @@ public abstract class AbstractStep implements Step, InitializingBean, BeanNameAw
 		}
 		catch (Throwable e) {
 			logger.error("Encountered an error executing the step", e);
-			stepExecution.setStatus(determineBatchStatus(e));
+			stepExecution.upgradeStatus(determineBatchStatus(e));
 			exitStatus = exitStatus.and(getDefaultExitStatusForFailure(e));
 			stepExecution.addFailureException(e);
 		}

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

@@ -9,7 +9,12 @@ import org.springframework.batch.core.UnexpectedJobExecutionException;
  * result in the step having a status of {@link BatchStatus#UNKNOWN}.
  */
 public class FatalException extends UnexpectedJobExecutionException {
+
 	public FatalException(String string, Throwable e) {
 		super(string, e);
 	}
+
+	public FatalException(String string) {
+		super(string);
+	}
 }

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

@@ -46,6 +46,11 @@ import org.springframework.transaction.PlatformTransactionManager;
 import org.springframework.transaction.TransactionStatus;
 import org.springframework.transaction.interceptor.DefaultTransactionAttribute;
 import org.springframework.transaction.interceptor.TransactionAttribute;
+import org.springframework.transaction.support.TransactionCallback;
+import org.springframework.transaction.support.TransactionSynchronization;
+import org.springframework.transaction.support.TransactionSynchronizationAdapter;
+import org.springframework.transaction.support.TransactionSynchronizationManager;
+import org.springframework.transaction.support.TransactionTemplate;
 import org.springframework.util.Assert;
 
 /**
@@ -248,142 +253,18 @@ public class TaskletStep extends AbstractStep {
 
 				StepExecution stepExecution = chunkContext.getStepContext().getStepExecution();
 
-				StepContribution contribution = stepExecution.createStepContribution();
-
 				// Before starting a new transaction, check for
 				// interruption.
 				interruptionPolicy.checkInterrupted(stepExecution);
 
-				RepeatStatus result = RepeatStatus.CONTINUABLE;
-
-				TransactionStatus transaction = transactionManager.getTransaction(transactionAttribute);
-
-				chunkListener.beforeChunk();
-
-				boolean locked = false;
-
-				Integer oldVersion = null;
-				boolean committed = true;
-
+				RepeatStatus result;
 				try {
-
-					try {
-						try {
-							result = tasklet.execute(contribution, chunkContext);
-							if (result == null) {
-								result = RepeatStatus.FINISHED;
-							}
-						}
-						catch (Exception e) {
-							if (transactionAttribute.rollbackOn(e)) {
-								throw e;
-							}
-						}
-
-					}
-					finally {
-
-						// If the step operations are asynchronous then we need
-						// to synchronize changes to the step execution (at a
-						// minimum). Take the lock *before* changing the step
-						// execution.
-						try {
-							semaphore.acquire();
-							locked = true;
-						}
-						catch (InterruptedException e) {
-							stepExecution.setStatus(BatchStatus.STOPPED);
-							Thread.currentThread().interrupt();
-						}
-
-						// In case we need to push it back to its old value
-						// after a commit fails...
-						oldVersion = stepExecution.getVersion();
-
-						// Apply the contribution to the step
-						// even if unsuccessful
-						logger.debug("Applying contribution: " + contribution);
-						stepExecution.apply(contribution);
-
-					}
-
-					stream.update(stepExecution.getExecutionContext());
-
-					try {
-						// Going to attempt a commit. If it fails this flag will
-						// stay false and we can use that later.
-						committed = false;
-						getJobRepository().updateExecutionContext(stepExecution);
-						stepExecution.incrementCommitCount();
-						/*
-						 * The step execution has to be saved before commit
-						 * because otherwise there is a deadlock between the
-						 * data source pool and the semaphore. As long as only
-						 * one connection is used inside the section of this
-						 * callback that is locked with the semaphore, the
-						 * deadlock is avoided.
-						 */
-						logger.debug("Saving step execution before commit: " + stepExecution);
-						getJobRepository().update(stepExecution);
-						transactionManager.commit(transaction);
-						committed = true;
-					}
-					catch (Exception e) {
-						throw new FatalException("Fatal failure detected", e);
-					}
-
+					result = (RepeatStatus) new TransactionTemplate(transactionManager, transactionAttribute)
+							.execute(new ChunkTransactionCallback(chunkContext));
 				}
-				catch (FatalException e) {
-					try {
-						logger.debug("Rollback for FatalException: " + e.getClass().getName() + ": " + e.getMessage());
-						rollback(stepExecution, transaction);
-					}
-					catch (Exception rollbackException) {
-						/*
-						 * Propagate the original fatal failure; only log the
-						 * failed rollback. The failure can be caused by
-						 * attempting a rollback when the commit has already
-						 * succeeded (which is normal so only logged at debug
-						 * level)
-						 */
-						logger.debug("Rollback caused by fatal failure failed", rollbackException);
-					}
-					throw e;
-				}
-				catch (Error e) {
-					try {
-						logger.debug("Rollback for Error: " + e.getClass().getName() + ": " + e.getMessage());
-						rollback(stepExecution, transaction);
-					}
-					catch (Exception rollbackException) {
-						logger.error("Fatal rollback failure, original exception that caused the rollback is", e);
-						throw new FatalException("Failed while processing rollback", rollbackException);
-					}
-					throw e;
-
-				}
-				catch (Exception e) {
-					try {
-						logger.debug("Rollback for Exception: " + e.getClass().getName() + ": " + e.getMessage());
-						rollback(stepExecution, transaction);
-					}
-					catch (Exception rollbackException) {
-						logger.error("Fatal rollback failure, original exception that caused the rollback is", e);
-						throw new FatalException("Failed while processing rollback", rollbackException);
-					}
-					throw e;
-				}
-				finally {
-					if (!committed && oldVersion != null) {
-						// Wah! the commit failed. We need to rescue the step
-						// execution data.
-						stepExecution.setVersion(oldVersion);
-					}
-					// only release the lock if we acquired it
-					if (locked) {
-						semaphore.release();
-					}
-					locked = false;
+				catch (TransactionException e) {
+					// Allow checked exceptions to be thrown inside callback
+					throw (Exception) e.getCause();
 				}
 
 				chunkListener.afterChunk();
@@ -408,8 +289,156 @@ public class TaskletStep extends AbstractStep {
 		stream.open(ctx);
 	}
 
-	private void rollback(StepExecution stepExecution, TransactionStatus transaction) {
-		transactionManager.rollback(transaction);
-		stepExecution.incrementRollbackCount();
+	private class ChunkTransactionCallback extends TransactionSynchronizationAdapter implements TransactionCallback {
+
+		private final StepExecution stepExecution;
+
+		private final ChunkContext chunkContext;
+
+		private boolean rolledBack = false;
+
+		private Integer oldVersion;
+
+		public ChunkTransactionCallback(ChunkContext chunkContext) {
+			this.chunkContext = chunkContext;
+			this.stepExecution = chunkContext.getStepContext().getStepExecution();
+		}
+
+		@Override
+		public void afterCompletion(int status) {
+			if (status != TransactionSynchronization.STATUS_COMMITTED) {
+				if (oldVersion != null) {
+					// Wah! the commit failed. We need to rescue the step
+					// execution data.
+					stepExecution.setVersion(oldVersion);
+				}				
+			}
+			if (status == TransactionSynchronization.STATUS_UNKNOWN) {
+				rollback(stepExecution);
+				stepExecution.upgradeStatus(BatchStatus.UNKNOWN);
+			}
+		}
+
+		public Object doInTransaction(TransactionStatus status) {
+
+			TransactionSynchronizationManager.registerSynchronization(this);
+
+			RepeatStatus result = RepeatStatus.CONTINUABLE;
+
+			StepContribution contribution = stepExecution.createStepContribution();
+
+			chunkListener.beforeChunk();
+
+			boolean locked = false;
+
+			try {
+
+				try {
+					try {
+						result = tasklet.execute(contribution, chunkContext);
+						if (result == null) {
+							result = RepeatStatus.FINISHED;
+						}
+					}
+					catch (Exception e) {
+						if (transactionAttribute.rollbackOn(e)) {
+							throw e;
+						}
+					}
+
+				}
+				finally {
+
+					// If the step operations are asynchronous then we need
+					// to synchronize changes to the step execution (at a
+					// minimum). Take the lock *before* changing the step
+					// execution.
+					try {
+						semaphore.acquire();
+						locked = true;
+					}
+					catch (InterruptedException e) {
+						stepExecution.setStatus(BatchStatus.STOPPED);
+						Thread.currentThread().interrupt();
+					}
+
+					// In case we need to push it back to its old value
+					// after a commit fails...
+					oldVersion = stepExecution.getVersion();
+
+					// Apply the contribution to the step
+					// even if unsuccessful
+					logger.debug("Applying contribution: " + contribution);
+					stepExecution.apply(contribution);
+
+				}
+
+				stream.update(stepExecution.getExecutionContext());
+
+				try {
+					// Going to attempt a commit. If it fails this flag will
+					// stay false and we can use that later.
+					getJobRepository().updateExecutionContext(stepExecution);
+					stepExecution.incrementCommitCount();
+					/*
+					 * The step execution has to be saved before commit because
+					 * otherwise there is a deadlock between the data source
+					 * pool and the semaphore. As long as only one connection is
+					 * used inside the section of this callback that is locked
+					 * with the semaphore, the deadlock is avoided.
+					 */
+					logger.debug("Saving step execution before commit: " + stepExecution);
+					getJobRepository().update(stepExecution);
+				}
+				catch (Exception e) {
+					throw new FatalException("Fatal failure detected", e);
+				}
+
+			}
+			catch (Error e) {
+				logger.debug("Rollback for Error: " + e.getClass().getName() + ": " + e.getMessage());
+				rollback(stepExecution);
+				throw e;
+
+			}
+			catch (RuntimeException e) {
+				logger.debug("Rollback for RuntimeException: " + e.getClass().getName() + ": " + e.getMessage());
+				rollback(stepExecution);
+				throw e;
+			}
+			catch (Exception e) {
+				logger.debug("Rollback for Exception: " + e.getClass().getName() + ": " + e.getMessage());
+				rollback(stepExecution);
+				// Allow checked exceptions
+				throw new TransactionException(e);
+			}
+			finally {
+				// only release the lock if we acquired it
+				if (locked) {
+					semaphore.release();
+				}
+				locked = false;
+			}
+
+			return result;
+
+		}
+
+		private void rollback(StepExecution stepExecution) {
+			if (!rolledBack) {
+				stepExecution.incrementRollbackCount();
+				rolledBack = true;
+			}
+		}
+
 	}
+
+	private static class TransactionException extends RuntimeException {
+
+		public TransactionException(Exception e) {
+			super(e);
+		}
+
+	}
+
 }

spring-batch-core/src/test/java/org/springframework/batch/core/step/item/TaskletStepExceptionTests.java

@@ -194,11 +194,14 @@ public class TaskletStepExceptionTests {
 	@Test
 	public void testCommitError() throws Exception {
 
-		final RuntimeException exception = new RuntimeException();
 		taskletStep.setTransactionManager(new ResourcelessTransactionManager() {
 			@Override
 			protected void doCommit(DefaultTransactionStatus status) throws TransactionException {
-				throw exception;
+				throw new RuntimeException("bar");
+			}
+			@Override
+			protected void doRollback(DefaultTransactionStatus status) throws TransactionException {
+				throw new RuntimeException("foo");
 			}
 		});
 
@@ -213,7 +216,7 @@ public class TaskletStepExceptionTests {
 		taskletStep.execute(stepExecution);
 		assertEquals(UNKNOWN, stepExecution.getStatus());
 		Throwable e = stepExecution.getFailureExceptions().get(0);
-		assertEquals(exception, e.getCause());
+		assertEquals("foo", e.getMessage());
 	}
 
 	@Test

spring-batch-core/src/test/java/org/springframework/batch/core/step/tasklet/ChunkOrientedStepIntegrationTests.java

@@ -17,7 +17,6 @@ package org.springframework.batch.core.step.tasklet;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
 
 import java.util.Arrays;
 import java.util.Collections;
@@ -115,19 +114,18 @@ public class ChunkOrientedStepIntegrationTests {
 
 		jobRepository.add(stepExecution);
 		step.execute(stepExecution);
-		assertEquals(BatchStatus.UNKNOWN, stepExecution.getStatus());
+		// Exception on commit is not necessarily fatal: it should fail and rollback
+		assertEquals(BatchStatus.FAILED, stepExecution.getStatus());
 		StepExecution lastStepExecution = jobRepository.getLastStepExecution(jobExecution.getJobInstance(), step
 				.getName());
 		assertEquals(lastStepExecution, stepExecution);
 		assertFalse(lastStepExecution == stepExecution);
 
 		// If the StepExecution is not saved after the failure it will be
-		// STARTED instead of UNKNOWN
-		assertEquals(BatchStatus.UNKNOWN, lastStepExecution.getStatus());
-		String msg = stepExecution.getExitStatus().getExitDescription();
-		assertTrue(msg.contains("Fatal failure detected"));
+		// STARTED instead of FAILED
+		assertEquals(BatchStatus.FAILED, lastStepExecution.getStatus());
 		// The original rollback was caused by this one:
-		assertEquals("Simulate commit failure", stepExecution.getFailureExceptions().get(0).getCause().getMessage());
+		assertEquals("Simulate commit failure", stepExecution.getFailureExceptions().get(0).getMessage());
 
 	}
 

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

@@ -673,7 +673,7 @@ public class TaskletStepTests {
 		String msg = stepExecution.getExitStatus().getExitDescription();
 		assertTrue("Message does not contain ResetFailedException: " + msg, msg.contains("ResetFailedException"));
 		// The original rollback was caused by this one:
-		assertEquals("Bar", stepExecution.getFailureExceptions().get(0).getCause().getMessage());
+		assertEquals("Bar", stepExecution.getFailureExceptions().get(0).getMessage());
 	}
 
 	@Test
@@ -682,6 +682,10 @@ public class TaskletStepTests {
 		step.setTransactionManager(new ResourcelessTransactionManager() {
 			protected void doCommit(DefaultTransactionStatus status) throws TransactionException {
 				// Simulate failure on commit
+				throw new RuntimeException("Foo");
+			}
+			@Override
+			protected void doRollback(DefaultTransactionStatus status) throws TransactionException {
 				throw new RuntimeException("Bar");
 			}
 		});
@@ -695,12 +699,10 @@ public class TaskletStepTests {
 		step.execute(stepExecution);
 		assertEquals(BatchStatus.UNKNOWN, stepExecution.getStatus());
 		String msg = stepExecution.getExitStatus().getExitDescription();
-		assertTrue(msg.contains("Fatal failure detected"));
 		Throwable ex = stepExecution.getFailureExceptions().get(0);
 		msg = ex.getMessage();
-		assertTrue(msg.contains("Fatal failure detected"));
-		// The original rollback was caused by this one:
-		assertEquals("Bar", ex.getCause().getMessage());
+		// The original rollback failed because of this one:
+		assertEquals("Bar", ex.getMessage());
 	}
 
 	@Test