Restores refreshable support using BeanDefinitionRegistryPostProcessor

spring-cloud-context/src/main/java/org/springframework/cloud/autoconfigure/RefreshAutoConfiguration.java

@@ -23,17 +23,17 @@ import java.util.Set;
 
 import javax.annotation.PostConstruct;
 
-import org.springframework.aop.framework.ProxyFactory;
+import org.springframework.aop.scope.ScopedProxyUtils;
 import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.BeanFactory;
 import org.springframework.beans.factory.ListableBeanFactory;
-import org.springframework.beans.factory.NoSuchBeanDefinitionException;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.config.BeanDefinition;
-import org.springframework.beans.factory.config.BeanPostProcessor;
+import org.springframework.beans.factory.config.BeanDefinitionHolder;
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
 import org.springframework.beans.factory.support.BeanDefinitionRegistry;
 import org.springframework.beans.factory.support.BeanDefinitionRegistryPostProcessor;
+import org.springframework.beans.factory.support.DefaultListableBeanFactory;
 import org.springframework.boot.autoconfigure.AutoConfigureBefore;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
@@ -100,9 +100,11 @@ public class RefreshAutoConfiguration {
 
 	@Component
 	protected static class RefreshScopeBeanDefinitionEnhancer
-			implements BeanPostProcessor, BeanDefinitionRegistryPostProcessor {
+			implements BeanDefinitionRegistryPostProcessor {
 
-		private BeanDefinitionRegistry registry;
+		private Environment environment;
+
+		private boolean bound = false;
 
 		/**
 		 * Class names for beans to post process into refresh scope. Useful when you don't
@@ -129,40 +131,42 @@ public class RefreshAutoConfiguration {
 		@Override
 		public void postProcessBeanFactory(ConfigurableListableBeanFactory beanFactory)
 				throws BeansException {
-			Environment environment = beanFactory.getBean(Environment.class);
-			if (environment == null) {
-				environment = new StandardEnvironment();
-			}
-			Binder.get(environment).bind(REFRESH_SCOPE_PREFIX, Bindable.ofInstance(this));
 		}
 
 		@Override
 		public void postProcessBeanDefinitionRegistry(BeanDefinitionRegistry registry)
 				throws BeansException {
-			this.registry = registry;
-		}
+			// temporarily allow bean definition overriding
+			boolean allowBeanOverriding = false;
+			if (registry instanceof DefaultListableBeanFactory) {
+				DefaultListableBeanFactory beanFactory = (DefaultListableBeanFactory) registry;
+				allowBeanOverriding = beanFactory.isAllowBeanDefinitionOverriding();
+				beanFactory.setAllowBeanDefinitionOverriding(true);
+			}
 
-		@Override
-		public Object postProcessAfterInitialization(Object bean, String beanName)
-				throws BeansException {
-			BeanDefinition definition = null;
-			try {
-				definition = registry.getBeanDefinition(beanName);
+			bindEnvironmentIfNeeded(registry);
+
+			for (String name : registry.getBeanDefinitionNames()) {
+				BeanDefinition definition = registry.getBeanDefinition(name);
+				if (isApplicable(registry, name, definition)) {
+					BeanDefinitionHolder holder = new BeanDefinitionHolder(definition,
+							name);
+					BeanDefinitionHolder proxy = ScopedProxyUtils
+							.createScopedProxy(holder, registry, true);
+					definition.setScope("refresh");
+					registry.registerBeanDefinition(proxy.getBeanName(),
+							proxy.getBeanDefinition());
+				}
 			}
-			catch (NoSuchBeanDefinitionException e) {
-				// just ignore and move on
-				return bean;
+			// restore bean definition overriding
+			if (registry instanceof DefaultListableBeanFactory) {
+				DefaultListableBeanFactory beanFactory = (DefaultListableBeanFactory) registry;
+				beanFactory.setAllowBeanDefinitionOverriding(allowBeanOverriding);
 			}
-			if (isApplicable(registry, beanName, definition)) {
-				definition.setScope(REFRESH_SCOPE_NAME);
-				ProxyFactory proxyFactory = new ProxyFactory(bean);
-				return proxyFactory.getProxy();
-			}
-			return bean;
 		}
 
 		private boolean isApplicable(BeanDefinitionRegistry registry, String name,
-				BeanDefinition definition) {
+									 BeanDefinition definition) {
 			String scope = definition.getScope();
 			if (REFRESH_SCOPE_NAME.equals(scope)) {
 				// Already refresh scoped
@@ -181,6 +185,21 @@ public class RefreshAutoConfiguration {
 			return false;
 		}
 
+		private void bindEnvironmentIfNeeded(BeanDefinitionRegistry registry) {
+			if (!bound) { // only bind once
+				if (this.environment == null && registry instanceof BeanFactory) {
+					this.environment = ((BeanFactory) registry)
+							.getBean(Environment.class);
+				}
+				if (this.environment == null) {
+					this.environment = new StandardEnvironment();
+				}
+				Binder.get(environment).bind("spring.cloud.refresh",
+						Bindable.ofInstance(this));
+				bound = true;
+			}
+		}
+
 	}
 
 	@Bean