ScheduledAnnotationBeanPostProcessor tracks individual bean instances of any scope

Issue: SPR-12216
Issue: SPR-12872

Author: jhoeller

spring-context/src/main/java/org/springframework/scheduling/annotation/ScheduledAnnotationBeanPostProcessor.java

@@ -17,7 +17,9 @@
 package org.springframework.scheduling.annotation;
 
 import java.lang.reflect.Method;
+import java.util.Collection;
 import java.util.Collections;
+import java.util.LinkedHashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.TimeZone;
@@ -35,7 +37,7 @@
 import org.springframework.beans.factory.NoSuchBeanDefinitionException;
 import org.springframework.beans.factory.NoUniqueBeanDefinitionException;
 import org.springframework.beans.factory.SmartInitializingSingleton;
-import org.springframework.beans.factory.config.BeanPostProcessor;
+import org.springframework.beans.factory.config.DestructionAwareBeanPostProcessor;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.ApplicationContextAware;
 import org.springframework.context.ApplicationListener;
@@ -48,6 +50,7 @@
 import org.springframework.scheduling.Trigger;
 import org.springframework.scheduling.config.CronTask;
 import org.springframework.scheduling.config.IntervalTask;
+import org.springframework.scheduling.config.ScheduledTask;
 import org.springframework.scheduling.config.ScheduledTaskRegistrar;
 import org.springframework.scheduling.support.CronTrigger;
 import org.springframework.scheduling.support.ScheduledMethodRunnable;
@@ -81,8 +84,8 @@
  * @see org.springframework.scheduling.config.ScheduledTaskRegistrar
  * @see AsyncAnnotationBeanPostProcessor
  */
-public class ScheduledAnnotationBeanPostProcessor implements BeanPostProcessor, Ordered,
-		EmbeddedValueResolverAware, BeanFactoryAware, ApplicationContextAware,
+public class ScheduledAnnotationBeanPostProcessor implements DestructionAwareBeanPostProcessor,
+		Ordered, EmbeddedValueResolverAware, BeanFactoryAware, ApplicationContextAware,
 		SmartInitializingSingleton, ApplicationListener<ContextRefreshedEvent>, DisposableBean {
 
 	/**
@@ -109,6 +112,9 @@ public class ScheduledAnnotationBeanPostProcessor implements BeanPostProcessor,
 	private final Set<Class<?>> nonAnnotatedClasses =
 			Collections.newSetFromMap(new ConcurrentHashMap<Class<?>, Boolean>(64));
 
+	private final Map<Object, Set<ScheduledTask>> scheduledTasks =
+			new ConcurrentHashMap<Object, Set<ScheduledTask>>(16);
+
 
 	@Override
 	public int getOrder() {
@@ -296,6 +302,9 @@ protected void processScheduled(Scheduled scheduled, Method method, Object bean)
 			String errorMessage =
 					"Exactly one of the 'cron', 'fixedDelay(String)', or 'fixedRate(String)' attributes is required";
 
+			Set<ScheduledTask> tasks =
+					new LinkedHashSet<ScheduledTask>(4);
+
 			// Determine initial delay
 			long initialDelay = scheduled.initialDelay();
 			String initialDelayString = scheduled.initialDelayString();
@@ -330,7 +339,7 @@ protected void processScheduled(Scheduled scheduled, Method method, Object bean)
 				else {
 					timeZone = TimeZone.getDefault();
 				}
-				this.registrar.addCronTask(new CronTask(runnable, new CronTrigger(cron, timeZone)));
+				tasks.add(this.registrar.scheduleCronTask(new CronTask(runnable, new CronTrigger(cron, timeZone))));
 			}
 
 			// At this point we don't need to differentiate between initial delay set or not anymore
@@ -343,7 +352,7 @@ protected void processScheduled(Scheduled scheduled, Method method, Object bean)
 			if (fixedDelay >= 0) {
 				Assert.isTrue(!processedSchedule, errorMessage);
 				processedSchedule = true;
-				this.registrar.addFixedDelayTask(new IntervalTask(runnable, fixedDelay, initialDelay));
+				tasks.add(this.registrar.scheduleFixedDelayTask(new IntervalTask(runnable, fixedDelay, initialDelay)));
 			}
 			String fixedDelayString = scheduled.fixedDelayString();
 			if (StringUtils.hasText(fixedDelayString)) {
@@ -359,15 +368,15 @@ protected void processScheduled(Scheduled scheduled, Method method, Object bean)
 					throw new IllegalArgumentException(
 							"Invalid fixedDelayString value \"" + fixedDelayString + "\" - cannot parse into integer");
 				}
-				this.registrar.addFixedDelayTask(new IntervalTask(runnable, fixedDelay, initialDelay));
+				tasks.add(this.registrar.scheduleFixedDelayTask(new IntervalTask(runnable, fixedDelay, initialDelay)));
 			}
 
 			// Check fixed rate
 			long fixedRate = scheduled.fixedRate();
 			if (fixedRate >= 0) {
 				Assert.isTrue(!processedSchedule, errorMessage);
 				processedSchedule = true;
-				this.registrar.addFixedRateTask(new IntervalTask(runnable, fixedRate, initialDelay));
+				tasks.add(this.registrar.scheduleFixedRateTask(new IntervalTask(runnable, fixedRate, initialDelay)));
 			}
 			String fixedRateString = scheduled.fixedRateString();
 			if (StringUtils.hasText(fixedRateString)) {
@@ -383,11 +392,12 @@ protected void processScheduled(Scheduled scheduled, Method method, Object bean)
 					throw new IllegalArgumentException(
 							"Invalid fixedRateString value \"" + fixedRateString + "\" - cannot parse into integer");
 				}
-				this.registrar.addFixedRateTask(new IntervalTask(runnable, fixedRate, initialDelay));
+				tasks.add(this.registrar.scheduleFixedRateTask(new IntervalTask(runnable, fixedRate, initialDelay)));
 			}
 
 			// Check whether we had any attribute set
 			Assert.isTrue(processedSchedule, errorMessage);
+			this.scheduledTasks.put(bean, tasks);
 		}
 		catch (IllegalArgumentException ex) {
 			throw new IllegalStateException(
@@ -396,8 +406,30 @@ protected void processScheduled(Scheduled scheduled, Method method, Object bean)
 	}
 
 
+	@Override
+	public void postProcessBeforeDestruction(Object bean, String beanName) {
+		Set<ScheduledTask> tasks = this.scheduledTasks.remove(bean);
+		if (tasks != null) {
+			for (ScheduledTask task : tasks) {
+				task.cancel();
+			}
+		}
+	}
+
+	@Override
+	public boolean requiresDestruction(Object bean) {
+		return this.scheduledTasks.containsKey(bean);
+	}
+
 	@Override
 	public void destroy() {
+		Collection<Set<ScheduledTask>> allTasks = this.scheduledTasks.values();
+		for (Set<ScheduledTask> tasks : allTasks) {
+			for (ScheduledTask task : tasks) {
+				task.cancel();
+			}
+		}
+		this.scheduledTasks.clear();
 		this.registrar.destroy();
 	}
 

spring-context/src/main/java/org/springframework/scheduling/config/ScheduledTask.java

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2002-2016 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.scheduling.config;
+
+import java.util.concurrent.ScheduledFuture;
+
+/**
+ * A representation of a scheduled task,
+ * used as a return value for scheduling methods.
+ *
+ * @author Juergen Hoeller
+ * @since 4.3
+ * @see ScheduledTaskRegistrar#scheduleTriggerTask
+ * @see ScheduledTaskRegistrar#scheduleFixedRateTask
+ */
+public final class ScheduledTask {
+
+	volatile ScheduledFuture<?> future;
+
+
+	ScheduledTask() {
+	}
+
+
+	/**
+	 * Trigger cancellation of this scheduled task.
+	 */
+	public void cancel() {
+		ScheduledFuture<?> future = this.future;
+		if (future != null) {
+			future.cancel(true);
+		}
+	}
+
+}