All Downloads are FREE. Search and download functionalities are using the official Maven repository.

org.springframework.batch.core.jsr.partition.JsrPartitionHandler Maven / Gradle / Ivy

/*
 * Copyright 2013-2018 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
 *
 *      https://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.core.jsr.partition;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Properties;
import java.util.Queue;
import java.util.Set;
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;
import javax.batch.api.partition.PartitionMapper;
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.BatchPropertyContext;
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;
import org.springframework.util.Assert;

/**
 * Executes a step instance per thread using a {@link ThreadPoolTaskExecutor} in
 * accordance with JSR-352.  The results from each step is aggregated into a
 * cumulative result.
 *
 * @author Michael Minella
 * @author Mahmoud Ben Hassine
 * @since 3.0
 */
public class JsrPartitionHandler implements PartitionHandler, InitializingBean {

	private static final int DEFAULT_POLLING_INTERVAL = 500;

	// 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;
	private Set partitionStepNames = new HashSet<>();
	private int pollingInterval = DEFAULT_POLLING_INTERVAL;

	/**
	 * @return the step that will be executed by each partition
	 */
	public Step getStep() {
		return step;
	}

	/**
	 * @return the names of each partitioned step
	 */
	public Collection getPartitionStepNames() {
		return partitionStepNames;
	}

	/**
	 * @param allowStartIfComplete flag stating if the step should restart if it
	 * 	was complete in a previous run
	 */
	public void setAllowStartIfComplete(boolean allowStartIfComplete) {
		this.allowStartIfComplete = allowStartIfComplete;
	}

	/**
	 * @param queue {@link Queue} to receive the output of the {@link PartitionCollector}
	 */
	public void setPartitionDataQueue(Queue queue) {
		this.partitionDataQueue = queue;
	}

	public void setPartitionLock(ReentrantLock lock) {
		this.lock = lock;
	}

	/**
	 * @param context {@link BatchPropertyContext} to resolve partition level step properties
	 */
	public void setPropertyContext(BatchPropertyContext context) {
		this.propertyContext = context;
	}

	/**
	 * @param mapper {@link PartitionMapper} used to configure partitioning
	 */
	public void setPartitionMapper(PartitionMapper mapper) {
		this.mapper = mapper;
	}

	/**
	 * @param step the step to be executed as a partitioned step
	 */
	public void setStep(Step step) {
		this.step = step;
	}

	/**
	 * @param analyzer {@link PartitionAnalyzer}
	 */
	public void setPartitionAnalyzer(PartitionAnalyzer analyzer) {
		this.analyzer = analyzer;
	}

	/**
	 * @param threads the number of threads to execute the partitions to be run
	 * within.  The default is the number of partitions.
	 */
	public void setThreads(int threads) {
		this.threads = threads;
	}

	/**
	 * @param partitions the number of partitions to be executed
	 */
	public void setPartitions(int partitions) {
		this.partitions = partitions;
	}

	/**
	 * @param jobRepository {@link JobRepository}
	 */
	public void setJobRepository(JobRepository jobRepository) {
		this.jobRepository = jobRepository;
	}

	/**
	 * @param pollingInterval the duration of partitions completion polling interval
	 *                       (in milliseconds). The default value is 500ms.
	 */
	public void setPollingInterval(int pollingInterval) {
		this.pollingInterval = pollingInterval;
	}

	/* (non-Javadoc)
	 * @see org.springframework.batch.core.partition.PartitionHandler#handle(org.springframework.batch.core.partition.StepExecutionSplitter, org.springframework.batch.core.StepExecution)
	 */
	@Override
	public Collection handle(StepExecutionSplitter stepSplitter,
			StepExecution stepExecution) throws Exception {
		final List> tasks = new ArrayList<>();
		final Set result = new HashSet<>();
		final ThreadPoolTaskExecutor taskExecutor = new ThreadPoolTaskExecutor();

		int stepExecutionCount = jobRepository.getStepExecutionCount(stepExecution.getJobExecution().getJobInstance(), stepExecution.getStepName());

		boolean isRestart = stepExecutionCount > 1;

		Set partitionStepExecutions = splitStepExecution(stepExecution, isRestart);

		for (StepExecution curStepExecution : partitionStepExecutions) {
			partitionStepNames.add(curStepExecution.getStepName());
		}

		taskExecutor.setCorePoolSize(threads);
		taskExecutor.setMaxPoolSize(threads);

		taskExecutor.initialize();

		try {
			for (final StepExecution curStepExecution : partitionStepExecutions) {
				final FutureTask task = createTask(step, curStepExecution);

				try {
					taskExecutor.execute(task);
					tasks.add(task);
				} catch (TaskRejectedException e) {
					// couldn't execute one of the tasks
					ExitStatus exitStatus = ExitStatus.FAILED
							.addExitDescription("TaskExecutor rejected the task for this step.");
				/*
				 * Set the status in case the caller is tracking it through the
				 * JobExecution.
				 */
					curStepExecution.setStatus(BatchStatus.FAILED);
					curStepExecution.setExitStatus(exitStatus);
					result.add(stepExecution);
				}
			}

			processPartitionResults(tasks, result);
		}
		finally {
			taskExecutor.shutdown();
		}

		return result;
	}

