Get Specification beans from LoadBalancerClientFactory. Sanitise generated class names. Fix generated statements.

spring-cloud-context/src/main/java/org/springframework/cloud/context/named/NamedContextFactory.java

@@ -215,6 +215,10 @@ public abstract class NamedContextFactory<C extends NamedContextFactory.Specific
 		return BeanFactoryUtils.beansOfTypeIncludingAncestors(context, type);
 	}
 
+	public Map<String, C> getConfigurations() {
+		return configurations;
+	}
+
 	/**
 	 * Specification with name and configuration.
 	 */

spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/aot/ChildContextMappings.java

@@ -1,66 +0,0 @@
-/*
- * Copyright 2012-2020 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
- *
- *      https://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.cloud.loadbalancer.aot;
-
-import java.util.ArrayList;
-import java.util.List;
-
-import org.springframework.beans.factory.config.BeanDefinition;
-import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
-import org.springframework.beans.factory.config.ConstructorArgumentValues;
-import org.springframework.cloud.loadbalancer.annotation.LoadBalancerClientSpecification;
-import org.springframework.context.ConfigurableApplicationContext;
-
-/**
- * @author Olga Maciaszek-Sharma
- */
-final class ChildContextMappings {
-
-	private ChildContextMappings() {
-		throw new IllegalStateException("Can't instantiate a utility class");
-	}
-
-	static List<Entry> getChildContexts(ConfigurableApplicationContext applicationContext) {
-		ConfigurableListableBeanFactory beanFactory = applicationContext.getBeanFactory();
-		String[] beanNames = applicationContext.getBeanNamesForType(LoadBalancerClientSpecification.class);
-		List<Entry> contexts = new ArrayList<>();
-		for (String beanName : beanNames) {
-			contexts.add(process(beanFactory.getMergedBeanDefinition(beanName)));
-		}
-		return contexts;
-	}
-
-	private static Entry process(BeanDefinition beanDefinition) {
-		ConstructorArgumentValues constructorArguments = beanDefinition.getConstructorArgumentValues();
-		ConstructorArgumentValues.ValueHolder nameValueHolder = constructorArguments.getIndexedArgumentValue(0, null);
-		if (nameValueHolder != null) {
-			String name = (String) nameValueHolder.getValue();
-			ConstructorArgumentValues.ValueHolder configurationsValueHolder = constructorArguments
-					.getIndexedArgumentValue(1, null);
-			if (configurationsValueHolder != null) {
-				String[] configurations = (String[]) configurationsValueHolder.getValue();
-				return new Entry(name, configurations);
-			}
-		}
-		throw new IllegalArgumentException("Invalid bean definition " + beanDefinition);
-	}
-
-	record Entry(String name, String[] configurations) {
-
-	}
-
-}

spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/aot/LoadBalancerChildContextInitializer.java

@@ -17,7 +17,6 @@
 package org.springframework.cloud.loadbalancer.aot;
 
 import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
@@ -88,14 +87,15 @@ public class LoadBalancerChildContextInitializer
 				&& registeredBean.getBeanFactory().equals(applicationBeanFactory))) {
 			return null;
 		}
-		List<ChildContextMappings.Entry> contextsMap = ChildContextMappings.getChildContexts(context);
+		// TODO: add context ids from properties
-		Set<ConfigurableApplicationContext> childContextAotContributions = contextsMap.stream()
+		Set<String> contextIds = loadBalancerClientFactory.getConfigurations().keySet();
+		Set<ConfigurableApplicationContext> childContextAotContributions = contextIds.stream()
 				.map(this::buildChildContext).collect(Collectors.toSet());
 		return new AotContribution(childContextAotContributions);
 	}
 
-	private ConfigurableApplicationContext buildChildContext(ChildContextMappings.Entry entry) {
+	private ConfigurableApplicationContext buildChildContext(String contextId) {
-		ConfigurableApplicationContext childContext = loadBalancerClientFactory.buildContext(entry.name());
+		ConfigurableApplicationContext childContext = loadBalancerClientFactory.buildContext(contextId);
 		registerBeans(childContext);
 		return childContext;
 	}
@@ -130,7 +130,9 @@ public class LoadBalancerChildContextInitializer
 		public void applyTo(GenerationContext generationContext, BeanRegistrationCode beanRegistrationCode) {
 			ApplicationContextAotGenerator contextAotGenerator = new ApplicationContextAotGenerator();
 			Map<String, ClassName> generatedInitializerClassNames = childContexts.stream().map(childContext -> {
-				GenerationContext childGenerationContext = generationContext.withName(childContext.getDisplayName());
+				String name = childContext.getDisplayName();
+				name = name.replaceAll("[-]", "_");
+				GenerationContext childGenerationContext = generationContext.withName(name);
 				ClassName initializerClassName = contextAotGenerator.generateApplicationContext(childContext,
 						childGenerationContext);
 				return Map.entry(childContext.getDisplayName(), initializerClassName);
@@ -143,7 +145,7 @@ public class LoadBalancerChildContextInitializer
 						builder.addParameter(LoadBalancerChildContextInitializer.class, "instance");
 						builder.returns(LoadBalancerChildContextInitializer.class);
 						builder.addStatement(
-								"Map<String, ApplicationContextInitializer<? extends ConfigurableApplicationContext> applicationContextInitializer> initializers = new HashMap();");
+								"Map<String, ApplicationContextInitializer<? extends ConfigurableApplicationContext>> initializers = new HashMap()");
 						generatedInitializerClassNames.keySet()
 								.forEach(contextId -> builder.addStatement("initializers.put($S, new $L())", contextId,
 										generatedInitializerClassNames.get(contextId)));