Introduce callback for singleton availability

Closes gh-21362

Author: jhoeller

Diff for: spring-beans/src/main/java/org/springframework/beans/factory/config/SingletonBeanRegistry.java

@@ -16,6 +16,8 @@
 
 package org.springframework.beans.factory.config;
 
+import java.util.function.Consumer;
+
 import org.springframework.lang.Nullable;
 
 /**
@@ -57,6 +59,17 @@ public interface SingletonBeanRegistry {
 	 */
 	void registerSingleton(String beanName, Object singletonObject);
 
+	/**
+	 * Add a callback to be triggered when the specified singleton becomes available
+	 * in the bean registry.
+	 * @param beanName the name of the bean
+	 * @param singletonConsumer a callback for reacting to the availability of the freshly
+	 * registered/created singleton instance (intended for follow-up steps before the bean is
+	 * actively used by other callers, not for modifying the given singleton instance itself)
+	 * @since 6.2
+	 */
+	void addSingletonCallback(String beanName, Consumer<Object> singletonConsumer);
+
 	/**
 	 * Return the (raw) singleton object registered under the given name.
 	 * <p>Only checks already instantiated singletons; does not return an Object

Diff for: spring-beans/src/main/java/org/springframework/beans/factory/support/DefaultSingletonBeanRegistry.java

@@ -17,7 +17,6 @@
 package org.springframework.beans.factory.support;
 
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
@@ -27,6 +26,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.Consumer;
 
 import org.springframework.beans.factory.BeanCreationException;
 import org.springframework.beans.factory.BeanCreationNotAllowedException;
@@ -79,8 +79,11 @@ public class DefaultSingletonBeanRegistry extends SimpleAliasRegistry implements
 	/** Cache of singleton objects: bean name to bean instance. */
 	private final Map<String, Object> singletonObjects = new ConcurrentHashMap<>(256);
 
-	/** Cache of singleton factories: bean name to ObjectFactory. */
-	private final Map<String, ObjectFactory<?>> singletonFactories = new HashMap<>(16);
+	/** Creation-time registry of singleton factories: bean name to ObjectFactory. */
+	private final Map<String, ObjectFactory<?>> singletonFactories = new ConcurrentHashMap<>(16);
+
+	/** Custom callbacks for singleton creation/registration. */
+	private final Map<String, Consumer<Object>> singletonCallbacks = new ConcurrentHashMap<>(16);
 
 	/** Cache of early singleton objects: bean name to bean instance. */
 	private final Map<String, Object> earlySingletonObjects = new ConcurrentHashMap<>(16);
@@ -133,8 +136,8 @@ public void registerSingleton(String beanName, Object singletonObject) throws Il
 	}
 
 	/**
-	 * Add the given singleton object to the singleton cache of this factory.
-	 * <p>To be called for eager registration of singletons.
+	 * Add the given singleton object to the singleton registry.
+	 * <p>To be called for exposure of freshly registered/created singletons.
 	 * @param beanName the name of the bean
 	 * @param singletonObject the singleton object
 	 */
@@ -147,12 +150,17 @@ protected void addSingleton(String beanName, Object singletonObject) {
 		this.singletonFactories.remove(beanName);
 		this.earlySingletonObjects.remove(beanName);
 		this.registeredSingletons.add(beanName);
+
+		Consumer<Object> callback = this.singletonCallbacks.get(beanName);
+		if (callback != null) {
+			callback.accept(singletonObject);
+		}
 	}
 
 	/**
 	 * Add the given singleton factory for building the specified singleton
 	 * if necessary.
-	 * <p>To be called for eager registration of singletons, e.g. to be able to
+	 * <p>To be called for early exposure purposes, e.g. to be able to
 	 * resolve circular references.
 	 * @param beanName the name of the bean
 	 * @param singletonFactory the factory for the singleton object
@@ -164,6 +172,11 @@ protected void addSingletonFactory(String beanName, ObjectFactory<?> singletonFa
 		this.registeredSingletons.add(beanName);
 	}
 
+	@Override
+	public void addSingletonCallback(String beanName, Consumer<Object> singletonConsumer) {
+		this.singletonCallbacks.put(beanName, singletonConsumer);
+	}
+
 	@Override
 	@Nullable
 	public Object getSingleton(String beanName) {
@@ -262,6 +275,7 @@ public Object getSingleton(String beanName, ObjectFactory<?> singletonFactory) {
 						}
 					}
 				}
+
 				if (this.singletonsCurrentlyInDestruction) {
 					throw new BeanCreationNotAllowedException(beanName,
 							"Singleton bean creation not allowed while singletons of this factory are in destruction " +
@@ -343,8 +357,8 @@ protected void onSuppressedException(Exception ex) {
 	}
 
 	/**
-	 * Remove the bean with the given name from the singleton cache of this factory,
-	 * to be able to clean up eager registration of a singleton if creation failed.
+	 * Remove the bean with the given name from the singleton registry, either on
+	 * regular destruction or on cleanup after early exposure when creation failed.
 	 * @param beanName the name of the bean
 	 */
 	protected void removeSingleton(String beanName) {

Diff for: spring-beans/src/test/java/org/springframework/beans/factory/support/DefaultSingletonBeanRegistryTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2023 the original author or authors.
+ * Copyright 2002-2024 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,6 +16,8 @@
 
 package org.springframework.beans.factory.support;
 
+import java.util.concurrent.atomic.AtomicBoolean;
+
 import org.junit.jupiter.api.Test;
 
 import org.springframework.beans.testfixture.beans.DerivedTestBean;
@@ -35,12 +37,18 @@ class DefaultSingletonBeanRegistryTests {
 
 	@Test
 	void singletons() {
+		AtomicBoolean tbFlag = new AtomicBoolean();
+		beanRegistry.addSingletonCallback("tb", instance -> tbFlag.set(true));
 		TestBean tb = new TestBean();
 		beanRegistry.registerSingleton("tb", tb);
 		assertThat(beanRegistry.getSingleton("tb")).isSameAs(tb);
+		assertThat(tbFlag.get()).isTrue();
 
+		AtomicBoolean tb2Flag = new AtomicBoolean();
+		beanRegistry.addSingletonCallback("tb2", instance -> tb2Flag.set(true));
 		TestBean tb2 = (TestBean) beanRegistry.getSingleton("tb2", TestBean::new);
 		assertThat(beanRegistry.getSingleton("tb2")).isSameAs(tb2);
+		assertThat(tb2Flag.get()).isTrue();
 
 		assertThat(beanRegistry.getSingleton("tb")).isSameAs(tb);
 		assertThat(beanRegistry.getSingleton("tb2")).isSameAs(tb2);