Moves CuratorFramework bootstrap registration to CuratorFactory

spring-cloud-zookeeper-config/src/main/java/org/springframework/cloud/zookeeper/config/ZookeeperConfigDataLocationResolver.java

@@ -21,17 +21,10 @@ import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
-import java.util.function.Supplier;
-import java.util.stream.Stream;
 
 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.framework.CuratorFramework;
 
 import org.springframework.boot.BootstrapRegistry.InstanceSupplier;
-import org.springframework.boot.ConfigurableBootstrapContext;
 import org.springframework.boot.context.config.ConfigDataLocationNotFoundException;
 import org.springframework.boot.context.config.ConfigDataLocationResolver;
 import org.springframework.boot.context.config.ConfigDataLocationResolverContext;
@@ -39,7 +32,6 @@ import org.springframework.boot.context.config.Profiles;
 import org.springframework.boot.context.properties.bind.Bindable;
 import org.springframework.boot.context.properties.bind.Binder;
 import org.springframework.cloud.zookeeper.CuratorFactory;
-import org.springframework.cloud.zookeeper.CuratorFrameworkCustomizer;
 import org.springframework.cloud.zookeeper.ZookeeperProperties;
 import org.springframework.core.env.MapPropertySource;
 import org.springframework.lang.Nullable;
@@ -86,14 +78,7 @@ public class ZookeeperConfigDataLocationResolver implements ConfigDataLocationRe
 		UriComponents locationUri = parseLocation(location);
 
 		// create curator
-		ZookeeperProperties zookeeperProperties = loadProperties(context.getBinder(), locationUri);
-		context.getBootstrapContext().register(ZookeeperProperties.class, InstanceSupplier.of(zookeeperProperties));
-
-		context.getBootstrapContext().registerIfAbsent(RetryPolicy.class,
-				InstanceSupplier.from(() -> CuratorFactory.retryPolicy(zookeeperProperties)));
-
-		context.getBootstrapContext().registerIfAbsent(CuratorFramework.class, InstanceSupplier
-				.from(() -> curatorFramework(context.getBootstrapContext(), zookeeperProperties, optional)));
+		CuratorFactory.registerCurator(context.getBootstrapContext(), locationUri, optional);
 
 		// create locations
 		ZookeeperConfigProperties properties = loadConfigProperties(context.getBinder());
