BATCHADM-49: loop back step contributions through writer

spring-batch-integration/src/main/java/org/springframework/batch/integration/chunk/ChunkMessageChannelItemWriter.java

@@ -16,7 +16,11 @@
 
 package org.springframework.batch.integration.chunk;
 
+import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
+import java.util.Queue;
+import java.util.concurrent.LinkedBlockingQueue;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -32,7 +36,8 @@ import org.springframework.batch.item.ItemWriter;
 import org.springframework.integration.gateway.MessagingGateway;
 import org.springframework.util.Assert;
 
-public class ChunkMessageChannelItemWriter<T> extends StepExecutionListenerSupport implements ItemWriter<T>, ItemStream {
+public class ChunkMessageChannelItemWriter<T> extends StepExecutionListenerSupport implements ItemWriter<T>,
+		ItemStream, StepContributionSource {
 
 	private static final Log logger = LogFactory.getLog(ChunkMessageChannelItemWriter.class);
 
@@ -57,7 +62,7 @@ public class ChunkMessageChannelItemWriter<T> extends StepExecutionListenerSuppo
 	 * result from the remote workers. This is a multiplier on the receive
 	 * timeout set separately on the gateway. The ideal value is a compromise
 	 * between allowing slow workers time to finish, and responsiveness if there
-	 * is a dead worker.  Defaults to 40.
+	 * is a dead worker. Defaults to 40.
 	 * 
 	 * @param maxWaitTimeouts the maximum number of wait timeouts
 	 */
@@ -119,6 +124,11 @@ public class ChunkMessageChannelItemWriter<T> extends StepExecutionListenerSuppo
 			stepExecution.setStatus(BatchStatus.FAILED);
 			return ExitStatus.FAILED.addExitDescription(e.getClass().getName() + ": " + e.getMessage());
 		}
+		finally {
+			for (StepContribution contribution : getStepContributions()) {
+				stepExecution.apply(contribution);
+			}
+		}
 		if (timedOut) {
 			stepExecution.setStatus(BatchStatus.FAILED);
 			throw new ItemStreamException("Timed out waiting for back log at end of step");
@@ -145,17 +155,32 @@ public class ChunkMessageChannelItemWriter<T> extends StepExecutionListenerSuppo
 		executionContext.putLong(ACTUAL, localState.actual);
 	}
 
+	public Collection<StepContribution> getStepContributions() {
+		return localState.pollContributions();
+	}
+
 	/**
 	 * Wait until all the results that are in the pipeline come back to the
 	 * reply channel.
 	 * 
 	 * @return true if successfully received a result, false if timed out
 	 */
-	private boolean waitForResults() {
+	private boolean waitForResults() throws AsynchronousFailureException {
 		int count = 0;
 		int maxCount = maxWaitTimeouts;
+		Throwable failure = null;
 		while (localState.getExpecting() > 0 && count++ < maxCount) {
-			getNextResult();
+			try {
+				getNextResult();
+			}
+			catch (Throwable t) {
+				logger.error("Detected error in remote result. Trying to recover " + localState.getExpecting()
+						+ " outstanding results before completing.", t);
+				failure = t;
+			}
+		}
+		if (failure != null) {
+			throw wrapIfNecessary(failure);
 		}
 		return count < maxCount;
 	}
@@ -178,7 +203,7 @@ public class ChunkMessageChannelItemWriter<T> extends StepExecutionListenerSuppo
 			Assert.state(jobInstanceId.equals(localState.getJobId()), "Message contained wrong job instance id ["
 					+ jobInstanceId + "] should have been [" + localState.getJobId() + "].");
 			localState.actual++;
-			// TODO: apply the skip count
+			localState.pushContribution(payload.getStepContribution());
 			if (!payload.isSuccessful()) {
 				throw new AsynchronousFailureException("Failure or interrupt detected in handler: "
 						+ payload.getMessage());
@@ -186,6 +211,22 @@ public class ChunkMessageChannelItemWriter<T> extends StepExecutionListenerSuppo
 		}
 	}
 
+	/**
+	 * Re-throws the original throwable if it is unchecked, wraps checked
+	 * exceptions into {@link AsynchronousFailureException}.
+	 */
+	private static AsynchronousFailureException wrapIfNecessary(Throwable throwable) {
+		if (throwable instanceof Error) {
+			throw (Error) throwable;
+		}
+		else if (throwable instanceof AsynchronousFailureException) {
+			return (AsynchronousFailureException) throwable;
+		}
+		else {
+			return new AsynchronousFailureException("Exception in remote process", throwable);
+		}
+	}
+
 	private static class LocalState {
 		private long actual;
 
@@ -193,10 +234,26 @@ public class ChunkMessageChannelItemWriter<T> extends StepExecutionListenerSuppo
 
 		private StepExecution stepExecution;
 
+		private Queue<StepContribution> contributions = new LinkedBlockingQueue<StepContribution>();
+
 		public long getExpecting() {
 			return expected - actual;
 		}
 
+		public Collection<StepContribution> pollContributions() {
+			Collection<StepContribution> set = new ArrayList<StepContribution>();
+			StepContribution item = contributions.poll();
+			while (item != null) {
+				set.add(item);
+				item = contributions.poll();
+			}
+			return set;
+		}
+
+		public void pushContribution(StepContribution stepContribution) {
+			contributions.add(stepContribution);
+		}
+
 		public StepContribution createStepContribution() {
 			return stepExecution.createStepContribution();
 		}

spring-batch-integration/src/main/java/org/springframework/batch/integration/chunk/ChunkProcessorChunkHandler.java

@@ -18,9 +18,15 @@ package org.springframework.batch.integration.chunk;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
+import org.springframework.batch.core.JobInterruptedException;
 import org.springframework.batch.core.StepContribution;
 import org.springframework.batch.core.step.item.Chunk;
 import org.springframework.batch.core.step.item.ChunkProcessor;
+import org.springframework.batch.core.step.item.FaultTolerantChunkProcessor;
+import org.springframework.batch.core.step.skip.NonSkippableReadException;
+import org.springframework.batch.core.step.skip.SkipLimitExceededException;
+import org.springframework.batch.core.step.skip.SkipListenerFailedException;
+import org.springframework.batch.retry.RetryException;
 import org.springframework.beans.factory.InitializingBean;
 import org.springframework.integration.annotation.MessageEndpoint;
 import org.springframework.integration.annotation.ServiceActivator;
@@ -63,16 +69,66 @@ public class ChunkProcessorChunkHandler<S> implements ChunkHandler<S>, Initializ
 
 		StepContribution stepContribution = chunkRequest.getStepContribution();
 		try {
-			chunkProcessor.process(stepContribution, new Chunk<S>(chunkRequest.getItems()));
+			process(chunkRequest, stepContribution);
 		}
 		catch (Exception e) {
 			logger.debug("Failed chunk", e);
 			return new ChunkResponse(false, chunkRequest.getJobId(), stepContribution, e.getClass().getName() + ": "
 					+ e.getMessage());
+		} catch (Throwable e) {
+			// The handler might throw an Error or other non-exception
+			logger.debug("Failed chunk with non-exception", e);
+			return new ChunkResponse(false, chunkRequest.getJobId(), stepContribution, e.getClass().getName() + ": "
+					+ e.getMessage());			
 		}
 
 		logger.debug("Completed chunk handling with " + stepContribution);
 		return new ChunkResponse(true, chunkRequest.getJobId(), stepContribution);
 
 	}
+
+	/**
+	 * @param chunkRequest the current request
+	 * @param stepContribution the step contribution to update
+	 * @throws Exception if there is a fatal exception
+	 */
+	private void process(ChunkRequest<S> chunkRequest, StepContribution stepContribution) throws Exception {
+
+		Chunk<S> chunk = new Chunk<S>(chunkRequest.getItems());
+
+		if (chunkProcessor instanceof FaultTolerantChunkProcessor<?, ?>) {
+
+			boolean processed = false;
+
+			while (!processed) {
+				try {
+					chunkProcessor.process(stepContribution, chunk);
+					processed = true;
+				}
+				catch (SkipLimitExceededException e) {
+					throw e;
+				}
+				catch (NonSkippableReadException e) {
+					throw e;
+				}
+				catch (SkipListenerFailedException e) {
+					throw e;
+				}
+				catch (RetryException e) {
+					throw e;
+				}
+				catch (JobInterruptedException e) {
+					throw e;
+				}
+				catch (Exception e) {
+					// try again...
+				}
+			}
+
+		}
+		else {
+			chunkProcessor.process(stepContribution, chunk);
+		}
+
+	}
 }

