RESOLVED - issue BATCH-409: StatefulRetryStepFactoryBean needs to co-ordinate exception handler with retry policy

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

Use a special exception handler that is aware of the retry.  Only works in the nested RepeatContext created by the ItemOrientedStep.

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

@@ -23,6 +23,7 @@ import org.springframework.batch.execution.step.ItemOrientedStep;
 import org.springframework.batch.item.ItemReader;
 import org.springframework.batch.item.ItemStream;
 import org.springframework.batch.item.ItemWriter;
+import org.springframework.batch.repeat.exception.handler.ExceptionHandler;
 import org.springframework.batch.repeat.exception.handler.SimpleLimitExceptionHandler;
 import org.springframework.batch.repeat.policy.SimpleCompletionPolicy;
 import org.springframework.batch.repeat.support.RepeatTemplate;
@@ -54,6 +55,8 @@ public class DefaultStepFactoryBean extends AbstractStepFactoryBean {
 
 	private RepeatTemplate stepOperations;
 
+	private SimpleLimitExceptionHandler exceptionHandler;
+
 	/**
 	 * Set the commit interval.
 	 * 
@@ -107,6 +110,22 @@ public class DefaultStepFactoryBean extends AbstractStepFactoryBean {
 		return stepOperations;
 	}
 
+	/**
+	 * Public setter for the SimpleLimitExceptionHandler.
+	 * @param exceptionHandler the exceptionHandler to set
+	 */
+	public void setExceptionHandler(SimpleLimitExceptionHandler exceptionHandler) {
+		this.exceptionHandler = exceptionHandler;
+	}
+
+	/**
+	 * Protected getter for the {@link ExceptionHandler}.
+	 * @return the {@link ExceptionHandler}
+	 */
+	protected ExceptionHandler getExceptionHandler() {
+		return exceptionHandler;
+	}
+
 	/**
 	 * Public setter for the {@link TaskExecutor}. If this is set, then it will
 	 * be used to execute the chunk processing inside the {@link Step}.
@@ -209,7 +228,8 @@ public class DefaultStepFactoryBean extends AbstractStepFactoryBean {
 			 * will never re-appear after a rollback.
 			 */
 			itemProcessor.setItemSkipPolicy(new LimitCheckingItemSkipPolicy(skipLimit));
-			stepOperations.setExceptionHandler(new SimpleLimitExceptionHandler(skipLimit));
+			setExceptionHandler(new SimpleLimitExceptionHandler(skipLimit));
+			stepOperations.setExceptionHandler(getExceptionHandler());
 			step.setStepOperations(stepOperations);
 		}
 		else {

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

@@ -0,0 +1,87 @@
+/*
+ * 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.execution.step.support;
+
+import org.springframework.batch.repeat.RepeatContext;
+import org.springframework.batch.repeat.exception.handler.ExceptionHandler;
+import org.springframework.batch.repeat.synch.RepeatSynchronizationManager;
+import org.springframework.batch.retry.RetryCallback;
+import org.springframework.batch.retry.RetryContext;
+import org.springframework.batch.retry.RetryPolicy;
+import org.springframework.batch.retry.listener.RetryListenerSupport;
+
+/**
+ * @author Dave Syer
+ * 
+ */
+public class SimpleRetryExceptionHandler extends RetryListenerSupport implements ExceptionHandler {
+
+	/**
+	 * Attribute key, whose existence signals an exhausted retry.
+	 */
+	private static final String EXHAUSTED = SimpleRetryExceptionHandler.class.getName() + ".RETRY_EXHAUSTED";
+
+	private RetryPolicy retryPolicy;
+
+	private ExceptionHandler exceptionHandler;
+
+	/**
+	 * @param retryPolicy
+	 * @param exceptionHandler
+	 */
+	public SimpleRetryExceptionHandler(RetryPolicy retryPolicy, ExceptionHandler exceptionHandler) {
+		this.retryPolicy = retryPolicy;
+		this.exceptionHandler = exceptionHandler;
+	}
+
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.batch.repeat.exception.handler.ExceptionHandler#handleException(org.springframework.batch.repeat.RepeatContext,
+	 * java.lang.Throwable)
+	 */
+	public void handleException(RepeatContext context, Throwable throwable) throws RuntimeException {
+		// Only bother to check the delegate exception handler if we know that
+		// retry is exhausted
+		if (context.hasAttribute(EXHAUSTED)) {
+			exceptionHandler.handleException(context, throwable);
+		}
+	}
+
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.batch.retry.RetryListener#close(org.springframework.batch.retry.RetryContext,
+	 * org.springframework.batch.retry.RetryCallback, java.lang.Throwable)
+	 */
+	public void close(RetryContext context, RetryCallback callback, Throwable throwable) {
+		if (!retryPolicy.canRetry(context)) {
+			getRepeatContext().setAttribute(EXHAUSTED, "true");
+		}
+	}
+
+	/**
+	 * Get the parent context (the retry is in an inner "chunk" loop and we want
+	 * the exception to be handled at the outer "step" level).
+	 * @return the {@link RepeatContext} that should hold the exhausted flag.
+	 */
+	private RepeatContext getRepeatContext() {
+		RepeatContext context = RepeatSynchronizationManager.getContext();
+		if (context.getParent() != null) {
+			return context.getParent();
+		}
+		return context;
+	}
+
+}

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

