diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/configuration/xml/StepParserStepFactoryBean.java b/spring-batch-core/src/main/java/org/springframework/batch/core/configuration/xml/StepParserStepFactoryBean.java index be2596477..fa1e0efbb 100644 --- a/spring-batch-core/src/main/java/org/springframework/batch/core/configuration/xml/StepParserStepFactoryBean.java +++ b/spring-batch-core/src/main/java/org/springframework/batch/core/configuration/xml/StepParserStepFactoryBean.java @@ -24,6 +24,7 @@ import java.util.LinkedHashSet; import java.util.Map; import java.util.Queue; import java.util.Set; +import java.util.concurrent.locks.ReentrantLock; import javax.batch.api.chunk.listener.RetryProcessListener; import javax.batch.api.chunk.listener.RetryReadListener; @@ -163,6 +164,8 @@ public class StepParserStepFactoryBean implements FactoryBean, BeanNameAwa private Queue partitionQueue; + private ReentrantLock partitionLock; + // // Tasklet Elements // @@ -248,6 +251,15 @@ public class StepParserStepFactoryBean implements FactoryBean, BeanNameAwa this.partitionQueue = queue; } + /** + * Used to coordinate access to the partition queue between the {@link PartitionCollector} and {@link PartitionAnalyzer} + * + * @param lock a lock that will be locked around accessing the partition queue + */ + public void setPartitionLock(ReentrantLock lock) { + this.partitionLock = lock; + } + /** * Create a {@link Step} from the configuration provided. * @@ -474,6 +486,10 @@ public class StepParserStepFactoryBean implements FactoryBean, BeanNameAwa enhanceCommonStep(builder); for (ChunkListener listener : chunkListeners) { + if(listener instanceof PartitionCollectorAdapter) { + ((PartitionCollectorAdapter) listener).setPartitionLock(partitionLock); + } + builder.listener(listener); } diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/configuration/xml/PartitionParser.java b/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/configuration/xml/PartitionParser.java index 376a45c6e..5cedbdbd3 100644 --- a/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/configuration/xml/PartitionParser.java +++ b/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/configuration/xml/PartitionParser.java @@ -17,6 +17,7 @@ package org.springframework.batch.core.jsr.configuration.xml; import java.util.List; import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.locks.ReentrantLock; import org.springframework.batch.core.jsr.configuration.support.BatchArtifact.BatchArtifactType; import org.springframework.batch.core.jsr.partition.JsrPartitionHandler; @@ -55,14 +56,17 @@ public class PartitionParser { private static final String LISTENERS_PROPERTY = "listeners"; private static final String THREADS_PROPERTY = "threads"; private static final String PARTITIONS_PROPERTY = "partitions"; + private static final String PARTITION_LOCK_PROPERTY = "partitionLock"; private final String name; + private boolean allowStartIfComplete = false; /** * @param stepName the name of the step that is being partitioned */ - public PartitionParser(String stepName) { + public PartitionParser(String stepName, boolean allowStartIfComplete) { this.name = stepName; + this.allowStartIfComplete = allowStartIfComplete; } public void parse(Element element, AbstractBeanDefinition bd, ParserContext parserContext, String stepName) { @@ -74,6 +78,8 @@ public class PartitionParser { MutablePropertyValues properties = partitionHandlerDefinition.getPropertyValues(); properties.addPropertyValue(PARTITION_CONTEXT_PROPERTY, new RuntimeBeanReference("batchPropertyContext")); + properties.addPropertyValue("jobRepository", new RuntimeBeanReference("jobRepository")); + properties.addPropertyValue("allowStartIfComplete", allowStartIfComplete); paserMapperElement(element, parserContext, properties); parsePartitionPlan(element, parserContext, stepName, properties); @@ -85,7 +91,6 @@ public class PartitionParser { String partitionHandlerBeanName = name + ".partitionHandler"; registry.registerBeanDefinition(partitionHandlerBeanName, partitionHandlerDefinition); factoryBeanProperties.add("partitionHandler", new RuntimeBeanReference(partitionHandlerBeanName)); - } private void parseCollectorElement(Element element, @@ -98,7 +103,9 @@ public class PartitionParser { // Only needed if a collector is used registerCollectorAnalyzerQueue(parserContext); properties.add(PARTITION_QUEUE_PROPERTY, new RuntimeBeanReference(name + "PartitionQueue")); + properties.add(PARTITION_LOCK_PROPERTY, new RuntimeBeanReference(name + "PartitionLock")); factoryBeanProperties.add("partitionQueue", new RuntimeBeanReference(name + "PartitionQueue")); + factoryBeanProperties.add("partitionLock", new RuntimeBeanReference(name + "PartitionLock")); String collectorName = collectorElement.getAttribute(REF); factoryBeanProperties.add(LISTENERS_PROPERTY, new RuntimeBeanReference(collectorName)); new PropertyParser(collectorName, parserContext, BatchArtifactType.STEP_ARTIFACT, name).parseProperties(collectorElement); @@ -142,8 +149,11 @@ public class PartitionParser { private void registerCollectorAnalyzerQueue(ParserContext parserContext) { AbstractBeanDefinition partitionQueueDefinition = BeanDefinitionBuilder.genericBeanDefinition(ConcurrentLinkedQueue.class) .getBeanDefinition(); + AbstractBeanDefinition partitionLockDefinition = BeanDefinitionBuilder.genericBeanDefinition(ReentrantLock.class) + .getBeanDefinition(); parserContext.getRegistry().registerBeanDefinition(name + "PartitionQueue", partitionQueueDefinition); + parserContext.getRegistry().registerBeanDefinition(name + "PartitionLock", partitionLockDefinition); } protected void parsePartitionPlan(Element element, diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/configuration/xml/StepFactoryBean.java b/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/configuration/xml/StepFactoryBean.java index 7f3b3a16b..d29e15798 100644 --- a/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/configuration/xml/StepFactoryBean.java +++ b/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/configuration/xml/StepFactoryBean.java @@ -31,7 +31,6 @@ import org.springframework.batch.core.jsr.step.builder.JsrBatchletStepBuilder; import org.springframework.batch.core.jsr.step.builder.JsrFaultTolerantStepBuilder; import org.springframework.batch.core.jsr.step.builder.JsrPartitionStepBuilder; import org.springframework.batch.core.jsr.step.builder.JsrSimpleStepBuilder; -import org.springframework.batch.core.partition.JsrStepExecutionSplitter; import org.springframework.batch.core.step.builder.FaultTolerantStepBuilder; import org.springframework.batch.core.step.builder.SimpleStepBuilder; import org.springframework.batch.core.step.builder.StepBuilder; @@ -179,8 +178,6 @@ public class StepFactoryBean extends StepParserStepFactoryBean { builder.reducer(reducer); } - builder.splitter(new JsrStepExecutionSplitter(getName(), getJobRepository())); - builder.aggregator(getStepExecutionAggergator()); return builder.build(); diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/configuration/xml/StepParser.java b/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/configuration/xml/StepParser.java index 64d36d473..dc38a6f03 100644 --- a/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/configuration/xml/StepParser.java +++ b/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/configuration/xml/StepParser.java @@ -68,9 +68,11 @@ public class StepParser extends AbstractSingleBeanDefinitionParser { } String allowStartIfComplete = element.getAttribute(ALLOW_START_IF_COMPLETE_ATTRIBUTE); + boolean allowStartIfCompletValue = false; if(StringUtils.hasText(allowStartIfComplete)) { bd.getPropertyValues().addPropertyValue("allowStartIfComplete", allowStartIfComplete); + allowStartIfCompletValue = Boolean.valueOf(allowStartIfComplete); } new ListenerParser(StepListenerFactoryBean.class, "listeners").parseListeners(element, parserContext, bd, stepName); @@ -91,7 +93,7 @@ public class StepParser extends AbstractSingleBeanDefinitionParser { } else if(name.equals(CHUNK_ELEMENT)) { new ChunkParser().parse(nestedElement, bd, parserContext, stepName); } else if(name.equals(PARTITION_ELEMENT)) { - new PartitionParser(stepName).parse(nestedElement, bd, parserContext, stepName); + new PartitionParser(stepName, allowStartIfCompletValue).parse(nestedElement, bd, parserContext, stepName); } } } diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/partition/JsrPartitionHandler.java b/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/partition/JsrPartitionHandler.java index bf23d72fc..00bc75c81 100644 --- a/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/partition/JsrPartitionHandler.java +++ b/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/partition/JsrPartitionHandler.java @@ -1,5 +1,5 @@ /* - * Copyright 2013 the original author or authors. + * Copyright 2013-2014 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. @@ -28,6 +28,7 @@ import java.util.concurrent.Callable; import java.util.concurrent.Future; import java.util.concurrent.FutureTask; import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.locks.ReentrantLock; import javax.batch.api.partition.PartitionAnalyzer; import javax.batch.api.partition.PartitionCollector; @@ -36,13 +37,17 @@ import javax.batch.api.partition.PartitionPlan; import org.springframework.batch.core.BatchStatus; import org.springframework.batch.core.ExitStatus; +import org.springframework.batch.core.JobExecutionException; import org.springframework.batch.core.Step; import org.springframework.batch.core.StepExecution; import org.springframework.batch.core.jsr.configuration.support.BatchArtifact.BatchArtifactType; import org.springframework.batch.core.jsr.configuration.support.BatchPropertyContext; import org.springframework.batch.core.jsr.configuration.support.BatchPropertyContext.BatchPropertyContextEntry; +import org.springframework.batch.core.partition.JsrStepExecutionSplitter; import org.springframework.batch.core.partition.PartitionHandler; import org.springframework.batch.core.partition.StepExecutionSplitter; +import org.springframework.batch.core.repository.JobRepository; +import org.springframework.batch.item.ExecutionContext; import org.springframework.beans.factory.InitializingBean; import org.springframework.core.task.TaskRejectedException; import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor; @@ -60,12 +65,19 @@ public class JsrPartitionHandler implements PartitionHandler, InitializingBean { // TODO: Replace with proper Channel and Messages once minimum support level for Spring is 4 private Queue partitionDataQueue; + private ReentrantLock lock; private Step step; private int partitions; private PartitionAnalyzer analyzer; private PartitionMapper mapper; private int threads; private BatchPropertyContext propertyContext; + private JobRepository jobRepository; + private boolean allowStartIfComplete = false; + + public void setAllowStartIfComplete(boolean allowStartIfComplete) { + this.allowStartIfComplete = allowStartIfComplete; + } /** * @param queue {@link Queue} to receive the output of the {@link PartitionCollector} @@ -74,6 +86,10 @@ public class JsrPartitionHandler implements PartitionHandler, InitializingBean { this.partitionDataQueue = queue; } + public void setPartitionLock(ReentrantLock lock) { + this.lock = lock; + } + /** * @param context {@link BatchPropertyContext} to resolve partition level step properties */ @@ -117,6 +133,13 @@ public class JsrPartitionHandler implements PartitionHandler, InitializingBean { this.partitions = partitions; } + /** + * @param jobRepository {@link JobRepository} + */ + public void setJobRepository(JobRepository jobRepository) { + this.jobRepository = jobRepository; + } + /* (non-Javadoc) * @see org.springframework.batch.core.partition.PartitionHandler#handle(org.springframework.batch.core.partition.StepExecutionSplitter, org.springframework.batch.core.StepExecution) */ @@ -127,24 +150,11 @@ public class JsrPartitionHandler implements PartitionHandler, InitializingBean { final Set result = new HashSet(); final ThreadPoolTaskExecutor taskExecutor = new ThreadPoolTaskExecutor(); - Set partitionStepExecutions; + int stepExecutionCount = jobRepository.getStepExecutionCount(stepExecution.getJobExecution().getJobInstance(), stepExecution.getStepName()); - if(mapper != null) { - PartitionPlan plan = mapper.mapPartitions(); - if(plan.getThreads() > 0) { - threads = plan.getThreads(); - } else if(plan.getPartitions() > 0) { - threads = plan.getPartitions(); - } else { - throw new IllegalArgumentException("Either a number of threads or partitions are required"); - } + boolean isRestart = stepExecutionCount > 1; - partitionStepExecutions = stepSplitter.split(stepExecution, plan.getPartitions()); - registerPartitionProperties(partitionStepExecutions, plan); - - } else { - partitionStepExecutions = stepSplitter.split(stepExecution, partitions); - } + Set partitionStepExecutions = splitStepExecution(stepExecution, isRestart); taskExecutor.setCorePoolSize(threads); taskExecutor.setMaxPoolSize(threads); @@ -171,8 +181,17 @@ public class JsrPartitionHandler implements PartitionHandler, InitializingBean { } } + processPartitionResults(tasks, result); + + return result; + } + + private void processPartitionResults( + final List> tasks, + final Set result) throws Exception { while(true) { - synchronized (partitionDataQueue) { + try { + lock.lock(); while(!partitionDataQueue.isEmpty()) { analyzer.analyzeCollectorData(partitionDataQueue.remove()); } @@ -182,10 +201,70 @@ public class JsrPartitionHandler implements PartitionHandler, InitializingBean { if(tasks.size() == 0) { break; } + } finally { + if(lock.isHeldByCurrentThread()) { + lock.unlock(); + } } } + } - return result; + private Set splitStepExecution(StepExecution stepExecution, + boolean isRestart) throws Exception, JobExecutionException { + Set partitionStepExecutions = new HashSet(); + if(isRestart) { + if(mapper != null) { + PartitionPlan plan = mapper.mapPartitions(); + + if(plan.getPartitionsOverride()) { + partitionStepExecutions = applyPartitionPlan(stepExecution, plan, false); + + for (StepExecution curStepExecution : partitionStepExecutions) { + curStepExecution.setExecutionContext(new ExecutionContext()); + } + } else { + Properties[] partitionProps = plan.getPartitionProperties(); + + plan = (PartitionPlanState) stepExecution.getExecutionContext().get("partitionPlanState"); + plan.setPartitionProperties(partitionProps); + + partitionStepExecutions = applyPartitionPlan(stepExecution, plan, true); + } + + } else { + StepExecutionSplitter stepSplitter = new JsrStepExecutionSplitter(jobRepository, allowStartIfComplete, stepExecution.getStepName(), true); + partitionStepExecutions = stepSplitter.split(stepExecution, partitions); + } + } else { + if(mapper != null) { + PartitionPlan plan = mapper.mapPartitions(); + partitionStepExecutions = applyPartitionPlan(stepExecution, plan, true); + } else { + StepExecutionSplitter stepSplitter = new JsrStepExecutionSplitter(jobRepository, allowStartIfComplete, stepExecution.getStepName(), true); + partitionStepExecutions = stepSplitter.split(stepExecution, partitions); + } + } + return partitionStepExecutions; + } + + private Set applyPartitionPlan(StepExecution stepExecution, + PartitionPlan plan, boolean restoreState) throws JobExecutionException { + StepExecutionSplitter stepSplitter; + Set partitionStepExecutions; + if(plan.getThreads() > 0) { + threads = plan.getThreads(); + } else if(plan.getPartitions() > 0) { + threads = plan.getPartitions(); + } else { + throw new IllegalArgumentException("Either a number of threads or partitions are required"); + } + + stepExecution.getExecutionContext().put("partitionPlanState", new PartitionPlanState(plan)); + + stepSplitter = new JsrStepExecutionSplitter(jobRepository, allowStartIfComplete, stepExecution.getStepName(), restoreState); + partitionStepExecutions = stepSplitter.split(stepExecution, plan.getPartitions()); + registerPartitionProperties(partitionStepExecutions, plan); + return partitionStepExecutions; } private void processFinishedPartitions( @@ -255,13 +334,110 @@ public class JsrPartitionHandler implements PartitionHandler, InitializingBean { }); } + /* (non-Javadoc) + * @see org.springframework.beans.factory.InitializingBean#afterPropertiesSet() + */ @Override public void afterPropertiesSet() throws Exception { Assert.notNull(propertyContext, "A BatchPropertyContext is required"); Assert.isTrue(mapper != null || threads > 0, "Either a mapper implementation or the number of partitions/threads is required"); + Assert.notNull(jobRepository, "A JobRepository is required"); if(partitionDataQueue == null) { partitionDataQueue = new LinkedBlockingQueue(); } + + if(lock == null) { + lock = new ReentrantLock(); + } + } + + /** + * Since a {@link PartitionPlan} could provide dynamic data (different results from run to run), + * the batch runtime needs to save off the results for restarts. This class serves as a container + * used to save off that state. + * + * @author Michael Minella + * @since 3.0 + */ + public static class PartitionPlanState implements PartitionPlan, Serializable { + + private static final long serialVersionUID = 1L; + private Properties[] partitionProperties; + private int partitions; + private int threads; + + /** + * @param plan the {@link PartitionPlan} that is the source of the state + */ + public PartitionPlanState(PartitionPlan plan) { + partitionProperties = plan.getPartitionProperties(); + partitions = plan.getPartitions(); + threads = plan.getThreads(); + } + + /* (non-Javadoc) + * @see javax.batch.api.partition.PartitionPlan#getPartitionProperties() + */ + @Override + public Properties[] getPartitionProperties() { + return partitionProperties; + } + + /* (non-Javadoc) + * @see javax.batch.api.partition.PartitionPlan#getPartitions() + */ + @Override + public int getPartitions() { + return partitions; + } + + /* (non-Javadoc) + * @see javax.batch.api.partition.PartitionPlan#getThreads() + */ + @Override + public int getThreads() { + return threads; + } + + /* (non-Javadoc) + * @see javax.batch.api.partition.PartitionPlan#setPartitions(int) + */ + @Override + public void setPartitions(int count) { + this.partitions = count; + } + + /* (non-Javadoc) + * @see javax.batch.api.partition.PartitionPlan#setPartitionsOverride(boolean) + */ + @Override + public void setPartitionsOverride(boolean override) { + // Intentional No-op + } + + /* (non-Javadoc) + * @see javax.batch.api.partition.PartitionPlan#getPartitionsOverride() + */ + @Override + public boolean getPartitionsOverride() { + return false; + } + + /* (non-Javadoc) + * @see javax.batch.api.partition.PartitionPlan#setThreads(int) + */ + @Override + public void setThreads(int count) { + this.threads = count; + } + + /* (non-Javadoc) + * @see javax.batch.api.partition.PartitionPlan#setPartitionProperties(java.util.Properties[]) + */ + @Override + public void setPartitionProperties(Properties[] props) { + this.partitionProperties = props; + } } } diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/partition/PartitionCollectorAdapter.java b/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/partition/PartitionCollectorAdapter.java index 1c6965fdc..f6c574884 100644 --- a/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/partition/PartitionCollectorAdapter.java +++ b/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/partition/PartitionCollectorAdapter.java @@ -1,5 +1,5 @@ /* - * Copyright 2013 the original author or authors. + * Copyright 2013-2014 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. @@ -17,6 +17,7 @@ package org.springframework.batch.core.jsr.partition; import java.io.Serializable; import java.util.Queue; +import java.util.concurrent.locks.ReentrantLock; import javax.batch.api.partition.PartitionCollector; import javax.batch.operations.BatchRuntimeException; @@ -38,6 +39,7 @@ public class PartitionCollectorAdapter implements ChunkListener { private PartitionCollector collector; private Queue partitionQueue; + private ReentrantLock lock; public PartitionCollectorAdapter(Queue queue, PartitionCollector collector) { Assert.notNull(queue, "A thread safe Queue is required"); @@ -47,6 +49,10 @@ public class PartitionCollectorAdapter implements ChunkListener { this.collector = collector; } + public void setPartitionLock(ReentrantLock lock) { + this.lock = lock; + } + @Override public void beforeChunk(ChunkContext context) { } @@ -54,22 +60,40 @@ public class PartitionCollectorAdapter implements ChunkListener { @Override public void afterChunk(ChunkContext context) { try { - synchronized (partitionQueue) { - partitionQueue.add(collector.collectPartitionData()); + if(context.isComplete()) { + lock.lock(); + Serializable collectPartitionData = collector.collectPartitionData(); + + if(collectPartitionData != null) { + partitionQueue.add(collectPartitionData); + } } } catch (Throwable e) { throw new BatchRuntimeException("An error occured while collecting data from the PartionCollector", e); + } finally { + if(lock.isHeldByCurrentThread()) { + lock.unlock(); + } } } @Override public void afterChunkError(ChunkContext context) { try { - synchronized (partitionQueue) { - partitionQueue.add(collector.collectPartitionData()); + lock.lock(); + if(context.isComplete()) { + Serializable collectPartitionData = collector.collectPartitionData(); + + if(collectPartitionData != null) { + partitionQueue.add(collectPartitionData); + } } } catch (Throwable e) { throw new BatchRuntimeException("An error occured while collecting data from the PartionCollector", e); + } finally { + if(lock.isHeldByCurrentThread()) { + lock.unlock(); + } } } } diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/step/PartitionStep.java b/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/step/PartitionStep.java index 8f0d0f547..ffde7aba7 100644 --- a/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/step/PartitionStep.java +++ b/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/step/PartitionStep.java @@ -70,7 +70,7 @@ public class PartitionStep extends org.springframework.batch.core.partition.supp } // Wait for task completion and then aggregate the results - Collection stepExecutions = getPartitionHandler().handle(getStepExecutionSplitter(), stepExecution); + Collection stepExecutions = getPartitionHandler().handle(null, stepExecution); stepExecution.upgradeStatus(BatchStatus.COMPLETED); stepExecutionAggregator.aggregate(stepExecution, stepExecutions); diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/step/batchlet/BatchletAdapter.java b/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/step/batchlet/BatchletAdapter.java index 993a8e70b..ddedb0613 100644 --- a/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/step/batchlet/BatchletAdapter.java +++ b/spring-batch-core/src/main/java/org/springframework/batch/core/jsr/step/batchlet/BatchletAdapter.java @@ -1,5 +1,5 @@ /* - * Copyright 2013 the original author or authors. + * Copyright 2013-2014 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. @@ -43,12 +43,18 @@ public class BatchletAdapter implements StoppableTasklet { @Override public RepeatStatus execute(StepContribution contribution, ChunkContext chunkContext) throws Exception { - String exitStatus = batchlet.process(); + String exitStatus; + try { + exitStatus = batchlet.process(); + } finally { + chunkContext.setComplete(); + } if(StringUtils.hasText(exitStatus)) { contribution.setExitStatus(new ExitStatus(exitStatus)); } + return RepeatStatus.FINISHED; } diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/partition/JsrStepExecutionSplitter.java b/spring-batch-core/src/main/java/org/springframework/batch/core/partition/JsrStepExecutionSplitter.java index 6b50a9323..ec20f4c33 100644 --- a/spring-batch-core/src/main/java/org/springframework/batch/core/partition/JsrStepExecutionSplitter.java +++ b/spring-batch-core/src/main/java/org/springframework/batch/core/partition/JsrStepExecutionSplitter.java @@ -1,5 +1,5 @@ /* - * Copyright 2013 the original author or authors. + * Copyright 2013-2014 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. @@ -23,7 +23,9 @@ import org.springframework.batch.core.JobExecution; import org.springframework.batch.core.JobExecutionException; import org.springframework.batch.core.StepExecution; import org.springframework.batch.core.jsr.launch.JsrJobOperator; +import org.springframework.batch.core.partition.support.SimpleStepExecutionSplitter; import org.springframework.batch.core.repository.JobRepository; +import org.springframework.batch.item.ExecutionContext; /** * Provides JSR-352 specific behavior for the splitting of {@link StepExecution}s. @@ -31,14 +33,17 @@ import org.springframework.batch.core.repository.JobRepository; * @author Michael Minella * @since 3.0 */ -public class JsrStepExecutionSplitter implements StepExecutionSplitter { +public class JsrStepExecutionSplitter extends SimpleStepExecutionSplitter { private String stepName; private JobRepository jobRepository; + private boolean restoreState; - public JsrStepExecutionSplitter(String stepName, JobRepository jobRepository) { + public JsrStepExecutionSplitter(JobRepository jobRepository, boolean allowStartIfComplete, String stepName, boolean restoreState) { + super(jobRepository, allowStartIfComplete, stepName, null); this.stepName = stepName; this.jobRepository = jobRepository; + this.restoreState = restoreState; } @Override @@ -78,7 +83,10 @@ public class JsrStepExecutionSplitter implements StepExecutionSplitter { String stepName = this.stepName + ":partition" + i; JobExecution curJobExecution = new JobExecution(jobExecution); StepExecution curStepExecution = new StepExecution(stepName, curJobExecution); - executions.add(curStepExecution); + + if(!restoreState || getStartable(curStepExecution, new ExecutionContext())) { + executions.add(curStepExecution); + } } jobRepository.addAll(executions); diff --git a/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/SimpleStepExecutionSplitter.java b/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/SimpleStepExecutionSplitter.java index f46f140fb..5a7993f1e 100644 --- a/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/SimpleStepExecutionSplitter.java +++ b/spring-batch-core/src/main/java/org/springframework/batch/core/partition/support/SimpleStepExecutionSplitter.java @@ -190,7 +190,7 @@ public class SimpleStepExecutionSplitter implements StepExecutionSplitter, Initi set.add(currentStepExecution); } } - + jobRepository.addAll(set); return set; @@ -235,7 +235,7 @@ public class SimpleStepExecutionSplitter implements StepExecutionSplitter, Initi return result; } - private boolean getStartable(StepExecution stepExecution, ExecutionContext context) throws JobExecutionException { + protected boolean getStartable(StepExecution stepExecution, ExecutionContext context) throws JobExecutionException { JobInstance jobInstance = stepExecution.getJobExecution().getJobInstance(); String stepName = stepExecution.getStepName(); diff --git a/spring-batch-core/src/test/java/org/springframework/batch/core/jsr/partition/JsrPartitionHandlerTests.java b/spring-batch-core/src/test/java/org/springframework/batch/core/jsr/partition/JsrPartitionHandlerTests.java index 32294dfb4..e6a6c32e9 100644 --- a/spring-batch-core/src/test/java/org/springframework/batch/core/jsr/partition/JsrPartitionHandlerTests.java +++ b/spring-batch-core/src/test/java/org/springframework/batch/core/jsr/partition/JsrPartitionHandlerTests.java @@ -1,5 +1,5 @@ /* - * Copyright 2013 the original author or authors. + * Copyright 2013-2014 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. @@ -24,19 +24,26 @@ import java.util.Properties; import java.util.Queue; import java.util.concurrent.ConcurrentLinkedQueue; +import javax.batch.api.BatchProperty; import javax.batch.api.partition.PartitionAnalyzer; +import javax.batch.api.partition.PartitionCollector; import javax.batch.api.partition.PartitionMapper; import javax.batch.api.partition.PartitionPlan; import javax.batch.api.partition.PartitionPlanImpl; +import javax.batch.api.partition.PartitionReducer; import javax.batch.runtime.BatchStatus; +import javax.inject.Inject; import org.junit.Before; import org.junit.Test; import org.springframework.batch.core.ExitStatus; import org.springframework.batch.core.JobExecution; +import org.springframework.batch.core.JobInstance; import org.springframework.batch.core.JobInterruptedException; import org.springframework.batch.core.StepExecution; +import org.springframework.batch.core.jsr.JsrTestUtils; import org.springframework.batch.core.jsr.configuration.support.BatchPropertyContext; +import org.springframework.batch.core.jsr.step.batchlet.BatchletSupport; import org.springframework.batch.core.partition.JsrStepExecutionSplitter; import org.springframework.batch.core.repository.JobRepository; import org.springframework.batch.core.repository.support.MapJobRepositoryFactoryBean; @@ -47,14 +54,17 @@ public class JsrPartitionHandlerTests { private JsrPartitionHandler handler; private JobRepository repository = new JobRepositorySupport(); - private StepExecution stepExecution = new StepExecution("step", new JobExecution(1L)); + private StepExecution stepExecution; private int count; private BatchPropertyContext propertyContext; private JsrStepExecutionSplitter stepSplitter; @Before public void setUp() throws Exception { - stepSplitter = new JsrStepExecutionSplitter("step1", repository); + JobExecution jobExecution = new JobExecution(1L); + jobExecution.setJobInstance(new JobInstance(1l, "job")); + stepExecution = new StepExecution("step1", jobExecution); + stepSplitter = new JsrStepExecutionSplitter(repository, false, "step1", true); Analyzer.collectorData = ""; Analyzer.status = ""; count = 0; @@ -70,6 +80,9 @@ public class JsrPartitionHandlerTests { propertyContext = new BatchPropertyContext(); handler.setPropertyContext(propertyContext); repository = new MapJobRepositoryFactoryBean().getJobRepository(); + handler.setJobRepository(repository); + MyPartitionReducer.reset(); + CountingPartitionCollector.reset(); } @Test @@ -93,6 +106,15 @@ public class JsrPartitionHandlerTests { } handler.setThreads(3); + + try { + handler.afterPropertiesSet(); + fail("JobRepository was not checked for"); + } catch(IllegalArgumentException iae) { + assertEquals("A JobRepository is required", iae.getMessage()); + } + + handler.setJobRepository(repository); handler.afterPropertiesSet(); } @@ -123,7 +145,7 @@ public class JsrPartitionHandlerTests { handler.afterPropertiesSet(); - Collection executions = handler.handle(new JsrStepExecutionSplitter("step1", repository), stepExecution); + Collection executions = handler.handle(new JsrStepExecutionSplitter(repository, false, "step1", true), stepExecution); assertEquals(3, executions.size()); assertEquals(3, count); @@ -144,7 +166,7 @@ public class JsrPartitionHandlerTests { handler.afterPropertiesSet(); - Collection executions = handler.handle(new JsrStepExecutionSplitter("step1", repository), stepExecution); + Collection executions = handler.handle(new JsrStepExecutionSplitter(repository, false, "step1", true), stepExecution); assertEquals(3, executions.size()); assertEquals(3, count); @@ -171,7 +193,7 @@ public class JsrPartitionHandlerTests { handler.afterPropertiesSet(); - Collection executions = handler.handle(new JsrStepExecutionSplitter("step1", repository), stepExecution); + Collection executions = handler.handle(new JsrStepExecutionSplitter(repository, false, "step1", true), stepExecution); assertEquals(3, executions.size()); assertEquals(3, count); @@ -191,7 +213,7 @@ public class JsrPartitionHandlerTests { handler.setPartitionAnalyzer(new Analyzer()); handler.afterPropertiesSet(); - Collection executions = handler.handle(new JsrStepExecutionSplitter("step1", repository), stepExecution); + Collection executions = handler.handle(new JsrStepExecutionSplitter(repository, false, "step1", true), stepExecution); assertEquals(2, executions.size()); assertEquals(2, count); @@ -199,6 +221,159 @@ public class JsrPartitionHandlerTests { assertEquals("COMPLETEDdone", Analyzer.status); } + @Test + public void testRestartNoOverride() throws Exception { + javax.batch.runtime.JobExecution execution1 = JsrTestUtils.runJob("jsrPartitionHandlerRestartWithOverrideJob", null, 1000000l); + assertEquals(BatchStatus.FAILED, execution1.getBatchStatus()); + assertEquals(1, MyPartitionReducer.beginCount); + assertEquals(0, MyPartitionReducer.beforeCount); + assertEquals(1, MyPartitionReducer.rollbackCount); + assertEquals(1, MyPartitionReducer.afterCount); + assertEquals(3, CountingPartitionCollector.collected); + + MyPartitionReducer.reset(); + CountingPartitionCollector.reset(); + + javax.batch.runtime.JobExecution execution2 = JsrTestUtils.restartJob(execution1.getExecutionId(), null, 1000000l); + assertEquals(BatchStatus.COMPLETED, execution2.getBatchStatus()); + assertEquals(1, MyPartitionReducer.beginCount); + assertEquals(1, MyPartitionReducer.beforeCount); + assertEquals(0, MyPartitionReducer.rollbackCount); + assertEquals(1, MyPartitionReducer.afterCount); + assertEquals(1, CountingPartitionCollector.collected); + } + + + @Test + public void testRestartOverride() throws Exception { + Properties jobParameters = new Properties(); + jobParameters.put("mapper.override", "true"); + + javax.batch.runtime.JobExecution execution1 = JsrTestUtils.runJob("jsrPartitionHandlerRestartWithOverrideJob", jobParameters, 1000000l); + assertEquals(BatchStatus.FAILED, execution1.getBatchStatus()); + assertEquals(1, MyPartitionReducer.beginCount); + assertEquals(0, MyPartitionReducer.beforeCount); + assertEquals(1, MyPartitionReducer.rollbackCount); + assertEquals(1, MyPartitionReducer.afterCount); + assertEquals(3, CountingPartitionCollector.collected); + + MyPartitionReducer.reset(); + CountingPartitionCollector.reset(); + + javax.batch.runtime.JobExecution execution2 = JsrTestUtils.restartJob(execution1.getExecutionId(), jobParameters, 1000000l); + assertEquals(BatchStatus.COMPLETED, execution2.getBatchStatus()); + assertEquals(1, MyPartitionReducer.beginCount); + assertEquals(1, MyPartitionReducer.beforeCount); + assertEquals(0, MyPartitionReducer.rollbackCount); + assertEquals(1, MyPartitionReducer.afterCount); + assertEquals(5, CountingPartitionCollector.collected); + } + + public static class CountingPartitionCollector implements PartitionCollector { + + public static int collected = 0; + + public static void reset() { + collected = 0; + } + + @Override + public Serializable collectPartitionData() throws Exception { + collected++; + + return null; + } + } + + public static class MyPartitionReducer implements PartitionReducer { + + public static int beginCount = 0; + public static int beforeCount = 0; + public static int rollbackCount = 0; + public static int afterCount = 0; + + public static void reset() { + beginCount = 0; + beforeCount = 0; + rollbackCount = 0; + afterCount = 0; + } + + @Override + public void beginPartitionedStep() throws Exception { + beginCount++; + } + + @Override + public void beforePartitionedStepCompletion() throws Exception { + beforeCount++; + } + + @Override + public void rollbackPartitionedStep() throws Exception { + rollbackCount++; + } + + @Override + public void afterPartitionedStepCompletion(PartitionStatus status) + throws Exception { + afterCount++; + } + } + + public static class MyPartitionMapper implements PartitionMapper { + + private static int count = 0; + + @Inject + @BatchProperty + String overrideString = "false"; + + @Override + public PartitionPlan mapPartitions() throws Exception { + count++; + + PartitionPlan plan = new PartitionPlanImpl(); + + if(count % 2 == 1) { + plan.setPartitions(3); + plan.setThreads(3); + } else { + plan.setPartitions(5); + plan.setThreads(5); + } + + plan.setPartitionsOverride(Boolean.valueOf(overrideString)); + + Properties[] props = new Properties[3]; + props[0] = new Properties(); + props[1] = new Properties(); + props[2] = new Properties(); + + if(count % 2 == 1) { + props[1].put("fail", "true"); + } + + plan.setPartitionProperties(props); + return plan; + } + } + + public static class MyBatchlet extends BatchletSupport { + @Inject + @BatchProperty + String fail; + + @Override + public String process() { + if("true".equalsIgnoreCase(fail)) { + throw new RuntimeException("Expected"); + } + + return null; + } + } + public static class Analyzer implements PartitionAnalyzer { public static String collectorData; diff --git a/spring-batch-core/src/test/java/org/springframework/batch/core/jsr/partition/JsrStepExecutionSplitterTests.java b/spring-batch-core/src/test/java/org/springframework/batch/core/jsr/partition/JsrStepExecutionSplitterTests.java index 8cde55318..d22da20b3 100644 --- a/spring-batch-core/src/test/java/org/springframework/batch/core/jsr/partition/JsrStepExecutionSplitterTests.java +++ b/spring-batch-core/src/test/java/org/springframework/batch/core/jsr/partition/JsrStepExecutionSplitterTests.java @@ -33,7 +33,7 @@ public class JsrStepExecutionSplitterTests { @Before public void setUp() throws Exception { - splitter = new JsrStepExecutionSplitter("step1", new JobRepositorySupport()); + splitter = new JsrStepExecutionSplitter(new JobRepositorySupport(), false, "step1", true); } @Test diff --git a/spring-batch-core/src/test/java/org/springframework/batch/core/jsr/partition/PartitionCollectorAdapterTests.java b/spring-batch-core/src/test/java/org/springframework/batch/core/jsr/partition/PartitionCollectorAdapterTests.java index 9a06d17bb..49c4a41b5 100644 --- a/spring-batch-core/src/test/java/org/springframework/batch/core/jsr/partition/PartitionCollectorAdapterTests.java +++ b/spring-batch-core/src/test/java/org/springframework/batch/core/jsr/partition/PartitionCollectorAdapterTests.java @@ -20,10 +20,12 @@ import static org.junit.Assert.assertEquals; import java.io.Serializable; import java.util.Queue; import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.locks.ReentrantLock; import javax.batch.api.partition.PartitionCollector; import org.junit.Test; +import org.springframework.batch.core.scope.context.ChunkContext; public class PartitionCollectorAdapterTests { @@ -43,9 +45,14 @@ public class PartitionCollectorAdapterTests { } }); - adapter.afterChunk(null); - adapter.afterChunkError(null); - adapter.afterChunk(null); + adapter.setPartitionLock(new ReentrantLock()); + + ChunkContext context = new ChunkContext(null); + context.setComplete(); + + adapter.afterChunk(context); + adapter.afterChunkError(context); + adapter.afterChunk(context); assertEquals(3, dataQueue.size()); assertEquals("0", dataQueue.remove()); diff --git a/spring-batch-core/src/test/java/org/springframework/batch/core/jsr/step/batchlet/BatchletAdapterTests.java b/spring-batch-core/src/test/java/org/springframework/batch/core/jsr/step/batchlet/BatchletAdapterTests.java index efa863379..fd7ff3ec0 100644 --- a/spring-batch-core/src/test/java/org/springframework/batch/core/jsr/step/batchlet/BatchletAdapterTests.java +++ b/spring-batch-core/src/test/java/org/springframework/batch/core/jsr/step/batchlet/BatchletAdapterTests.java @@ -14,6 +14,7 @@ import org.mockito.Mock; import org.mockito.MockitoAnnotations; import org.springframework.batch.core.ExitStatus; import org.springframework.batch.core.StepContribution; +import org.springframework.batch.core.scope.context.ChunkContext; import org.springframework.batch.repeat.RepeatStatus; public class BatchletAdapterTests { @@ -38,7 +39,7 @@ public class BatchletAdapterTests { @Test public void testExecuteNoExitStatus() throws Exception { - assertEquals(RepeatStatus.FINISHED, adapter.execute(contribution, null)); + assertEquals(RepeatStatus.FINISHED, adapter.execute(contribution, new ChunkContext(null))); verify(delegate).process(); } @@ -47,7 +48,7 @@ public class BatchletAdapterTests { public void testExecuteWithExitStatus() throws Exception { when(delegate.process()).thenReturn("my exit status"); - assertEquals(RepeatStatus.FINISHED, adapter.execute(contribution, null)); + assertEquals(RepeatStatus.FINISHED, adapter.execute(contribution, new ChunkContext(null))); verify(delegate).process(); verify(contribution).setExitStatus(new ExitStatus("my exit status")); diff --git a/spring-batch-core/src/test/resources/META-INF/batch-jobs/jsrPartitionHandlerRestartWithOverrideJob.xml b/spring-batch-core/src/test/resources/META-INF/batch-jobs/jsrPartitionHandlerRestartWithOverrideJob.xml new file mode 100644 index 000000000..f83c508d4 --- /dev/null +++ b/spring-batch-core/src/test/resources/META-INF/batch-jobs/jsrPartitionHandlerRestartWithOverrideJob.xml @@ -0,0 +1,15 @@ + + + + + + + + + + + + + + +