spring-batch-integration/src/main/java/org/springframework/batch/integration/chunk/RemoteChunkHandlerFactoryBean.java

@@ -0,0 +1,216 @@
+/*
+ * Copyright 2006-2010 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.integration.chunk;
+
+import java.lang.reflect.Field;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.springframework.batch.core.StepContribution;
+import org.springframework.batch.core.StepExecutionListener;
+import org.springframework.batch.core.step.item.Chunk;
+import org.springframework.batch.core.step.item.ChunkOrientedTasklet;
+import org.springframework.batch.core.step.item.ChunkProcessor;
+import org.springframework.batch.core.step.item.SimpleChunkProcessor;
+import org.springframework.batch.core.step.tasklet.Tasklet;
+import org.springframework.batch.core.step.tasklet.TaskletStep;
+import org.springframework.batch.item.ItemWriter;
+import org.springframework.batch.item.support.PassThroughItemProcessor;
+import org.springframework.beans.factory.FactoryBean;
+import org.springframework.util.Assert;
+import org.springframework.util.ReflectionUtils;
+
+/**
+ * Convenient factory bean for a chunk handler that also converts an existing
+ * chunk-oriented step into a remote chunk master. The idea is to lift the
+ * existing chunk processor out of a step that works locally, and replace it
+ * with a chunk writer that is already configured to write chunks into a message
+ * channel. The existing step hands its business chunk processing responsibility
+ * over to the handler produced by the factory, which then needs to be set up as
+ * a remote worker on the other end of the channel the chunks are being sent to.
+ * 
+ * @author Dave Syer
+ * 
+ */
+public class RemoteChunkHandlerFactoryBean<T> implements FactoryBean<ChunkHandler<T>> {
+
+	private static Log logger = LogFactory.getLog(RemoteChunkHandlerFactoryBean.class);
+
+	private TaskletStep step;
+
+	private ItemWriter<T> chunkWriter;
+
+	private StepContributionSource stepContributionSource;
+
+	/**
+	 * @param step the step to set
+	 */
+	public void setStep(TaskletStep step) {
+		this.step = step;
+	}
+
+	/**
+	 * @param chunkWriter the chunk writer to set
+	 */
+	public void setChunkWriter(ItemWriter<T> chunkWriter) {
+		this.chunkWriter = chunkWriter;
+	}
+
+	/**
+	 * @param stepContributionSource the step contribution source to set
+	 * (defaults to the chunk writer)
+	 */
+	public void setStepContributionSource(StepContributionSource stepContributionSource) {
+		this.stepContributionSource = stepContributionSource;
+	}
+
+	public Class<?> getObjectType() {
+		return ChunkHandler.class;
+	}
+
+	public boolean isSingleton() {
+		return true;
+	}
+
+	public ChunkHandler<T> getObject() throws Exception {
+
+		if (stepContributionSource == null) {
+			Assert.state(chunkWriter instanceof StepContributionSource,
+					"The chunk writer must be a StepContributionSource or else the source must be provided explicitly");
+			stepContributionSource = (StepContributionSource) chunkWriter;
+		}
+
+		Assert.state(step instanceof TaskletStep, "Step [" + step.getName() + "] must be a TaskletStep");
+		logger.debug("Converting TaskletStep with name=" + step.getName());
+
+		Tasklet tasklet = getTasklet((TaskletStep) step);
+		Assert.state(tasklet instanceof ChunkOrientedTasklet<?>, "Tasklet must be ChunkOrientedTasklet in step="
+				+ step.getName());
+
+		ChunkProcessor<T> chunkProcessor = getChunkProcessor((ChunkOrientedTasklet<?>) tasklet);
+		Assert.state(chunkProcessor != null, "ChunkProcessor must be accessible in Tasklet in step=" + step.getName());
+
+		ItemWriter<T> itemWriter = getItemWriter(chunkProcessor);
+		Assert.state(!(itemWriter instanceof ChunkMessageChannelItemWriter<?>), "Cannot adapt step [" + step.getName()
+				+ "] because it already has a remote chunk writer.  Use a local writer in the step.");
+
+		replaceChunkProcessor((ChunkOrientedTasklet<?>) tasklet, chunkWriter, stepContributionSource);
+		if (chunkWriter instanceof StepExecutionListener) {
+			step.registerStepExecutionListener((StepExecutionListener) chunkWriter);
+		}
+
+		ChunkProcessorChunkHandler<T> handler = new ChunkProcessorChunkHandler<T>();
+		handler.setChunkProcessor(chunkProcessor);
+		// TODO: create step context for the processor in case it has scope="step" dependencies
+		handler.afterPropertiesSet();
+
+		return handler;
+
+	}
+
+	/**
+	 * @param tasklet
+	 * @param chunkWriter
+	 */
+	private void replaceChunkProcessor(ChunkOrientedTasklet<?> tasklet, ItemWriter<T> chunkWriter,
+			final StepContributionSource stepContributionSource) {
+		setField(tasklet, "chunkProcessor", new SimpleChunkProcessor<T, T>(new PassThroughItemProcessor<T>(),
+				chunkWriter) {
+			@Override
+			protected void write(StepContribution contribution, Chunk<T> inputs, Chunk<T> outputs) throws Exception {
+				doWrite(outputs.getItems());
+				// Do not update the step contribution until the chunks are
+				// actually processed
+				updateStepContribution(contribution, stepContributionSource);
+			}
+		});
+	}
+
+	/**
+	 * @param contribution
+	 * @param chunkWriter
+	 */
+	private void updateStepContribution(StepContribution contribution, StepContributionSource stepContributionSource) {
+		for (StepContribution result : stepContributionSource.getStepContributions()) {
+			contribution.incrementFilterCount(result.getFilterCount());
+			contribution.incrementWriteCount(result.getWriteCount());
+			for (int i = 0; i < result.getProcessSkipCount(); i++) {
+				contribution.incrementProcessSkipCount();
+			}
+			for (int i = 0; i < result.getWriteSkipCount(); i++) {
+				contribution.incrementWriteSkipCount();
+			}
+			contribution.setExitStatus(contribution.getExitStatus().and(result.getExitStatus()));
+		}
+	}
+
+	/**
+	 * @param chunkProcessor
+	 * @return
+	 */
+	@SuppressWarnings("unchecked")
+	private ItemWriter<T> getItemWriter(ChunkProcessor<T> chunkProcessor) {
+		return (ItemWriter<T>) getField(chunkProcessor, "itemWriter");
+	}
+
+	/**
+	 * @param tasklet
+	 * @return
+	 */
+	@SuppressWarnings("unchecked")
+	private ChunkProcessor<T> getChunkProcessor(ChunkOrientedTasklet<?> tasklet) {
+		return (ChunkProcessor<T>) getField(tasklet, "chunkProcessor");
+	}
+
+	/**
+	 * @param bean
+	 * @return
+	 */
+	private Tasklet getTasklet(TaskletStep bean) {
+		return (Tasklet) getField(bean, "tasklet");
+	}
+
+	private static Object getField(Object target, String name) {
+		Assert.notNull(target, "Target object must not be null");
+		Field field = ReflectionUtils.findField(target.getClass(), name);
+		if (field == null) {
+			logger.debug("Could not find field [" + name + "] on target [" + target + "]");
+			return null;
+		}
+
+		if (logger.isDebugEnabled()) {
+			logger.debug("Getting field [" + name + "] from target [" + target + "]");
+		}
+		ReflectionUtils.makeAccessible(field);
+		return ReflectionUtils.getField(field, target);
+	}
+
+	private static void setField(Object target, String name, Object value) {
+		Assert.notNull(target, "Target object must not be null");
+		Field field = ReflectionUtils.findField(target.getClass(), name);
+		if (field == null) {
+			throw new IllegalStateException("Could not find field [" + name + "] on target [" + target + "]");
+		}
+
+		if (logger.isDebugEnabled()) {
+			logger.debug("Getting field [" + name + "] from target [" + target + "]");
+		}
+		ReflectionUtils.makeAccessible(field);
+		ReflectionUtils.setField(field, target, value);
+	}
+
+}

spring-batch-integration/src/main/java/org/springframework/batch/integration/chunk/StepContributionSource.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2006-2010 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.integration.chunk;
+
+import java.util.Collection;
+
+import org.springframework.batch.core.StepContribution;
+
+/**
+ * @author Dave Syer
+ *
+ */
+public interface StepContributionSource {
+
+	Collection<StepContribution> getStepContributions();
+
+}