@@ -106,9 +91,6 @@ public class ZookeeperConfigDataLocationResolver implements ConfigDataLocationRe
 
 		// promote beans to context
 		context.getBootstrapContext().addCloseListener(event -> {
-			CuratorFramework curatorFramework = event.getBootstrapContext().get(CuratorFramework.class);
-			event.getApplicationContext().getBeanFactory().registerSingleton("configDataCuratorFramework",
-					curatorFramework);
 			HashMap<String, Object> source = new HashMap<>();
 			source.put("spring.cloud.zookeeper.config.property-source-contexts", contexts);
 			MapPropertySource propertySource = new MapPropertySource("zookeeperConfigData", source);
@@ -145,51 +127,6 @@ public class ZookeeperConfigDataLocationResolver implements ConfigDataLocationRe
 		return UriComponentsBuilder.fromUriString(uri).build();
 	}
 
-	protected CuratorFramework curatorFramework(ConfigurableBootstrapContext context, ZookeeperProperties properties,
-			boolean optional) {
-
-		try {
-			Supplier<Stream<CuratorFrameworkCustomizer>> customizers;
-			if (context.isRegistered(CuratorFrameworkCustomizer.class)) {
-				customizers = () -> Stream.of(context.get(CuratorFrameworkCustomizer.class));
-			}
-			else {
-				customizers = () -> null;
-			}
-			return CuratorFactory.curatorFramework(properties, context.get(RetryPolicy.class), customizers,
-					supplier(context, EnsembleProvider.class), supplier(context, TracerDriver.class));
-		}
-		catch (Exception e) {
-			if (!optional) {
-				log.error("Unable to connect to zookeeper", e);
-				throw new ConfigDataLocationNotFoundException("Unable to connect to zookeeper", null, e);
-			}
-			else if (log.isDebugEnabled()) {
-				log.debug("Unable to connect to zookeeper", e);
-			}
-		}
-		return null;
-	}
-
-	private <T> Supplier<T> supplier(ConfigurableBootstrapContext context, Class<T> type) {
-		return () -> context.isRegistered(type) ? context.get(type) : null;
-	}
-
-	protected ZookeeperProperties loadProperties(Binder binder, UriComponents location) {
-		ZookeeperProperties properties = binder.bind(ZookeeperProperties.PREFIX, Bindable.of(ZookeeperProperties.class))
-				.orElse(new ZookeeperProperties());
-
-		if (location != null && StringUtils.hasText(location.getHost())) {
-			if (location.getPort() < 0) {
-				throw new IllegalArgumentException(
-						"zookeeper port must be greater than or equal to zero: " + location.getPort());
-			}
-			properties.setConnectString(location.getHost() + ":" + location.getPort());
-		}
-
-		return properties;
-	}
-
 	protected ZookeeperConfigProperties loadConfigProperties(Binder binder) {
 		ZookeeperConfigProperties properties = binder
 				.bind(ZookeeperConfigProperties.PREFIX, Bindable.of(ZookeeperConfigProperties.class))

spring-cloud-zookeeper-config/src/test/java/org/springframework/cloud/zookeeper/config/ZookeeperConfigDataLocationResolverTests.java

@@ -27,10 +27,8 @@ import org.springframework.boot.ConfigurableBootstrapContext;
 import org.springframework.boot.context.config.ConfigDataLocationResolverContext;
 import org.springframework.boot.context.config.Profiles;
 import org.springframework.boot.context.properties.bind.Binder;
-import org.springframework.cloud.zookeeper.ZookeeperProperties;
 import org.springframework.mock.env.MockEnvironment;
 import org.springframework.web.util.UriComponents;
-import org.springframework.web.util.UriComponentsBuilder;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.mock;
@@ -69,14 +67,6 @@ public class ZookeeperConfigDataLocationResolverTests {
 				"config/testapp", "config/application,dev", "config/application");
 	}
 
-	@Test
-	public void testLoadProperties() {
-		ZookeeperProperties properties = createResolver().loadProperties(
-				Binder.get(new MockEnvironment()),
-				UriComponentsBuilder.fromUriString("zookeeper://myhost:8502").build());
-		assertThat(properties.getConnectString()).isEqualTo("myhost:8502");
-	}
-
 	private List<String> toContexts(List<ZookeeperConfigDataLocation> locations) {
 		return locations.stream().map(ZookeeperConfigDataLocation::getContext)
 				.collect(Collectors.toList());

spring-cloud-zookeeper-core/src/main/java/org/springframework/cloud/zookeeper/CuratorFactory.java

@@ -28,16 +28,22 @@ import org.apache.curator.framework.CuratorFramework;
 import org.apache.curator.framework.CuratorFrameworkFactory;
 import org.apache.curator.retry.ExponentialBackoffRetry;
 
+import org.springframework.boot.BootstrapContext;
+import org.springframework.boot.BootstrapRegistry;
+import org.springframework.boot.context.config.ConfigDataLocationNotFoundException;
+import org.springframework.boot.context.properties.bind.Bindable;
+import org.springframework.boot.context.properties.bind.Binder;
+import org.springframework.util.StringUtils;
+import org.springframework.web.util.UriComponents;
 
 public abstract class CuratorFactory {
+
 	private static final Log log = LogFactory.getLog(ZookeeperAutoConfiguration.class);
 
-	public static CuratorFramework curatorFramework(
-			ZookeeperProperties properties,
-			RetryPolicy retryPolicy,
+	public static CuratorFramework curatorFramework(ZookeeperProperties properties, RetryPolicy retryPolicy,
 			Supplier<Stream<CuratorFrameworkCustomizer>> optionalCuratorFrameworkCustomizerProvider,
-			Supplier<EnsembleProvider> optionalEnsembleProvider,
-			Supplier<TracerDriver> optionalTracerDriverProvider) throws Exception {
+			Supplier<EnsembleProvider> optionalEnsembleProvider, Supplier<TracerDriver> optionalTracerDriverProvider)
+			throws Exception {
 		CuratorFrameworkFactory.Builder builder = CuratorFrameworkFactory.builder();
 
 		EnsembleProvider ensembleProvider = optionalEnsembleProvider.get();
@@ -48,15 +54,11 @@ public abstract class CuratorFactory {
 			builder.connectString(properties.getConnectString());
 		}
 		builder.sessionTimeoutMs((int) properties.getSessionTimeout().toMillis())
-				.connectionTimeoutMs((int) properties.getConnectionTimeout().toMillis())
-				.retryPolicy(retryPolicy);
+				.connectionTimeoutMs((int) properties.getConnectionTimeout().toMillis()).retryPolicy(retryPolicy);
 
-		Stream<CuratorFrameworkCustomizer> customizers = optionalCuratorFrameworkCustomizerProvider
-				.get();
+		Stream<CuratorFrameworkCustomizer> customizers = optionalCuratorFrameworkCustomizerProvider.get();
 		if (customizers != null) {
-			customizers
-					.forEach(curatorFrameworkCustomizer -> curatorFrameworkCustomizer
-							.customize(builder));
+			customizers.forEach(curatorFrameworkCustomizer -> curatorFrameworkCustomizer.customize(builder));
 		}
 
 		CuratorFramework curator = builder.build();
@@ -67,12 +69,10 @@ public abstract class CuratorFactory {
 
 		curator.start();
 		if (log.isTraceEnabled()) {
-			log.trace("blocking until connected to zookeeper for "
-					+ properties.getBlockUntilConnectedWait()
+			log.trace("blocking until connected to zookeeper for " + properties.getBlockUntilConnectedWait()
 					+ properties.getBlockUntilConnectedUnit());
 		}
-		curator.blockUntilConnected(properties.getBlockUntilConnectedWait(),
-				properties.getBlockUntilConnectedUnit());
+		curator.blockUntilConnected(properties.getBlockUntilConnectedWait(), properties.getBlockUntilConnectedUnit());
 		if (log.isTraceEnabled()) {
 			log.trace("connected to zookeeper");
 		}
@@ -83,4 +83,77 @@ public abstract class CuratorFactory {
 		return new ExponentialBackoffRetry(properties.getBaseSleepTimeMs(), properties.getMaxRetries(),
 				properties.getMaxSleepMs());
 	}
+
+	public static void registerCurator(BootstrapRegistry registery, UriComponents location, boolean optional) {
+		registery.registerIfAbsent(ZookeeperProperties.class,
+				context -> loadProperties(context.get(Binder.class), location));
+
+		registery.registerIfAbsent(RetryPolicy.class, context -> retryPolicy(context.get(ZookeeperProperties.class)));
+
+		registery.registerIfAbsent(CuratorFramework.class,
+				context -> curatorFramework(context, context.get(ZookeeperProperties.class), optional));
+
+		// promote beans to context
+		registery.addCloseListener(event -> {
+			CuratorFramework curatorFramework = event.getBootstrapContext().get(CuratorFramework.class);
+			event.getApplicationContext().getBeanFactory().registerSingleton("configDataCuratorFramework",
+					curatorFramework);
+		});
+
+	}
+
+	static ZookeeperProperties loadProperties(Binder binder, UriComponents location) {
+		ZookeeperProperties properties = binder.bind(ZookeeperProperties.PREFIX, Bindable.of(ZookeeperProperties.class))
+				.orElse(new ZookeeperProperties());
+
+		if (location != null && StringUtils.hasText(location.getHost())) {
+			if (location.getPort() < 0) {
+				throw new IllegalArgumentException(
+						"zookeeper port must be greater than or equal to zero: " + location.getPort());
+			}
+			properties.setConnectString(location.getHost() + ":" + location.getPort());
+		}
+
+		return properties;
+	}
+
+	private static CuratorFramework curatorFramework(BootstrapContext context, ZookeeperProperties properties,
+			boolean optional) {
+
+		Supplier<Stream<CuratorFrameworkCustomizer>> customizers;
+		// TODO: use new apis after milestone release
+		try {
+			CuratorFrameworkCustomizer customizer = context.get(CuratorFrameworkCustomizer.class);
+			customizers = () -> Stream.of(customizer);
+		}
+		catch (IllegalStateException e) {
+			customizers = () -> null;
+		}
+		try {
+			return CuratorFactory.curatorFramework(properties, context.get(RetryPolicy.class), customizers,
+					supplier(context, EnsembleProvider.class), supplier(context, TracerDriver.class));
+		}
+		catch (Exception e) {
+			if (!optional) {
+				log.error("Unable to connect to zookeeper", e);
+				throw new ConfigDataLocationNotFoundException("Unable to connect to zookeeper", null, e);
+			}
+			if (log.isDebugEnabled()) {
+				log.debug("Unable to connect to zookeeper", e);
+			}
+		}
+		return null;
+	}
+
+	private static <T> Supplier<T> supplier(BootstrapContext context, Class<T> type) {
+		try {
+			// TODO: use new apis after milestone release
+			T instance = context.get(type);
+			return () -> instance;
+		}
+		catch (IllegalStateException e) {
+			return () -> null;
+		}
+	}
+
 }

spring-cloud-zookeeper-core/src/test/java/org/springframework/cloud/zookeeper/CuratorFactoryTests.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2015-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.zookeeper;
+
+import org.junit.jupiter.api.Test;
+
+import org.springframework.boot.context.properties.bind.Binder;
+import org.springframework.mock.env.MockEnvironment;
+import org.springframework.web.util.UriComponentsBuilder;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class CuratorFactoryTests {
+
+	@Test
+	public void testLoadProperties() {
+		ZookeeperProperties properties = CuratorFactory.loadProperties(
+				Binder.get(new MockEnvironment()),
+				UriComponentsBuilder.fromUriString("zookeeper://myhost:8502").build());
+		assertThat(properties.getConnectString()).isEqualTo("myhost:8502");
+	}
+
+}

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

@@ -16,21 +16,12 @@
 
 package org.springframework.cloud.zookeeper.discovery.configclient;
 
-import java.util.function.Supplier;
-import java.util.stream.Stream;
-
-import org.apache.commons.logging.Log;
-import org.apache.commons.logging.LogFactory;
-import org.apache.curator.RetryPolicy;
-import org.apache.curator.drivers.TracerDriver;
-import org.apache.curator.ensemble.EnsembleProvider;
 import org.apache.curator.framework.CuratorFramework;
 import org.apache.curator.x.discovery.ServiceDiscovery;
 import org.apache.curator.x.discovery.ServiceDiscoveryBuilder;
 import org.apache.curator.x.discovery.details.InstanceSerializer;
 import org.apache.curator.x.discovery.details.JsonInstanceSerializer;
 
-import org.springframework.boot.BootstrapContext;
 import org.springframework.boot.BootstrapRegistry;
 import org.springframework.boot.Bootstrapper;
 import org.springframework.boot.context.properties.bind.Binder;
@@ -39,8 +30,6 @@ import org.springframework.cloud.commons.util.InetUtilsProperties;
 import org.springframework.cloud.config.client.ConfigClientProperties;
 import org.springframework.cloud.config.client.ConfigServerInstanceProvider;
 import org.springframework.cloud.zookeeper.CuratorFactory;
-import org.springframework.cloud.zookeeper.CuratorFrameworkCustomizer;
-import org.springframework.cloud.zookeeper.ZookeeperProperties;
 import org.springframework.cloud.zookeeper.discovery.ZookeeperDiscoveryClient;
 import org.springframework.cloud.zookeeper.discovery.ZookeeperDiscoveryProperties;
 import org.springframework.cloud.zookeeper.discovery.ZookeeperInstance;
@@ -51,8 +40,6 @@ import org.springframework.util.ClassUtils;
 
 public class ZookeeperConfigServerBootstrapper implements Bootstrapper {
 
-	private final Log log = LogFactory.getLog(getClass());
-
 	@Override
 	@SuppressWarnings("unchecked")
 	public void intitialize(BootstrapRegistry registry) {
@@ -60,16 +47,7 @@ public class ZookeeperConfigServerBootstrapper implements Bootstrapper {
 			return;
 		}
 		// create curator
-		registry.registerIfAbsent(ZookeeperProperties.class, context -> context.get(Binder.class)
-				.bind(ZookeeperProperties.PREFIX, ZookeeperProperties.class).orElseGet(ZookeeperProperties::new));
-		registry.registerIfAbsent(RetryPolicy.class, context -> {
-			ZookeeperProperties properties = context.get(ZookeeperProperties.class);
-			return CuratorFactory.retryPolicy(properties);
-		});
-		registry.registerIfAbsent(CuratorFramework.class, context -> {
-			ZookeeperProperties zookeeperProperties = context.get(ZookeeperProperties.class);
-			return curatorFramework(context, zookeeperProperties);
-		});
+		CuratorFactory.registerCurator(registry, null, true);
 
 		// create discovery
 		registry.registerIfAbsent(ZookeeperDiscoveryProperties.class,
@@ -123,38 +101,4 @@ public class ZookeeperConfigServerBootstrapper implements Bootstrapper {
 		return binder.bind(ConfigClientProperties.CONFIG_DISCOVERY_ENABLED, Boolean.class).orElse(false);
 	}
 
-	protected CuratorFramework curatorFramework(BootstrapContext context, ZookeeperProperties properties) {
-
-		Supplier<Stream<CuratorFrameworkCustomizer>> customizers;
-		// TODO: use new apis after milestone release
-		try {
-			CuratorFrameworkCustomizer customizer = context.get(CuratorFrameworkCustomizer.class);
-			customizers = () -> Stream.of(customizer);
-		}
-		catch (IllegalStateException e) {
-			customizers = () -> null;
-		}
-		try {
-			return CuratorFactory.curatorFramework(properties, context.get(RetryPolicy.class), customizers,
-					supplier(context, EnsembleProvider.class), supplier(context, TracerDriver.class));
-		}
-		catch (Exception e) {
-			if (log.isDebugEnabled()) {
-				log.debug("Unable to connect to zookeeper", e);
-			}
-		}
-		return null;
-	}
-
-	private <T> Supplier<T> supplier(BootstrapContext context, Class<T> type) {
-		try {
-			// TODO: use new apis after milestone release
-			T instance = context.get(type);
-			return () -> instance;
-		}
-		catch (IllegalStateException e) {
-			return () -> null;
-		}
-	}
-
 }