Support to use classes registered in bootstrap context

spring-cloud-zookeeper-discovery/src/main/java/org/springframework/cloud/zookeeper/discovery/configclient/ZookeeperConfigServerBootstrapper.java

@@ -168,22 +168,22 @@ public List<ServiceInstance> apply(String serviceId, Binder binder, BindHandler
  return Collections.emptyList();
  }
 
- ZookeeperProperties properties = binder.bind(ZookeeperProperties.PREFIX, Bindable.of(ZookeeperProperties.class))
- .orElse(new ZookeeperProperties());
- RetryPolicy retryPolicy = new ExponentialBackoffRetry(properties.getBaseSleepTimeMs(), properties.getMaxRetries(),
- properties.getMaxSleepMs());
+ ZookeeperProperties properties = context.getOrElse(ZookeeperProperties.class, binder.bind(ZookeeperProperties.PREFIX, Bindable.of(ZookeeperProperties.class))
+ .orElse(new ZookeeperProperties()));
+ RetryPolicy retryPolicy = context.getOrElse(RetryPolicy.class, new ExponentialBackoffRetry(properties.getBaseSleepTimeMs(), properties.getMaxRetries(),
+ properties.getMaxSleepMs()));
  try {
- CuratorFramework curatorFramework = CuratorFactory.curatorFramework(properties, retryPolicy, Stream::of,
- () -> null, () -> null);
- InstanceSerializer<ZookeeperInstance> serializer = new JsonInstanceSerializer<>(ZookeeperInstance.class);
- ZookeeperDiscoveryProperties discoveryProperties = binder.bind(ZookeeperDiscoveryProperties.PREFIX, Bindable
+ CuratorFramework curatorFramework = context.getOrElse(CuratorFramework.class, CuratorFactory.curatorFramework(properties, retryPolicy, Stream::of,
+ () -> null, () -> null));
+ InstanceSerializer<ZookeeperInstance> serializer = context.getOrElse(InstanceSerializer.class, new JsonInstanceSerializer<>(ZookeeperInstance.class));
+ ZookeeperDiscoveryProperties discoveryProperties = context.getOrElse(ZookeeperDiscoveryProperties.class, binder.bind(ZookeeperDiscoveryProperties.PREFIX, Bindable
  .of(ZookeeperDiscoveryProperties.class), bindHandler)
- .orElseGet(() -> new ZookeeperDiscoveryProperties(new InetUtils(new InetUtilsProperties())));
- DefaultServiceDiscoveryCustomizer customizer = new DefaultServiceDiscoveryCustomizer(curatorFramework, discoveryProperties, serializer);
+ .orElseGet(() -> new ZookeeperDiscoveryProperties(new InetUtils(new InetUtilsProperties()))));
+ ServiceDiscoveryCustomizer customizer = context.getOrElse(ServiceDiscoveryCustomizer.class, new DefaultServiceDiscoveryCustomizer(curatorFramework, discoveryProperties, serializer));
  ServiceDiscovery<ZookeeperInstance> serviceDiscovery = customizer.customize(ServiceDiscoveryBuilder.builder(ZookeeperInstance.class));
- ZookeeperDependencies dependencies = binder.bind(ZookeeperDependencies.PREFIX, Bindable
+ ZookeeperDependencies dependencies = context.getOrElse(ZookeeperDependencies.class, binder.bind(ZookeeperDependencies.PREFIX, Bindable
  .of(ZookeeperDependencies.class), bindHandler)
- .orElseGet(ZookeeperDependencies::new);
+ .orElseGet(ZookeeperDependencies::new));
  return new ZookeeperDiscoveryClient(serviceDiscovery, dependencies, discoveryProperties).getInstances(serviceId);
  }
  catch (Exception e) {

spring-cloud-zookeeper-discovery/src/test/java/org/springframework/cloud/zookeeper/discovery/configclient/ZookeeperConfigServerBootstrapperCustomizerTests.java

@@ -0,0 +1,171 @@
+/*
+ * Copyright 2015-2023 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.zookeeper.discovery.configclient;
+
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.apache.commons.logging.Log;
+import org.apache.curator.RetryPolicy;
+import org.apache.curator.drivers.TracerDriver;
+import org.apache.curator.ensemble.EnsembleProvider;
+import org.apache.curator.ensemble.fixed.FixedEnsembleProvider;
+import org.apache.curator.framework.CuratorFrameworkFactory;
+import org.apache.curator.retry.ExponentialBackoffRetry;
+import org.apache.curator.utils.DefaultTracerDriver;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+
+import org.springframework.boot.BootstrapRegistry;
+import org.springframework.boot.BootstrapRegistryInitializer;
+import org.springframework.boot.SpringBootConfiguration;
+import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
+import org.springframework.boot.builder.SpringApplicationBuilder;
+import org.springframework.boot.context.properties.bind.BindContext;
+import org.springframework.boot.context.properties.bind.BindHandler;
+import org.springframework.boot.context.properties.bind.Bindable;
+import org.springframework.boot.context.properties.bind.Binder;
+import org.springframework.boot.context.properties.source.ConfigurationPropertyName;
+import org.springframework.cloud.config.client.ConfigServerInstanceProvider;
+import org.springframework.cloud.zookeeper.CuratorFrameworkCustomizer;
+import org.springframework.cloud.zookeeper.ZookeeperProperties;
+import org.springframework.cloud.zookeeper.discovery.ZookeeperDiscoveryClient;
+import org.springframework.cloud.zookeeper.test.ZookeeperTestingServer;
+import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.core.Ordered;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+
+/**
+ * @author Ryan Baxter
+ */
+public class ZookeeperConfigServerBootstrapperCustomizerTests {
+ private ConfigurableApplicationContext context;
+
+ @AfterEach
+ public void after() {
+ if (context != null) {
+ context.close();
+ }
+ }
+
+ @Test
+ public void enabledAddsInstanceProviderFn() {
+ AtomicReference<ZookeeperDiscoveryClient> bootstrapDiscoveryClient = new AtomicReference<>();
+ BindHandlerBootstrapper bindHandlerBootstrapper = new BindHandlerBootstrapper();
+ context = new SpringApplicationBuilder(ZookeeperConfigServerBootstrapperTests.TestConfig.class)
+ .listeners(new ZookeeperTestingServer())
+ .properties("--server.port=0", "spring.cloud.config.discovery.enabled=true",
+ "spring.cloud.zookeeper.discovery.metadata[mymetadataprop]=mymetadataval",
+ "spring.cloud.service-registry.auto-registration.enabled=false")
+ .addBootstrapRegistryInitializer(bindHandlerBootstrapper)
+ .addBootstrapRegistryInitializer(registry -> registry.addCloseListener(event -> {
+ ConfigServerInstanceProvider.Function providerFn = event.getBootstrapContext()
+ .get(ConfigServerInstanceProvider.Function.class);
+ assertThat(providerFn.apply("id", event.getBootstrapContext()
+ .get(Binder.class), event.getBootstrapContext()
+ .get(BindHandler.class), mock(Log.class))).as("Should return empty list.")
+ .isNotNull();
+ bootstrapDiscoveryClient.set(event.getBootstrapContext().get(ZookeeperDiscoveryClient.class));
+ CuratorFrameworkCustomizer curatorFrameworkCustomizer = event.getBootstrapContext()
+ .get(CuratorFrameworkCustomizer.class);
+ assertThat(curatorFrameworkCustomizer).isInstanceOf(MyCuratorFrameworkCustomizer.class);
+ RetryPolicy retryPolicy = event.getBootstrapContext().get(RetryPolicy.class);
+ assertThat(retryPolicy).isInstanceOf(RetryPolicy.class);
+ EnsembleProvider ensembleProvider = event.getBootstrapContext().get(EnsembleProvider.class);
+ assertThat(ensembleProvider).isInstanceOf(MyEnsembleProvider.class);
+ TracerDriver tracerDriver = event.getBootstrapContext().get(TracerDriver.class);
+ assertThat(tracerDriver).isInstanceOf(MyTracerDriver.class);
+ })).run();
+
+ ZookeeperDiscoveryClient discoveryClient = context.getBean(ZookeeperDiscoveryClient.class);
+
+ assertThat(discoveryClient == bootstrapDiscoveryClient.get()).isTrue();
+ assertThat(bindHandlerBootstrapper.onSuccessCount).isGreaterThan(0);
+ }
+
+ @SpringBootConfiguration
+ @EnableAutoConfiguration
+ static class TestConfig {
+
+ }
+
+ static class BindHandlerBootstrapper implements BootstrapRegistryInitializer, Ordered {
+
+ private int onSuccessCount = 0;
+
+ private static <T> void registerIfAbsentAndEnabled(
+ BootstrapRegistry registry, Class<T> type, BootstrapRegistry.InstanceSupplier<T> supplier) {
+ registry.registerIfAbsent(type, supplier);
+ }
+
+ @Override
+ public void initialize(BootstrapRegistry registry) {
+ registry.register(BindHandler.class, context -> new BindHandler() {
+ @Override
+ public Object onSuccess(ConfigurationPropertyName name, Bindable<?> target, BindContext context,
+ Object result) {
+ onSuccessCount++;
+ return result;
+ }
+ });
+
+ registerIfAbsentAndEnabled(registry, RetryPolicy.class, context ->
+ new MyRetryPolicy(context.get(ZookeeperProperties.class)));
+
+ registerIfAbsentAndEnabled(registry, CuratorFrameworkCustomizer.class, context ->
+ new MyCuratorFrameworkCustomizer());
+
+ registerIfAbsentAndEnabled(registry, EnsembleProvider.class, context ->
+ new MyEnsembleProvider(context.get(ZookeeperProperties.class)));
+
+ registerIfAbsentAndEnabled(registry, TracerDriver.class, context ->
+ new MyTracerDriver());
+ }
+
+ @Override
+ public int getOrder() {
+ return Ordered.HIGHEST_PRECEDENCE;
+ }
+ }
+
+ static class MyRetryPolicy extends ExponentialBackoffRetry {
+
+ MyRetryPolicy(ZookeeperProperties properties) {
+ super(properties.getBaseSleepTimeMs(), properties.getMaxRetries(), properties.getMaxSleepMs());
+ }
+ }
+
+ static class MyCuratorFrameworkCustomizer implements CuratorFrameworkCustomizer {
+ @Override
+ public void customize(CuratorFrameworkFactory.Builder builder) {
+
+ }
+ }
+
+ static class MyEnsembleProvider extends FixedEnsembleProvider {
+ MyEnsembleProvider(ZookeeperProperties properties) {
+ super(properties.getConnectString());
+ }
+
+ }
+
+ static class MyTracerDriver extends DefaultTracerDriver {
+
+ }
+
+}