ScheduledAnnotationBeanPostProcessor registers tasks in ContextRefreshedEvent phase (again)

Issue: SPR-12641

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2014 the original author or authors.
+ * Copyright 2002-2015 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.
@@ -39,7 +39,10 @@ import org.springframework.beans.factory.NoUniqueBeanDefinitionException;
 import org.springframework.beans.factory.SmartInitializingSingleton;
 import org.springframework.beans.factory.config.BeanPostProcessor;
 import org.springframework.context.ApplicationContext;
+import org.springframework.context.ApplicationContextAware;
+import org.springframework.context.ApplicationListener;
 import org.springframework.context.EmbeddedValueResolverAware;
+import org.springframework.context.event.ContextRefreshedEvent;
 import org.springframework.core.Ordered;
 import org.springframework.core.annotation.AnnotationUtils;
 import org.springframework.scheduling.TaskScheduler;
@@ -80,7 +83,8 @@ import org.springframework.util.StringValueResolver;
  * @see org.springframework.scheduling.config.ScheduledTaskRegistrar
  */
 public class ScheduledAnnotationBeanPostProcessor implements BeanPostProcessor, Ordered,
-		EmbeddedValueResolverAware, BeanFactoryAware, SmartInitializingSingleton, DisposableBean {
+		EmbeddedValueResolverAware, BeanFactoryAware, ApplicationContextAware,
+		SmartInitializingSingleton, ApplicationListener<ContextRefreshedEvent>, DisposableBean {
 
 	protected final Log logger = LogFactory.getLog(getClass());
 
@@ -90,6 +94,8 @@ public class ScheduledAnnotationBeanPostProcessor implements BeanPostProcessor,
 
 	private BeanFactory beanFactory;
 
+	private ApplicationContext applicationContext;
+
 	private final ScheduledTaskRegistrar registrar = new ScheduledTaskRegistrar();
 
 	private final Set<Class<?>> nonAnnotatedClasses =
@@ -126,10 +132,13 @@ public class ScheduledAnnotationBeanPostProcessor implements BeanPostProcessor,
 	}
 
 	/**
-	 * @deprecated as of Spring 4.1, in favor of {@link #setBeanFactory}
+	 * Setting an {@link ApplicationContext} is optional: If set, registered
+	 * tasks will be activated in the {@link ContextRefreshedEvent} phase;
+	 * if not set, it will happen at {@link #afterSingletonsInstantiated} time.
 	 */
-	@Deprecated
+	@Override
 	public void setApplicationContext(ApplicationContext applicationContext) {
+		this.applicationContext = applicationContext;
 		if (this.beanFactory == null) {
 			this.beanFactory = applicationContext;
 		}
@@ -138,6 +147,23 @@ public class ScheduledAnnotationBeanPostProcessor implements BeanPostProcessor,
 
 	@Override
 	public void afterSingletonsInstantiated() {
+		if (this.applicationContext == null) {
+			// Not running in an ApplicationContext -> register tasks early...
+			finishRegistration();
+		}
+	}
+
+	@Override
+	public void onApplicationEvent(ContextRefreshedEvent event) {
+		if (event.getApplicationContext() == this.applicationContext) {
+			// Running in an ApplicationContext -> register tasks this late...
+			// giving other ContextRefreshedEvent listeners a chance to perform
+			// their work at the same time (e.g. Spring Batch's job registration).
+			finishRegistration();
+		}
+	}
+
+	private void finishRegistration() {
 		if (this.scheduler != null) {
 			this.registrar.setScheduler(this.scheduler);
 		}