Merge pull request #173 from willschipp/BATCH-2009

* BATCH-2009:
  BATCH-2009 - Stoppable tasklet * Added the StoppableTasklet interface (an extension of the Tasklet interface). * Updated JobOperator#stop(long executionId) to call StoppableTasklet#stop() on any currently running tasklets (local executions only). * Updated the SystemCommandTasklet to implement StoppableTasklet

Author: mminella

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

@@ -33,6 +33,7 @@
 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;
@@ -52,6 +53,11 @@
 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.NoSuchStepException;
+import org.springframework.batch.core.step.StepLocator;
+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 {
@@ -394,6 +401,33 @@ public boolean stop(long executionId) throws NoSuchJobExecutionException, JobExe
 		jobExecution.setStatus(BatchStatus.STOPPING);
 		jobRepository.update(jobExecution);
 
+		try {
+			Job job = jobRegistry.getJob(jobExecution.getJobInstance().getJobName());
+			if (job instanceof StepLocator) {//can only process as StepLocator is the only way to get the step object
+				//get the current stepExecution
+				for (StepExecution stepExecution : jobExecution.getStepExecutions()) {
+					if (stepExecution.getStatus().isRunning()) {
+						try {
+							//have the step execution that's running -> need to 'stop' it
+							Step step = ((StepLocator)job).getStep(stepExecution.getStepName());
+							if (step instanceof TaskletStep) {
+								Tasklet tasklet = ((TaskletStep)step).getTasklet();
+								if (tasklet instanceof StoppableTasklet) {
+									((StoppableTasklet)tasklet).stop();
+								}
+							}
+						}
+						catch (NoSuchStepException e) {
+							logger.warn("Step not found",e);
+						}
+					}
+				}
+			}
+		}
+		catch (NoSuchJobException e) {
+			logger.warn("Cannot find Job object",e);
+		}
+
 		return true;
 	}
 

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

@@ -0,0 +1,38 @@
+/*
+ * Copyright 2013 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.core.step.tasklet;
+
+import org.springframework.batch.core.launch.JobOperator;
+
+/**
+ * An extension to the {@link Tasklet} interface to allow users to
+ * add logic for stopping a tasklet.  It is up to each implementation
+ * as to how the stop will behave.  The only guarantee provided by the
+ * framework is that a call to {@link JobOperator#stop(long)} will
+ * attempt to call the stop method on any currently running
+ * StoppableTasklet.
+ *
+ * @author Will Schipp
+ * @since 3.0
+ */
+public interface StoppableTasklet extends Tasklet {
+
+	/**
+	 * Used to signal that the job this {@link Tasklet} is executing
+	 * within has been requested 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,8 +115,12 @@ else if (execution.isTerminateOnly()) {
 				systemCommandTask.cancel(interruptOnCancel);
 				throw new JobInterruptedException("Job interrupted while executing system command '" + command + "'");
 			}
+			else if (stopped) {
+				systemCommandTask.cancel(interruptOnCancel);
+				contribution.setExitStatus(ExitStatus.STOPPED);
+				return RepeatStatus.FINISHED;
+			}
 		}
-
 	}
 
 	/**
@@ -208,4 +215,18 @@ public void setInterruptOnCancel(boolean interruptOnCancel) {
 		this.interruptOnCancel = interruptOnCancel;
 	}
 
+	/**
+	 * Will interrupt the thread executing the system command only if
+	 * {@link #setInterruptOnCancel(boolean)} has been set to true.  Otherwise
+	 * the underlying command will be allowed to finish before the tasklet
+	 * ends.
+	 *
+	 * @since 3.0
+	 * @see StoppableTasklet#stop()
+	 */
+	@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 the {@link Tasklet} instance being executed within this step
+	 */
+	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

@@ -19,10 +19,12 @@
 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;
@@ -35,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;
@@ -51,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;
 
 /**
@@ -254,6 +265,7 @@ public void testFindRunningExecutionsSunnyDay() throws Exception {
 	}
 
 	@Test
+	@SuppressWarnings("unchecked")
 	public void testFindRunningExecutionsNoSuchJob() throws Exception {
 		jobParameters = new JobParameters();
 		when(jobExplorer.findRunningJobExecutions("no-such-job")).thenReturn(Collections.EMPTY_SET);
@@ -348,6 +360,67 @@ 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, null);
+		StoppableTasklet tasklet = mock(StoppableTasklet.class);
+		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());
+	}
+
+	@Test
+	public void testStopTaskletException() throws Exception {
+		JobInstance jobInstance = new JobInstance(123L, job.getName());
+		JobExecution jobExecution = new JobExecution(jobInstance, 111L, jobParameters, null);
+		StoppableTasklet tasklet = new StoppableTasklet() {
+
+			@Override
+			public RepeatStatus execute(StepContribution contribution,
+					ChunkContext chunkContext) throws Exception {
+				return null;
+			}
+
+			@Override
+			public void stop() {
+				throw new IllegalStateException();
+			}};
+			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());
+	}
+
 	@Test
 	public void testAbort() throws Exception {
 		JobInstance jobInstance = new JobInstance(123L, job.getName());
@@ -369,4 +442,25 @@ 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 {
+
+		}
+
+	}
 }

spring-batch-core/src/test/java/org/springframework/batch/core/step/tasklet/SystemCommandTaskletIntegrationTests.java

@@ -221,6 +221,26 @@ public void testWorkingDirectory() throws Exception {
 		tasklet.setWorkingDirectory(directory.getCanonicalPath());
 	}
 
+	/*
+	 * test stopping a tasklet
+	 */
+	@Test
+	public void testStopped() throws Exception {
+		String command = System.getProperty("os.name").toLowerCase().indexOf("win") >= 0 ?
+				"ping 1.1.1.1 -n 1 -w 5000" :
+					"sleep 5";
+		tasklet.setCommand(command);
+		tasklet.setTerminationCheckInterval(10);
+		tasklet.afterPropertiesSet();
+
+		StepContribution contribution = stepExecution.createStepContribution();
+		//send stop
+		tasklet.stop();
+		tasklet.execute(contribution, null);
+
+		assertEquals(contribution.getExitStatus().getExitCode(),ExitStatus.STOPPED.getExitCode());
+	}
+
 	/**
 	 * Exit code mapper containing mapping logic expected by the tests. 0 means
 	 * finished successfully, other value means failure.