RESOLVED - issue BATCH-1272: Write skips do not work in a multi-threaded step

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

spring-batch-core/src/main/java/org/springframework/batch/core/scope/context/StepContextRepeatCallback.java

@@ -17,7 +17,10 @@ package org.springframework.batch.core.scope.context;
 
 import java.util.Queue;
 import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.atomic.AtomicInteger;
 
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
 import org.springframework.batch.core.Step;
 import org.springframework.batch.core.StepExecution;
 import org.springframework.batch.repeat.RepeatCallback;
@@ -34,9 +37,13 @@ import org.springframework.batch.repeat.RepeatStatus;
 public abstract class StepContextRepeatCallback implements RepeatCallback {
 
 	private final Queue<ChunkContext> attributeQueue = new LinkedBlockingQueue<ChunkContext>();
+	
+	private final AtomicInteger workerCount = new AtomicInteger(0);
 
 	private final StepExecution stepExecution;
 
+	private final Log logger = LogFactory.getLog(StepContextRepeatCallback.class);
+
 	/**
 	 * @param stepExecution
 	 */
@@ -45,14 +52,18 @@ public abstract class StepContextRepeatCallback implements RepeatCallback {
 	}
 
 	/**
-	 * Manage the {@link StepContext} lifecycle to ensure that the current
-	 * thread has a reference to the context, even if the callback is executed
-	 * in a pooled thread. Handles the registration and de-registration of the
-	 * step context, so clients should not duplicate those calls.
+	 * Manage the {@link StepContext} lifecycle. Business processing should be
+	 * delegated to {@link #doInChunkContext(RepeatContext, ChunkContext)}. This
+	 * is to ensure that the current thread has a reference to the context, even
+	 * if the callback is executed in a pooled thread. Handles the registration
+	 * and de-registration of the step context, so clients should not duplicate
+	 * those calls.
 	 * 
 	 * @see RepeatCallback#doInIteration(RepeatContext)
 	 */
 	public RepeatStatus doInIteration(RepeatContext context) throws Exception {
+		
+		workerCount.incrementAndGet();
 
 		// The StepContext has to be the same for all chunks,
 		// otherwise step-scoped beans will be re-initialised for each chunk.
@@ -64,9 +75,12 @@ public abstract class StepContextRepeatCallback implements RepeatCallback {
 		}
 
 		try {
-			return doInChunkContext(context, chunkContext);
+			logger.debug("Chunk execution starting: worker count="+workerCount.get()+", queue size="+attributeQueue.size());
+			return RepeatStatus.continueIf(doInChunkContext(context, chunkContext).isContinuable()
+					|| (attributeQueue.isEmpty() && workerCount.get()>1));
 		}
 		finally {
+			workerCount.decrementAndGet();
 			// Still some stuff to do with the data in this chunk,
 			// pass it back.
 			if (!chunkContext.isComplete()) {
@@ -77,14 +91,14 @@ public abstract class StepContextRepeatCallback implements RepeatCallback {
 	}
 
 	/**
-	 * Do the work required for this chunk of the step. The
-	 * {@link ChunkContext} provided is managed by the base class, so that if
-	 * there is still work to do for the task in hand state can be stored here.
-	 * In a multi-threaded client, the base class ensures that only one thread
-	 * at a time can be working on each instance of {@link ChunkContext}. Workers
-	 * should signal that they are finished with a context by removing all the
-	 * attributes they have added. If a worker does not remove them another
-	 * thread might see stale state.
+	 * Do the work required for this chunk of the step. The {@link ChunkContext}
+	 * provided is managed by the base class, so that if there is still work to
+	 * do for the task in hand state can be stored here. In a multi-threaded
+	 * client, the base class ensures that only one thread at a time can be
+	 * working on each instance of {@link ChunkContext}. Workers should signal
+	 * that they are finished with a context by removing all the attributes they
+	 * have added. If a worker does not remove them another thread might see
+	 * stale state.
 	 * 
 	 * @param context the current {@link RepeatContext}
 	 * @param chunkContext the chunk context in which to carry out the work

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

@@ -361,7 +361,7 @@ public class FaultTolerantChunkProcessor<I, O> extends SimpleChunkProcessor<I, O
 	private void scan(final StepContribution contribution, final Chunk<I> inputs, final Chunk<O> outputs,
 			ChunkMonitor chunkMonitor) throws Exception {
 
-		logger.debug("Scanning for failed item on write.");
+		logger.debug("Scanning for failed item on write: "+inputs);
 		if (outputs.isEmpty()) {
 			inputs.setBusy(false);
 			return;

spring-batch-core/src/test/java/org/springframework/batch/core/scope/context/StepContextRepeatCallbackTests.java

@@ -15,10 +15,13 @@
  */
 package org.springframework.batch.core.scope.context;
 
+import static org.junit.Assert.fail;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
+import java.util.concurrent.CountDownLatch;
+
 import org.junit.After;
 import org.junit.Test;
 import org.springframework.batch.core.ExitStatus;
@@ -28,6 +31,7 @@ import org.springframework.batch.core.scope.context.StepContextRepeatCallback;
 import org.springframework.batch.core.scope.context.StepSynchronizationManager;
 import org.springframework.batch.repeat.RepeatContext;
 import org.springframework.batch.repeat.RepeatStatus;
+import org.springframework.batch.repeat.context.RepeatContextSupport;
 
 /**
  * @author Dave Syer
@@ -58,7 +62,7 @@ public class StepContextRepeatCallbackTests {
 	}
 
 	@Test
-	public void testUnfinishedWork() throws Exception {
+	public void testAddingAttributes() throws Exception {
 		StepSynchronizationManager.register(stepExecution);
 		StepContextRepeatCallback callback = new StepContextRepeatCallback(stepExecution) {
 			@Override
@@ -73,11 +77,37 @@ public class StepContextRepeatCallbackTests {
 				return RepeatStatus.FINISHED;
 			}
 		};
-		callback.doInIteration(null);
+		assertEquals(RepeatStatus.FINISHED, callback.doInIteration(null));		
 		assertTrue(addedAttribute);
 		callback.doInIteration(null);
 		assertTrue(removedAttribute);
 		callback.doInIteration(null);
 		assertFalse(removedAttribute);
 	}
+
+	@Test
+	public void testUnfinishedWork() throws Exception {
+		StepSynchronizationManager.register(stepExecution);
+		final CountDownLatch latch = new CountDownLatch(2);
+		final StepContextRepeatCallback callback = new StepContextRepeatCallback(stepExecution) {
+			@Override
+			public RepeatStatus doInChunkContext(RepeatContext context, ChunkContext chunkContext) throws Exception {
+				if (context==null) latch.await();
+				return RepeatStatus.FINISHED;
+			}
+		};
+		new Thread(new Runnable() {
+			public void run() {
+				try {
+					callback.doInIteration(null);
+				}
+				catch (Exception e) {
+					fail(e.getMessage());
+				}
+			}
+		}).start();
+		assertEquals(RepeatStatus.CONTINUABLE, callback.doInIteration(new RepeatContextSupport(null)));
+		latch.countDown();
+	}
+
 }

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

@@ -12,7 +12,6 @@ import java.util.List;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.junit.Before;
-import org.junit.Ignore;
 import org.junit.Test;
 import org.springframework.batch.core.BatchStatus;
 import org.springframework.batch.core.JobExecution;
@@ -358,7 +357,6 @@ public class FaultTolerantStepFactoryBeanRollbackTests {
 	}
 
 	@Test
-	@Ignore
 	public void testMultithreadedSkipInWriter() throws Exception {
 		writer.setFailures("1", "2", "3", "4", "5");
 		factory.setCommitInterval(3);