OPEN - issue BATCH-771: Refactor Listeners for chunk changes

Add attribute accessor for synchronizing unfinished work between transactions

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

@@ -30,7 +30,7 @@ import org.springframework.batch.repeat.ExitStatus;
 import org.springframework.batch.repeat.RepeatCallback;
 import org.springframework.batch.repeat.RepeatContext;
 import org.springframework.batch.repeat.RepeatOperations;
-import org.springframework.transaction.support.TransactionSynchronizationManager;
+import org.springframework.core.AttributeAccessor;
 
 /**
  * Simplest possible implementation of {@link StepHandler} with no skipping or
@@ -46,6 +46,8 @@ import org.springframework.transaction.support.TransactionSynchronizationManager
  */
 public class ItemOrientedStepHandler<T, S> implements StepHandler {
 
+	private static final String ITEM_BUFFER_KEY = ItemOrientedStepHandler.class.getName() + ".ITEM_BUFFER_KEY";
+
 	protected final Log logger = LogFactory.getLog(getClass());
 
 	private final ItemReader<? extends T> itemReader;
@@ -78,11 +80,12 @@ public class ItemOrientedStepHandler<T, S> implements StepHandler {
 	 * {@link ItemProcessor} returns null, the write is omitted and another item
 	 * taken from the reader.
 	 * 
-	 * @see org.springframework.batch.core.step.item.StepHandler#handle(org.springframework.batch.core.StepContribution)
+	 * @see org.springframework.batch.core.step.item.StepHandler#handle(org.springframework.batch.core.StepContribution,
+	 * AttributeAccessor)
 	 */
-	public ExitStatus handle(final StepContribution contribution) throws Exception {
+	public ExitStatus handle(final StepContribution contribution, AttributeAccessor attributes) throws Exception {
 
-		final List<ReadWrapper<T>> buffer = getItemBuffer();
+		final List<ReadWrapper<T>> buffer = getItemBuffer(attributes);
 
 		ExitStatus result = ExitStatus.CONTINUABLE;
 
@@ -91,10 +94,10 @@ public class ItemOrientedStepHandler<T, S> implements StepHandler {
 			result = repeatOperations.iterate(new RepeatCallback() {
 				public ExitStatus doInIteration(final RepeatContext context) throws Exception {
 					ReadWrapper<T> item = read(contribution);
-					if (item == null) {
+					contribution.incrementReadSkipCount(item.getSkipCount());
+					if (item.getItem() == null) {
 						return ExitStatus.FINISHED;
 					}
-					contribution.incrementReadSkipCount(item.getSkipCount());
 					buffer.add(item);
 					return ExitStatus.CONTINUABLE;
 				}
@@ -127,19 +130,28 @@ public class ItemOrientedStepHandler<T, S> implements StepHandler {
 		for (S data : processed) {
 			write(data, contribution);
 		}
-		buffer.clear();
+
+		// On successful completion clear the attributes to signal that there is
+		// no more processing
+		clearAll(attributes);
 
 		logger.info("Contribution: " + contribution);
 		return result;
 
 	}
 
-	private List<ReadWrapper<T>> getItemBuffer() {
-		if (!TransactionSynchronizationManager.hasResource(this)) {
-			TransactionSynchronizationManager.bindResource(this, new ArrayList<ReadWrapper<T>>());
+	private void clearAll(AttributeAccessor attributes) {
+		for (String key : attributes.attributeNames()) {
+			attributes.removeAttribute(key);
+		}
+	}
+
+	private List<ReadWrapper<T>> getItemBuffer(AttributeAccessor attributes) {
+		if (!attributes.hasAttribute(ITEM_BUFFER_KEY)) {
+			attributes.setAttribute(ITEM_BUFFER_KEY, new ArrayList<ReadWrapper<T>>());
 		}
 		@SuppressWarnings("unchecked")
-		List<ReadWrapper<T>> resource = (List<ReadWrapper<T>>) TransactionSynchronizationManager.getResource(this);
+		List<ReadWrapper<T>> resource = (List<ReadWrapper<T>>) attributes.getAttribute(ITEM_BUFFER_KEY);
 		return resource;
 	}
 
@@ -148,8 +160,7 @@ public class ItemOrientedStepHandler<T, S> implements StepHandler {
 	 * @return next item for writing
 	 */
 	protected ReadWrapper<T> read(StepContribution contribution) throws Exception {
-		T item = doRead();
-		return item==null ? null : new ReadWrapper<T>(item);
+		return new ReadWrapper<T>(doRead());
 	}
 
 	/**

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

@@ -374,8 +374,7 @@ public class SkipLimitStepFactoryBean<T, S> extends SimpleStepFactoryBean<T, S>
 
 			while (true) {
 				try {
-					T item = doRead();
-					return item==null ? null : new ReadWrapper<T>(item, skipCount);
+					return new ReadWrapper<T>(doRead(), skipCount);
 				}
 				catch (Exception e) {
 					try {

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

@@ -19,6 +19,7 @@ import org.springframework.batch.core.StepContribution;
 import org.springframework.batch.item.ItemReader;
 import org.springframework.batch.item.ItemWriter;
 import org.springframework.batch.repeat.ExitStatus;
+import org.springframework.core.AttributeAccessor;
 
 /**
  * Strategy for processing in a step. Bears a resemblance to {@link ItemReader}
@@ -37,10 +38,12 @@ public interface StepHandler {
 	 * not null process the item and return {@link ExitStatus#CONTINUABLE}. On
 	 * failure throws an exception.
 	 * 
-	 * @param contribution the current step context
+	 * @param contribution mutable state to be passed back to update the current
+	 * step execution
+	 * @param attributes attributes shared between invocations
 	 * @return an {@link ExitStatus} indicating whether processing is
 	 * continuable.
 	 */
-	ExitStatus handle(StepContribution contribution) throws Exception;
+	ExitStatus handle(StepContribution contribution, AttributeAccessor attributes) throws Exception;
 
 }

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

@@ -15,6 +15,9 @@
  */
 package org.springframework.batch.core.step.item;
 
+import java.util.Queue;
+import java.util.concurrent.LinkedBlockingQueue;
+
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.springframework.batch.core.BatchStatus;
@@ -38,6 +41,8 @@ import org.springframework.batch.repeat.RepeatCallback;
 import org.springframework.batch.repeat.RepeatContext;
 import org.springframework.batch.repeat.RepeatOperations;
 import org.springframework.batch.repeat.support.RepeatTemplate;
+import org.springframework.core.AttributeAccessor;
+import org.springframework.core.AttributeAccessorSupport;
 import org.springframework.transaction.PlatformTransactionManager;
 import org.springframework.transaction.TransactionStatus;
 import org.springframework.transaction.interceptor.DefaultTransactionAttribute;
@@ -206,12 +211,16 @@ public class StepHandlerStep extends AbstractStep {
 		stream.update(stepExecution.getExecutionContext());
 		getJobRepository().updateExecutionContext(stepExecution);
 
-		final ExceptionHolder fatalException = new ExceptionHolder();
-
 		return stepOperations.iterate(new RepeatCallback() {
 
+			final Queue<AttributeAccessor> attributeQueue = new LinkedBlockingQueue<AttributeAccessor>();
+
 			public ExitStatus doInIteration(RepeatContext context) throws Exception {
-				final StepContribution contribution = stepExecution.createStepContribution();
+
+				ExceptionHolder fatalException = new ExceptionHolder();
+
+				StepContribution contribution = stepExecution.createStepContribution();
+
 				// Before starting a new transaction, check for
 				// interruption.
 				interruptionPolicy.checkInterrupted(stepExecution);
@@ -222,10 +231,16 @@ public class StepHandlerStep extends AbstractStep {
 
 				boolean locked = false;
 
+				AttributeAccessor attributes = attributeQueue.poll();
+				if (attributes == null) {
+					attributes = new AttributeAccessorSupport() {
+					};
+				}
+
 				try {
 
 					try {
-						exitStatus = itemHandler.handle(contribution);
+						exitStatus = itemHandler.handle(contribution, attributes);
 					}
 					catch (Error e) {
 						if (transactionAttribute.rollbackOn(e)) {
@@ -237,6 +252,13 @@ public class StepHandlerStep extends AbstractStep {
 							throw e;
 						}
 					}
+					finally {
+						// Still some stuff to do with the data in this chunk,
+						// pass it back
+						if (attributes.attributeNames().length > 0) {
+							attributeQueue.add(attributes);
+						}
+					}
 
 					contribution.incrementCommitCount();
 

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

@@ -33,8 +33,10 @@ import org.springframework.batch.item.NoWorkFoundException;
 import org.springframework.batch.item.ParseException;
 import org.springframework.batch.item.UnexpectedInputException;
 import org.springframework.batch.item.support.PassthroughItemProcessor;
+import org.springframework.batch.repeat.context.RepeatContextSupport;
 import org.springframework.batch.repeat.policy.SimpleCompletionPolicy;
 import org.springframework.batch.repeat.support.RepeatTemplate;
+import org.springframework.core.AttributeAccessor;
 
 /**
  * @author Dave Syer
@@ -48,6 +50,8 @@ public class ItemOrientedStepHandlerTests {
 	
 	private RepeatTemplate repeatTemplate = new RepeatTemplate();
 	
+	private AttributeAccessor context = new RepeatContextSupport(null);
+	
 	@Before
 	public void setUp() {
 		repeatTemplate.setCompletionPolicy(new SimpleCompletionPolicy(2));
@@ -59,7 +63,7 @@ public class ItemOrientedStepHandlerTests {
 				new PassthroughItemProcessor<String>(), itemWriter, repeatTemplate);
 		StepContribution contribution = new StepContribution(new StepExecution("foo", new JobExecution(new JobInstance(
 				123L, new JobParameters(), "job"))));
-		handler.handle(contribution);
+		handler.handle(contribution, context);
 		assertEquals(2, itemReader.count);
 		assertEquals("12", itemWriter.values);
 	}
@@ -70,7 +74,7 @@ public class ItemOrientedStepHandlerTests {
 				new AgrgegateItemProcessor(), itemWriter, repeatTemplate);
 		StepContribution contribution = new StepContribution(new StepExecution("foo", new JobExecution(new JobInstance(
 				123L, new JobParameters(), "job"))));
-		handler.handle(contribution);
+		handler.handle(contribution, context);
 		assertEquals(2, itemReader.count);
 		assertEquals("12", itemWriter.values);
 	}

spring-batch-infrastructure/src/test/java/org/springframework/batch/item/database/JdbcPagingItemReaderCommonTests.java

@@ -1,7 +1,6 @@
 package org.springframework.batch.item.database;
 
 import org.junit.runner.RunWith;
-import org.junit.Test;
 import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.batch.item.CommonItemStreamItemReaderTests;
@@ -11,7 +10,6 @@ import org.springframework.batch.item.sample.Foo;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.jdbc.core.simple.ParameterizedRowMapper;
 
-import javax.persistence.EntityManagerFactory;
 import javax.sql.DataSource;
 import java.sql.ResultSet;
 import java.sql.SQLException;