	/**
	 * Blocks until all partitioned steps have completed.  As each step completes
	 * the PartitionAnalyzer analyzes the collector data received from each
	 * partition (if there is any).
	 *
	 * @param tasks The {@link Future} that contains the reference to the executing step
	 * @param result Set of completed {@link StepExecution}s
	 * @throws Exception
	 */
	private void processPartitionResults(
			final List> tasks,
			final Set result) throws Exception {
		while(true) {
			Thread.sleep(pollingInterval);
			try {
				lock.lock();
				while(!partitionDataQueue.isEmpty()) {
					analyzer.analyzeCollectorData(partitionDataQueue.remove());
				}

				processFinishedPartitions(tasks, result);

				if(tasks.size() == 0) {
					break;
				}
			} finally {
				if(lock.isHeldByCurrentThread()) {
					lock.unlock();
				}
			}
		}
	}

	/**
	 * Uses either the {@link PartitionMapper} or the hard coded configuration to split
	 * the supplied manager StepExecution into the worker StepExecutions.
	 *
	 * @param stepExecution manager {@link StepExecution}
	 * @param isRestart true if this step is being restarted
	 * @return a {@link Set} of {@link StepExecution}s to be executed
	 * @throws Exception
	 * @throws JobExecutionException
	 */
	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");
		}

		PartitionPlanState partitionPlanState = new PartitionPlanState();
		partitionPlanState.setPartitionPlan(plan);

		stepExecution.getExecutionContext().put("partitionPlanState", partitionPlanState);

		stepSplitter = new JsrStepExecutionSplitter(jobRepository, allowStartIfComplete, stepExecution.getStepName(), restoreState);
		partitionStepExecutions = stepSplitter.split(stepExecution, plan.getPartitions());
		registerPartitionProperties(partitionStepExecutions, plan);
		return partitionStepExecutions;
	}

	private void processFinishedPartitions(
			final List> tasks,
			final Set result) throws Exception {
		for(int i = 0; i < tasks.size(); i++) {
			Future curTask = tasks.get(i);

			if(curTask.isDone()) {
				StepExecution curStepExecution = curTask.get();

				if(analyzer != null) {
					analyzer.analyzeStatus(curStepExecution.getStatus().getBatchStatus(), curStepExecution.getExitStatus().getExitCode());
				}

				result.add(curStepExecution);

				tasks.remove(i);
				i--;
			}
		}
	}

	private void registerPartitionProperties(
			Set partitionStepExecutions, PartitionPlan plan) {
		Properties[] partitionProperties = plan.getPartitionProperties();
		if(partitionProperties != null) {
			Iterator executions = partitionStepExecutions.iterator();

			int i = 0;
			while(executions.hasNext()) {
				StepExecution curExecution = executions.next();

				if(i < partitionProperties.length) {
					Properties partitionPropertyValues = partitionProperties[i];
					if(partitionPropertyValues != null) {
						propertyContext.setStepProperties(curExecution.getStepName(), partitionPropertyValues);
					}

					i++;
				} else {
					break;
				}
			}
		}
	}

	/**
	 * Creates the task executing the given step in the context of the given execution.
	 *
	 * @param step the step to execute
	 * @param stepExecution the given execution
	 * @return the task executing the given step
	 */
	protected FutureTask createTask(final Step step,
			final StepExecution stepExecution) {
		return new FutureTask<>(new Callable() {
			@Override
			public StepExecution call() throws Exception {
				step.execute(stepExecution);
				return stepExecution;
			}
		});
	}

	/* (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 || partitions > 0), "Either a mapper implementation or the number of partitions/threads is required");
		Assert.notNull(jobRepository, "A JobRepository is required");
		Assert.isTrue(pollingInterval >= 0, "The polling interval must be positive");

		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();
		}

		public PartitionPlanState() {
		}

		public void setPartitionPlan(PartitionPlan plan) {
			this.partitionProperties = plan.getPartitionProperties();
			this.partitions = plan.getPartitions();
			this.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;
		}
	}
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy