Revised scheduling lifecycle integration

ScheduledAnnotationBeanPostProcessor uses getBean(Class) for TaskScheduler/ScheduledExecutorService retrieval, allowing for a scheduler bean to be flagged as primary, and for a TaskScheduler bean to override a ScheduledExecutorService bean.

ContextLifecycleScheduledTaskRegistrar hooks into SmartInitializingSingleton's afterSingletonsInstantiated callback instead of ContextRefreshedEvent, as a natural consequence of SmartInitializingSingleton's introduction in Spring Framework 4.1 GA.

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

@@ -19,7 +19,6 @@ package org.springframework.scheduling.annotation;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.LinkedHashSet;
 import java.util.Map;
 import java.util.Set;
@@ -35,6 +34,8 @@ import org.springframework.beans.factory.BeanFactory;
 import org.springframework.beans.factory.BeanFactoryAware;
 import org.springframework.beans.factory.DisposableBean;
 import org.springframework.beans.factory.ListableBeanFactory;
+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.context.ApplicationContext;
@@ -87,7 +88,7 @@ public class ScheduledAnnotationBeanPostProcessor implements BeanPostProcessor,
 
 	private StringValueResolver embeddedValueResolver;
 
-	private ListableBeanFactory beanFactory;
+	private BeanFactory beanFactory;
 
 	private final ScheduledTaskRegistrar registrar = new ScheduledTaskRegistrar();
 
@@ -121,7 +122,7 @@ public class ScheduledAnnotationBeanPostProcessor implements BeanPostProcessor,
 	 */
 	@Override
 	public void setBeanFactory(BeanFactory beanFactory) {
-		this.beanFactory = (beanFactory instanceof ListableBeanFactory ? (ListableBeanFactory) beanFactory : null);
+		this.beanFactory = beanFactory;
 	}
 
 	/**
@@ -141,8 +142,9 @@ public class ScheduledAnnotationBeanPostProcessor implements BeanPostProcessor,
 			this.registrar.setScheduler(this.scheduler);
 		}
 
-		if (this.beanFactory != null) {
-			Map<String, SchedulingConfigurer> configurers = this.beanFactory.getBeansOfType(SchedulingConfigurer.class);
+		if (this.beanFactory instanceof ListableBeanFactory) {
+			Map<String, SchedulingConfigurer> configurers =
+					((ListableBeanFactory) this.beanFactory).getBeansOfType(SchedulingConfigurer.class);
 			for (SchedulingConfigurer configurer : configurers.values()) {
 				configurer.configureTasks(this.registrar);
 			}
@@ -150,21 +152,30 @@ public class ScheduledAnnotationBeanPostProcessor implements BeanPostProcessor,
 
 		if (this.registrar.hasTasks() && this.registrar.getScheduler() == null) {
 			Assert.state(this.beanFactory != null, "BeanFactory must be set to find scheduler by type");
-			Map<String, ? super Object> schedulers = new HashMap<String, Object>();
-			schedulers.putAll(this.beanFactory.getBeansOfType(TaskScheduler.class));
-			schedulers.putAll(this.beanFactory.getBeansOfType(ScheduledExecutorService.class));
-			if (schedulers.size() == 0) {
-				// do nothing -> fall back to default scheduler
+			try {
+				// Search for TaskScheduler bean...
+				this.registrar.setScheduler(this.beanFactory.getBean(TaskScheduler.class));
 			}
-			else if (schedulers.size() == 1) {
-				this.registrar.setScheduler(schedulers.values().iterator().next());
+			catch (NoUniqueBeanDefinitionException ex) {
+				throw new IllegalStateException("More than one TaskScheduler exists within the context. " +
+						"Remove all but one of the beans; or implement the SchedulingConfigurer interface and call " +
+						"ScheduledTaskRegistrar#setScheduler explicitly within the configureTasks() callback.", ex);
 			}
-			else if (schedulers.size() >= 2){
-				throw new IllegalStateException("More than one TaskScheduler and/or ScheduledExecutorService  " +
-						"exist within the context. Remove all but one of the beans; or implement the " +
-						"SchedulingConfigurer interface and call ScheduledTaskRegistrar#setScheduler " +
-						"explicitly within the configureTasks() callback. Found the following beans: " +
-						schedulers.keySet());
+			catch (NoSuchBeanDefinitionException ex) {
+				logger.debug("Could not find default TaskScheduler bean", ex);
+				// Search for ScheduledExecutorService bean next...
+				try {
+					this.registrar.setScheduler(this.beanFactory.getBean(ScheduledExecutorService.class));
+				}
+				catch (NoUniqueBeanDefinitionException ex2) {
+					throw new IllegalStateException("More than one ScheduledExecutorService exists within the context. " +
+							"Remove all but one of the beans; or implement the SchedulingConfigurer interface and call " +
+							"ScheduledTaskRegistrar#setScheduler explicitly within the configureTasks() callback.", ex);
+				}
+				catch (NoSuchBeanDefinitionException ex2) {
+					logger.debug("Could not find default ScheduledExecutorService bean", ex);
+					// Giving up -> falling back to default scheduler within the registrar...
+				}
 			}
 		}
 

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2013 the original author or authors.
+ * Copyright 2002-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.
@@ -16,51 +16,24 @@
 
 package org.springframework.scheduling.config;
 
-import org.springframework.context.ApplicationContext;
-import org.springframework.context.ApplicationContextAware;
-import org.springframework.context.ApplicationListener;
-import org.springframework.context.event.ContextRefreshedEvent;
+import org.springframework.beans.factory.SmartInitializingSingleton;
 
 /**
- * {@link ScheduledTaskRegistrar} subclass that redirects the actual scheduling
- * of tasks to the {@link ContextRefreshedEvent} callback. Falls back to regular
- * {@code ScheduledTaskRegistrar} behavior when not running within an ApplicationContext.
+ * {@link ScheduledTaskRegistrar} subclass which redirects the actual scheduling
+ * of tasks to the {@link #afterSingletonsInstantiated()} callback (as of 4.1.2).
  *
  * @author Juergen Hoeller
  * @since 3.2.1
  */
-public class ContextLifecycleScheduledTaskRegistrar extends ScheduledTaskRegistrar
-		implements ApplicationContextAware, ApplicationListener<ContextRefreshedEvent> {
+public class ContextLifecycleScheduledTaskRegistrar extends ScheduledTaskRegistrar implements SmartInitializingSingleton {
 
-	private ApplicationContext applicationContext;
-
-
-	@Override
-	public void setApplicationContext(ApplicationContext applicationContext) {
-		this.applicationContext = applicationContext;
-	}
-
-
-	/**
-	 * If we're running within an ApplicationContext, don't schedule the tasks
-	 * right here; wait for this context's ContextRefreshedEvent instead.
-	 */
 	@Override
 	public void afterPropertiesSet() {
-		if (this.applicationContext == null) {
-			scheduleTasks();
-		}
+		// no-op
 	}
 
-	/**
-	 * Actually schedule the tasks at the right time of the context lifecycle,
-	 * if we're running within an ApplicationContext.
-	 */
 	@Override
-	public void onApplicationEvent(ContextRefreshedEvent event) {
-		if (event.getApplicationContext() != this.applicationContext) {
-			return;
-		}
+	public void afterSingletonsInstantiated() {
 		scheduleTasks();
 	}