@@ -23,7 +23,6 @@ import org.springframework.batch.item.ItemReader;
 import org.springframework.batch.item.ItemRecoverer;
 import org.springframework.batch.item.ItemWriter;
 import org.springframework.batch.repeat.ExitStatus;
-import org.springframework.batch.repeat.exception.handler.SimpleLimitExceptionHandler;
 import org.springframework.batch.retry.RetryOperations;
 import org.springframework.batch.retry.RetryPolicy;
 import org.springframework.batch.retry.backoff.BackOffPolicy;
@@ -49,15 +48,16 @@ public class StatefulRetryStepFactoryBean extends DefaultStepFactoryBean {
 	private ItemKeyGenerator itemKeyGenerator;
 
 	private ItemRecoverer itemRecoverer;
-	
+
 	private int retryLimit;
-	
+
 	private Class[] retryableExceptionClasses;
-	
+
 	private BackOffPolicy backOffPolicy;
 
 	/**
-	 * Public setter for the retry limit.  Each item can be retried up to this limit.
+	 * Public setter for the retry limit. Each item can be retried up to this
+	 * limit.
 	 * @param retryLimit the retry limit to set
 	 */
 	public void setRetryLimit(int retryLimit) {
@@ -114,16 +114,16 @@ public class StatefulRetryStepFactoryBean extends DefaultStepFactoryBean {
 
 		super.applyConfiguration(step);
 
-		if (retryLimit>0) {
-			
+		if (retryLimit > 0) {
+
 			SimpleRetryPolicy retryPolicy = new SimpleRetryPolicy(retryLimit);
-			if (retryableExceptionClasses!=null) {
+			if (retryableExceptionClasses != null) {
 				retryPolicy.setRetryableExceptionClasses(retryableExceptionClasses);
 			}
 
-			// TODO: actually we need to co-ordinate the retry policy with the
-			// exception handler limit, so this is a hack for now.
-			getStepOperations().setExceptionHandler(new SimpleLimitExceptionHandler(Integer.MAX_VALUE));
+			// Co-ordinate the retry policy with the exception handler:
+			getStepOperations()
+					.setExceptionHandler(new SimpleRetryExceptionHandler(retryPolicy, getExceptionHandler()));
 
 			ItemReaderRetryCallback retryCallback = new ItemReaderRetryCallback(getItemReader(), itemKeyGenerator,
 					getItemWriter());
@@ -132,7 +132,7 @@ public class StatefulRetryStepFactoryBean extends DefaultStepFactoryBean {
 
 			RetryTemplate retryTemplate = new RetryTemplate();
 			retryTemplate.setRetryPolicy(itemProviderRetryPolicy);
-			if (backOffPolicy!=null) {
+			if (backOffPolicy != null) {
 				retryTemplate.setBackOffPolicy(backOffPolicy);
 			}
 

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

@@ -0,0 +1,112 @@
+/*
+ * 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.execution.step.support;
+
+import junit.framework.TestCase;
+
+import org.springframework.batch.repeat.RepeatContext;
+import org.springframework.batch.repeat.context.RepeatContextSupport;
+import org.springframework.batch.repeat.exception.handler.SimpleLimitExceptionHandler;
+import org.springframework.batch.repeat.synch.RepeatSynchronizationManager;
+import org.springframework.batch.retry.RetryContext;
+import org.springframework.batch.retry.RetryPolicy;
+import org.springframework.batch.retry.policy.AlwaysRetryPolicy;
+import org.springframework.batch.retry.policy.NeverRetryPolicy;
+
+/**
+ * @author Dave Syer
+ * 
+ */
+public class SimpleRetryExceptionHandlerTests extends TestCase {
+
+	private RepeatContext context = new RepeatContextSupport(new RepeatContextSupport(null));
+
+	/*
+	 * (non-Javadoc)
+	 * @see junit.framework.TestCase#setUp()
+	 */
+	protected void setUp() throws Exception {
+		RepeatSynchronizationManager.register(context);
+	}
+	
+	/* (non-Javadoc)
+	 * @see junit.framework.TestCase#tearDown()
+	 */
+	protected void tearDown() throws Exception {
+		RepeatSynchronizationManager.clear();
+	}
+
+	/**
+	 * Test method for
+	 * {@link org.springframework.batch.execution.step.support.SimpleRetryExceptionHandler#handleException(org.springframework.batch.repeat.RepeatContext, java.lang.Throwable)}.
+	 */
+	public void testRethrowWhenRetryExhausted() {
+
+		RetryPolicy retryPolicy = new NeverRetryPolicy();
+		RuntimeException ex = new RuntimeException("foo");
+
+		SimpleRetryExceptionHandler handler = getHandlerAfterRetry(retryPolicy, ex);
+
+		// Then pretend to handle the exception in the parent context...
+		try {
+			handler.handleException(context.getParent(), ex);
+			fail("Expected RuntimeException");
+		}
+		catch (RuntimeException e) {
+			assertEquals(ex, e);
+		}
+
+		assertEquals(0, context.attributeNames().length);
+		// One for the retry exhausted flag and one for the counter in the delegate exception handler
+		assertEquals(2, context.getParent().attributeNames().length);
+	}
+
+	/**
+	 * Test method for
+	 * {@link org.springframework.batch.execution.step.support.SimpleRetryExceptionHandler#handleException(org.springframework.batch.repeat.RepeatContext, java.lang.Throwable)}.
+	 */
+	public void testNoRethrowWhenRetryNotExhausted() {
+
+		RetryPolicy retryPolicy = new AlwaysRetryPolicy();
+		RuntimeException ex = new RuntimeException("foo");
+
+		SimpleRetryExceptionHandler handler = getHandlerAfterRetry(retryPolicy, ex);
+
+		// Then pretend to handle the exception in the parent context...
+		handler.handleException(context.getParent(), ex);
+		
+		assertEquals(0, context.attributeNames().length);
+		assertEquals(0, context.getParent().attributeNames().length);
+	}
+
+	/**
+	 * @param retryPolicy
+	 * @param ex
+	 * @return
+	 */
+	private SimpleRetryExceptionHandler getHandlerAfterRetry(RetryPolicy retryPolicy, RuntimeException ex) {
+		
+		// Always rethrow if the retry is exhausted
+		SimpleRetryExceptionHandler handler = new SimpleRetryExceptionHandler(retryPolicy, new SimpleLimitExceptionHandler(0));
+
+		// Simulate a failed retry...
+		RetryContext retryContext = retryPolicy.open(null);
+		retryPolicy.registerThrowable(retryContext, ex);
+		handler.close(retryContext, null, ex);
+		return handler;
+	}
+
+}