BATCH-2009 - stoppable tasklet

spring-batch-core/src/main/java/org/springframework/batch/core/launch/support/SimpleJobOperator.java

@@ -16,6 +16,7 @@
 package org.springframework.batch.core.launch.support;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Date;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
@@ -33,13 +34,15 @@
 import org.springframework.batch.core.JobParameters;
 import org.springframework.batch.core.JobParametersIncrementer;
 import org.springframework.batch.core.JobParametersInvalidException;
+import org.springframework.batch.core.Step;
 import org.springframework.batch.core.StepExecution;
 import org.springframework.batch.core.UnexpectedJobExecutionException;
 import org.springframework.batch.core.configuration.JobRegistry;
 import org.springframework.batch.core.configuration.ListableJobLocator;
 import org.springframework.batch.core.converter.DefaultJobParametersConverter;
 import org.springframework.batch.core.converter.JobParametersConverter;
 import org.springframework.batch.core.explore.JobExplorer;
+import org.springframework.batch.core.job.AbstractJob;
 import org.springframework.batch.core.launch.JobExecutionNotRunningException;
 import org.springframework.batch.core.launch.JobInstanceAlreadyExistsException;
 import org.springframework.batch.core.launch.JobLauncher;
@@ -52,6 +55,9 @@
 import org.springframework.batch.core.repository.JobInstanceAlreadyCompleteException;
 import org.springframework.batch.core.repository.JobRepository;
 import org.springframework.batch.core.repository.JobRestartException;
+import org.springframework.batch.core.step.tasklet.StoppableTasklet;
+import org.springframework.batch.core.step.tasklet.Tasklet;
+import org.springframework.batch.core.step.tasklet.TaskletStep;
 import org.springframework.batch.support.PropertiesConverter;
 import org.springframework.beans.factory.InitializingBean;
 import org.springframework.transaction.annotation.Transactional;
@@ -71,6 +77,7 @@
  *
  * @author Dave Syer
  * @author Lucas Ward
+ * @author Will Schipp
  * @since 2.0
  */
 public class SimpleJobOperator implements JobOperator, InitializingBean {
@@ -392,6 +399,32 @@ public boolean stop(long executionId) throws NoSuchJobExecutionException, JobExe
 		}
 		jobExecution.setStatus(BatchStatus.STOPPING);
 		jobRepository.update(jobExecution);
+		//implementation to support Tasklet.stop()
+		//TODO manage this is an 'scoped' proxy test
+		//find the job object
+		Job job;
+		try {
+			job = jobRegistry.getJob(jobExecution.getJobInstance().getJobName());
+			//get the steps for the job
+			if (job instanceof AbstractJob) {
+				//retrieve the steps
+				Collection<String> stepNames = ((AbstractJob)job).getStepNames();
+				//go through the step names are retrieve each step
+				for (String stepName : stepNames) {
+					Step step = ((AbstractJob)job).getStep(stepName);
+					//determine type
+					if (step instanceof TaskletStep) {
+						//invoke stop --> reflection?
+						Tasklet tasklet = ((TaskletStep) step).getTasklet();
+						if (tasklet instanceof StoppableTasklet) {
+							((StoppableTasklet)tasklet).stop();
+						}//end if
+					}//end if
+				}//end for
+			}//end if
+		} catch (NoSuchJobException e) {
+			logger.error("Couldn't find Job for the execution:" + jobExecution.getJobInstance().getJobName());
+		}
 
 		return true;
 	}

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

@@ -18,8 +18,10 @@
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
+import org.springframework.batch.core.ExitStatus;
 import org.springframework.batch.core.StepContribution;
 import org.springframework.batch.core.scope.context.ChunkContext;
+import org.springframework.batch.core.step.tasklet.StoppableTasklet;
 import org.springframework.batch.core.step.tasklet.Tasklet;
 import org.springframework.batch.repeat.RepeatStatus;
 
@@ -28,10 +30,11 @@
  * handling.
  *
  * @author Dave Syer
+ * @author Will Schipp
  *
  * @param <I> input item type
  */
