diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/repeat/callback/NestedRepeatCallback.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/repeat/callback/NestedRepeatCallback.java index af5d3b9bb..ee943e982 100644 --- a/spring-batch-infrastructure/src/main/java/org/springframework/batch/repeat/callback/NestedRepeatCallback.java +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/repeat/callback/NestedRepeatCallback.java @@ -16,15 +16,15 @@ package org.springframework.batch.repeat.callback; -import org.springframework.batch.repeat.RepeatStatus; import org.springframework.batch.repeat.RepeatCallback; import org.springframework.batch.repeat.RepeatContext; import org.springframework.batch.repeat.RepeatOperations; +import org.springframework.batch.repeat.RepeatStatus; /** - * Callback that delegates to another callback, via a {@link RepeatOperations} instance. - * Useful when nesting or composing batches in one another, e.g. for breaking a - * batch down into chunks. + * Callback that delegates to another callback, via a {@link RepeatOperations} + * instance. Useful when nesting or composing batches in one another, e.g. for + * breaking a batch down into chunks. * * @author Dave Syer * @@ -38,8 +38,8 @@ public class NestedRepeatCallback implements RepeatCallback { /** * Constructor setting mandatory fields. * - * @param template the {@link RepeatOperations} to use when calling the delegate - * callback + * @param template the {@link RepeatOperations} to use when calling the + * delegate callback * @param callback the {@link RepeatCallback} delegate */ public NestedRepeatCallback(RepeatOperations template, RepeatCallback callback) { diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/repeat/support/TaskExecutorRepeatTemplate.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/repeat/support/TaskExecutorRepeatTemplate.java index 49a25313c..abb5a65d4 100644 --- a/spring-batch-infrastructure/src/main/java/org/springframework/batch/repeat/support/TaskExecutorRepeatTemplate.java +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/repeat/support/TaskExecutorRepeatTemplate.java @@ -51,7 +51,8 @@ public class TaskExecutorRepeatTemplate extends RepeatTemplate { /** * Default limit for maximum number of concurrent unfinished results allowed * by the template. - * {@link #getNextResult(RepeatContext, RepeatCallback, RepeatInternalState)}. + * {@link #getNextResult(RepeatContext, RepeatCallback, RepeatInternalState)} + * . */ public static final int DEFAULT_THROTTLE_LIMIT = 4; @@ -71,10 +72,10 @@ public class TaskExecutorRepeatTemplate extends RepeatTemplate { } /** - * Use the {@link #setTaskExecutor(TaskExecutor)} to generate a result. The internal state in - * this case is a queue of unfinished result holders of type - * {@link ResultHolder}. The holder with the return value should not be on - * the queue when this method exits. The queue is scoped in the calling + * Use the {@link #setTaskExecutor(TaskExecutor)} to generate a result. The + * internal state in this case is a queue of unfinished result holders of + * type {@link ResultHolder}. The holder with the return value should not be + * on the queue when this method exits. The queue is scoped in the calling * method so there is no need to synchronize access. * */ @@ -185,7 +186,7 @@ public class TaskExecutorRepeatTemplate extends RepeatTemplate { * */ private static class ExecutingRunnable implements Runnable, ResultHolder { - + private final RepeatCallback callback; private final RepeatContext context; @@ -226,13 +227,21 @@ public class TaskExecutorRepeatTemplate extends RepeatTemplate { * @see java.lang.Runnable#run() */ public void run() { + boolean clearContext = false; try { + if (RepeatSynchronizationManager.getContext() == null) { + clearContext = true; + RepeatSynchronizationManager.register(context); + } result = callback.doInIteration(context); } catch (Exception e) { error = e; } finally { + if (clearContext) { + RepeatSynchronizationManager.clear(); + } queue.put(this); } } @@ -264,12 +273,12 @@ public class TaskExecutorRepeatTemplate extends RepeatTemplate { /** * @author Dave Syer - * + * */ private static class ResultQueueInternalState extends RepeatInternalStateSupport { private final ResultQueue results; - + /** * @param throttleLimit the throttle limit for the result queue */ @@ -277,8 +286,12 @@ public class TaskExecutorRepeatTemplate extends RepeatTemplate { super(); this.results = new ThrottleLimitResultQueue(throttleLimit); } - /* (non-Javadoc) - * @see org.springframework.batch.repeat.support.RepeatInternalState#getResultQueue() + + /* + * (non-Javadoc) + * + * @seeorg.springframework.batch.repeat.support.RepeatInternalState# + * getResultQueue() */ public ResultQueue getResultQueue() { return results; diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/repeat/support/TaskExecutorRepeatTemplateAsynchronousTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/repeat/support/TaskExecutorRepeatTemplateAsynchronousTests.java index 59b0ac304..e55d8fc3f 100644 --- a/spring-batch-infrastructure/src/test/java/org/springframework/batch/repeat/support/TaskExecutorRepeatTemplateAsynchronousTests.java +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/repeat/support/TaskExecutorRepeatTemplateAsynchronousTests.java @@ -17,7 +17,9 @@ package org.springframework.batch.repeat.support; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNotSame; +import static org.junit.Assert.assertSame; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; @@ -29,12 +31,12 @@ import org.junit.Test; import org.springframework.batch.repeat.RepeatCallback; import org.springframework.batch.repeat.RepeatContext; import org.springframework.batch.repeat.RepeatStatus; +import org.springframework.batch.repeat.callback.NestedRepeatCallback; import org.springframework.batch.repeat.exception.ExceptionHandler; import org.springframework.batch.repeat.policy.SimpleCompletionPolicy; import org.springframework.core.task.SimpleAsyncTaskExecutor; -public class TaskExecutorRepeatTemplateAsynchronousTests extends - AbstractTradeBatchTests { +public class TaskExecutorRepeatTemplateAsynchronousTests extends AbstractTradeBatchTests { RepeatTemplate template = getRepeatTemplate(); @@ -44,6 +46,8 @@ public class TaskExecutorRepeatTemplateAsynchronousTests extends public RepeatTemplate getRepeatTemplate() { TaskExecutorRepeatTemplate template = new TaskExecutorRepeatTemplate(); template.setTaskExecutor(new SimpleAsyncTaskExecutor()); + // Set default completion above number of items in input file + template.setCompletionPolicy(new SimpleCompletionPolicy(8)); return template; } @@ -57,19 +61,19 @@ public class TaskExecutorRepeatTemplateAsynchronousTests extends template.setTaskExecutor(taskExecutor); try { template.iterate(new RepeatCallback() { - public RepeatStatus doInIteration(RepeatContext context) - throws Exception { + public RepeatStatus doInIteration(RepeatContext context) throws Exception { count++; throw new IllegalStateException("foo!"); } }); fail("Expected IllegalStateException"); - } catch (IllegalStateException e) { + } + catch (IllegalStateException e) { assertEquals("foo!", e.getMessage()); } - assertTrue("Too few attempts: "+count, count>=2); - assertTrue("Too many attempts: "+count, count<=10); + assertTrue("Too few attempts: " + count, count >= 1); + assertTrue("Too many attempts: " + count, count <= 10); } @@ -83,19 +87,46 @@ public class TaskExecutorRepeatTemplateAsynchronousTests extends template.setTaskExecutor(taskExecutor); template.setExceptionHandler(new ExceptionHandler() { - public void handleException(RepeatContext context, - Throwable throwable) throws Throwable { + public void handleException(RepeatContext context, Throwable throwable) throws Throwable { count++; } }); template.iterate(new RepeatCallback() { - public RepeatStatus doInIteration(RepeatContext context) - throws Exception { + public RepeatStatus doInIteration(RepeatContext context) throws Exception { throw new IllegalStateException("foo!"); } }); - assertEquals(4, count); + assertTrue("Too few attempts: " + count, count >= 1); + assertTrue("Too many attempts: " + count, count <= 10); + + } + + @Test + public void testNestedSession() throws Exception { + + RepeatTemplate outer = getRepeatTemplate(); + RepeatTemplate inner = new RepeatTemplate(); + + outer.iterate(new NestedRepeatCallback(inner, new RepeatCallback() { + public RepeatStatus doInIteration(RepeatContext context) throws Exception { + count++; + assertNotNull(context); + assertNotSame("Nested batch should have new session", context, context.getParent()); + assertSame(context, RepeatSynchronizationManager.getContext()); + return RepeatStatus.FINISHED; + } + }) { + public RepeatStatus doInIteration(RepeatContext context) throws Exception { + count++; + assertNotNull(context); + assertSame(context, RepeatSynchronizationManager.getContext()); + return super.doInIteration(context); + } + }); + + assertTrue("Too few attempts: " + count, count >= 1); + assertTrue("Too many attempts: " + count, count <= 10); } @@ -113,8 +144,7 @@ public class TaskExecutorRepeatTemplateAsynchronousTests extends final Set threadNames = new HashSet(); final RepeatCallback callback = new RepeatCallback() { - public RepeatStatus doInIteration(RepeatContext context) - throws Exception { + public RepeatStatus doInIteration(RepeatContext context) throws Exception { assertNotSame(threadName, Thread.currentThread().getName()); threadNames.add(Thread.currentThread().getName()); Thread.sleep(100); @@ -145,8 +175,7 @@ public class TaskExecutorRepeatTemplateAsynchronousTests extends final Set threadNames = new HashSet(); final RepeatCallback callback = new RepeatCallback() { - public RepeatStatus doInIteration(RepeatContext context) - throws Exception { + public RepeatStatus doInIteration(RepeatContext context) throws Exception { assertNotSame(threadName, Thread.currentThread().getName()); threadNames.add(Thread.currentThread().getName()); Trade item = provider.read(); @@ -181,10 +210,8 @@ public class TaskExecutorRepeatTemplateAsynchronousTests extends final String threadName = Thread.currentThread().getName(); final Set threadNames = new HashSet(); - final RepeatCallback stepCallback = new ItemReaderRepeatCallback( - provider, processor) { - public RepeatStatus doInIteration(RepeatContext context) - throws Exception { + final RepeatCallback stepCallback = new ItemReaderRepeatCallback(provider, processor) { + public RepeatStatus doInIteration(RepeatContext context) throws Exception { assertNotSame(threadName, Thread.currentThread().getName()); threadNames.add(Thread.currentThread().getName()); Thread.sleep(100); @@ -192,8 +219,7 @@ public class TaskExecutorRepeatTemplateAsynchronousTests extends } }; RepeatCallback jobCallback = new RepeatCallback() { - public RepeatStatus doInIteration(RepeatContext context) - throws Exception { + public RepeatStatus doInIteration(RepeatContext context) throws Exception { stepTemplate.iterate(stepCallback); return RepeatStatus.FINISHED; }