-public class ChunkOrientedTasklet<I> implements Tasklet {
+public class ChunkOrientedTasklet<I> implements StoppableTasklet {
 
 	private static final String INPUTS_KEY = "INPUTS";
 
@@ -40,6 +43,11 @@ public class ChunkOrientedTasklet<I> implements Tasklet {
 	private final ChunkProvider<I> chunkProvider;
 
 	private boolean buffering = true;
+
+	/**
+	 * support for stoppable tasklet
+	 */
+	private boolean stopped = false;
 
 	private static Log logger = LogFactory.getLog(ChunkOrientedTasklet.class);
 
@@ -63,6 +71,13 @@ public void setBuffering(boolean buffering) {
 	@Override
 	public RepeatStatus execute(StepContribution contribution, ChunkContext chunkContext) throws Exception {
 
+		//check for stopped at the beginning of a chunk
+		if (stopped) {
+			stopped = false;//reset
+			contribution.setExitStatus(ExitStatus.STOPPED);
+			return RepeatStatus.FINISHED;
+		}//end if
+
 		@SuppressWarnings("unchecked")
 		Chunk<I> inputs = (Chunk<I>) chunkContext.getAttribute(INPUTS_KEY);
 		if (inputs == null) {
@@ -90,4 +105,9 @@ public RepeatStatus execute(StepContribution contribution, ChunkContext chunkCon
 
 	}
 
+	@Override
+	public void stop() {
+		stopped = true;
+	}
+
 }

spring-batch-core/src/main/java/org/springframework/batch/core/step/tasklet/CallableTaskletAdapter.java

@@ -30,7 +30,7 @@
  * @author Dave Syer
  *
  */
-public class CallableTaskletAdapter implements Tasklet, InitializingBean {
+public class CallableTaskletAdapter implements StoppableTasklet, InitializingBean {
 
 	private Callable<RepeatStatus> callable;
 
@@ -62,4 +62,7 @@ public RepeatStatus execute(StepContribution contribution, ChunkContext chunkCon
 		return callable.call();
 	}
 
+	@Override
+	public void stop() { }
+
 }

spring-batch-core/src/main/java/org/springframework/batch/core/step/tasklet/MethodInvokingTaskletAdapter.java

@@ -31,9 +31,10 @@
  * @see AbstractMethodInvokingDelegator
  *
  * @author Dave Syer
+ * @author Will Schipp
  *
  */
-public class MethodInvokingTaskletAdapter extends AbstractMethodInvokingDelegator<Object> implements Tasklet {
+public class MethodInvokingTaskletAdapter extends AbstractMethodInvokingDelegator<Object> implements StoppableTasklet {
 
 	/**
 	 * Delegate execution to the target object and translate the return value to
@@ -62,4 +63,7 @@ protected ExitStatus mapResult(Object result) {
 		return ExitStatus.COMPLETED;
 	}
 
+	@Override
+	public void stop() { }
+
 }

spring-batch-core/src/main/java/org/springframework/batch/core/step/tasklet/StoppableTasklet.java

@@ -0,0 +1,20 @@
+package org.springframework.batch.core.step.tasklet;
+
+/**
+ * 
+ * JSR-352 compatible tasklet that provides the 'stop' function.
+ * The Spring Batch Tasklet is analogous to 'batchlet' in the JSR
+ * terminology
+ * 
+ * @author Will Schipp
+ *
+ */
+public interface StoppableTasklet extends Tasklet {
+
+	/**
+	 * method to signal a long running/looping tasklet to stop
+	 * 
+	 */
+	void stop();
+
+}

spring-batch-core/src/main/java/org/springframework/batch/core/step/tasklet/SystemCommandTasklet.java

@@ -54,8 +54,9 @@
  * still running when tasklet exits (abnormally).
  *
  * @author Robert Kasanicky
+ * @author Will Schipp
  */
-public class SystemCommandTasklet extends StepExecutionListenerSupport implements Tasklet, InitializingBean {
+public class SystemCommandTasklet extends StepExecutionListenerSupport implements StoppableTasklet, InitializingBean {
 
 	protected static final Log logger = LogFactory.getLog(SystemCommandTasklet.class);
 
@@ -77,6 +78,8 @@ public class SystemCommandTasklet extends StepExecutionListenerSupport implement
 
 	private boolean interruptOnCancel = false;
 
+	private boolean stopped = false;
+
 	/**
 	 * Execute system command and map its exit code to {@link ExitStatus} using
 	 * {@link SystemProcessExitCodeMapper}.
@@ -99,7 +102,7 @@ public Integer call() throws Exception {
 		taskExecutor.execute(systemCommandTask);
 
 		while (true) {
-			Thread.sleep(checkInterval);
+			Thread.sleep(checkInterval);//moved to the end of the logic			
 			if (systemCommandTask.isDone()) {
 				contribution.setExitStatus(systemProcessExitCodeMapper.getExitStatus(systemCommandTask.get()));
 				return RepeatStatus.FINISHED;
@@ -112,6 +115,13 @@ else if (execution.isTerminateOnly()) {
 				systemCommandTask.cancel(interruptOnCancel);
 				throw new JobInterruptedException("Job interrupted while executing system command '" + command + "'");
 			}
+			else if (stopped) {
+				stopped = false;//reset
+				systemCommandTask.cancel(interruptOnCancel);
+				contribution.setExitStatus(ExitStatus.STOPPED);
+				return RepeatStatus.FINISHED;
+			}
+
 		}
 
 	}
@@ -208,4 +218,9 @@ public void setInterruptOnCancel(boolean interruptOnCancel) {
 		this.interruptOnCancel = interruptOnCancel;
 	}
 
+	@Override
+	public void stop() {
+		stopped = true;
+	}
+
 }

spring-batch-core/src/main/java/org/springframework/batch/core/step/tasklet/TaskletStep.java

@@ -71,6 +71,7 @@
  * @author Ben Hale
  * @author Robert Kasanicky
  * @author Michael Minella
+ * @author Will Schipp
  */
 @SuppressWarnings("serial")
 public class TaskletStep extends AbstractStep {
@@ -306,6 +307,14 @@ protected void open(ExecutionContext ctx) throws Exception {
 		stream.open(ctx);
 	}
 
+	/**
+	 * retrieve the tasklet - helper method for JobOperator
+	 * @return
+	 */
+	public Tasklet getTasklet() {
+		return tasklet;
+	}
+
 	/**
 	 * A callback for the transactional work inside a chunk. Also detects
 	 * failures in the transaction commit and rollback, only panicking if the

spring-batch-core/src/test/java/org/springframework/batch/core/launch/support/SimpleJobOperatorTests.java

@@ -15,15 +15,16 @@
  */
 package org.springframework.batch.core.launch.support;
 
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
-
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.mockito.Matchers.anyString;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
@@ -36,12 +37,17 @@
 import org.springframework.batch.core.BatchStatus;
 import org.springframework.batch.core.Job;
 import org.springframework.batch.core.JobExecution;
+import org.springframework.batch.core.JobExecutionException;
 import org.springframework.batch.core.JobInstance;
 import org.springframework.batch.core.JobParameters;
 import org.springframework.batch.core.JobParametersIncrementer;
+import org.springframework.batch.core.Step;
+import org.springframework.batch.core.StepContribution;
+import org.springframework.batch.core.configuration.JobRegistry;
 import org.springframework.batch.core.configuration.support.MapJobRegistry;
 import org.springframework.batch.core.converter.DefaultJobParametersConverter;
 import org.springframework.batch.core.explore.JobExplorer;
+import org.springframework.batch.core.job.AbstractJob;
 import org.springframework.batch.core.job.JobSupport;
 import org.springframework.batch.core.launch.JobInstanceAlreadyExistsException;
 import org.springframework.batch.core.launch.JobLauncher;
@@ -52,6 +58,10 @@
 import org.springframework.batch.core.repository.JobInstanceAlreadyCompleteException;
 import org.springframework.batch.core.repository.JobRepository;
 import org.springframework.batch.core.repository.JobRestartException;
+import org.springframework.batch.core.scope.context.ChunkContext;
+import org.springframework.batch.core.step.tasklet.StoppableTasklet;
+import org.springframework.batch.core.step.tasklet.TaskletStep;
+import org.springframework.batch.repeat.RepeatStatus;
 import org.springframework.batch.support.PropertiesConverter;
 
 /**
@@ -349,6 +359,32 @@ public void testStop() throws Exception{
 		assertEquals(BatchStatus.STOPPING, jobExecution.getStatus());
 	}
 
+	@Test
+	public void testStopTasklet() throws Exception {
+		JobInstance jobInstance = new JobInstance(123L, job.getName());
+		JobExecution jobExecution = new JobExecution(jobInstance, 111L, jobParameters);
+		MockStoppableTasklet tasklet = new MockStoppableTasklet();
+		TaskletStep taskletStep = new TaskletStep();
+		taskletStep.setTasklet(tasklet);
+		MockJob job = new MockJob();
+		job.taskletStep = taskletStep;
+
+		JobRegistry jobRegistry = mock(JobRegistry.class);
+		TaskletStep step = mock(TaskletStep.class);
+
+		when(step.getTasklet()).thenReturn(tasklet);
+		when(step.getName()).thenReturn("test_job.step1");
+		when(jobRegistry.getJob(anyString())).thenReturn(job);
+		when(jobExplorer.getJobExecution(111L)).thenReturn(jobExecution);
+
+		jobOperator.setJobRegistry(jobRegistry);
+		jobExplorer.getJobExecution(111L);
+		jobRepository.update(jobExecution);
+		jobOperator.stop(111L);
+		assertEquals(BatchStatus.STOPPING, jobExecution.getStatus());	
+		assertTrue(tasklet.stopped);
+	}
+
 	@Test
 	public void testAbort() throws Exception {
 		JobInstance jobInstance = new JobInstance(123L, job.getName());
@@ -370,4 +406,42 @@ public void testAbortNonStopping() throws Exception {
 		jobRepository.update(jobExecution);
 		jobOperator.abandon(123L);
 	}
+
+	class MockJob extends AbstractJob {
+
+		private TaskletStep taskletStep;
+
+		@Override
+		public Step getStep(String stepName) {
+			return taskletStep;
+		}
+
+		@Override
+		public Collection<String> getStepNames() {
+			return Arrays.asList("test_job.step1");
+		}
+
+		@Override
+		protected void doExecute(JobExecution execution) throws JobExecutionException {
+
+		}
+
+	}
+
+	class MockStoppableTasklet implements StoppableTasklet {
+
+		boolean stopped = Boolean.FALSE;
+
+		@Override
+		public RepeatStatus execute(StepContribution contribution,
+				ChunkContext chunkContext) throws Exception {
+			return null;
+		}
+
+		@Override
+		public void stop() {
+			stopped = Boolean.TRUE;
+		}
+
+	}
 }