Updates to use new Bootstrap apis in boot.

Creates a ZookeeperBootstrapper for users to customize Zookeeper.

Moves CuratorFramework create to CuratorFactory that is used in auto configuration and ConfigData

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

@@ -0,0 +1,82 @@
+/*
+ * 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.config;
+
+import java.util.function.Function;
+
+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.BootstrapContext;
+import org.springframework.boot.BootstrapRegistry;
+import org.springframework.boot.Bootstrapper;
+import org.springframework.cloud.zookeeper.CuratorFrameworkCustomizer;
+
+public class ZookeeperBootstrapper implements Bootstrapper {
+
+	private Function<BootstrapContext, RetryPolicy> retryPolicy;
+
+	private Function<BootstrapContext, EnsembleProvider> ensembleProvider;
+
+	private Function<BootstrapContext, TracerDriver> tracerDriver;
+
+	private Function<BootstrapContext, CuratorFrameworkCustomizer> curatorFrameworkCustomizer;
+
+	static Bootstrapper fromBootstrapContext(Function<BootstrapContext, CuratorFramework> factory) {
+		return registry -> registry.register(CuratorFramework.class, factory::apply);
+	}
+
+	static ZookeeperBootstrapper create() {
+		return new ZookeeperBootstrapper();
+	}
+
+	public ZookeeperBootstrapper retryPolicy(Function<BootstrapContext, RetryPolicy> retryPolicy) {
+		this.retryPolicy = retryPolicy;
+		return this;
+	}
+
+	public ZookeeperBootstrapper ensembleProvider(Function<BootstrapContext, EnsembleProvider> ensembleProvider) {
+		this.ensembleProvider = ensembleProvider;
+		return this;
+	}
+
+	public ZookeeperBootstrapper tracerDriver(Function<BootstrapContext, TracerDriver> tracerDriver) {
+		this.tracerDriver = tracerDriver;
+		return this;
+	}
+
+	public ZookeeperBootstrapper curatorFrameworkCustomizer(Function<BootstrapContext, CuratorFrameworkCustomizer> curatorFrameworkCustomizer) {
+		this.curatorFrameworkCustomizer = curatorFrameworkCustomizer;
+		return this;
+	}
+
+	@Override
+	public void intitialize(BootstrapRegistry registry) {
+		register(registry, RetryPolicy.class, retryPolicy);
+		register(registry, EnsembleProvider.class, ensembleProvider);
+		register(registry, TracerDriver.class, tracerDriver);
+		register(registry, CuratorFrameworkCustomizer.class, curatorFrameworkCustomizer);
+	}
+
+	private <T> void register(BootstrapRegistry registry, Class<T> type, Function<BootstrapContext, T> factory) {
+		if (this.retryPolicy != null) {
+			registry.register(type, factory::apply);
+		}
+	}
+}

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

@@ -34,14 +34,15 @@ public class ZookeeperConfigDataLoader implements ConfigDataLoader<ZookeeperConf
 	@Override
 	public ConfigData load(ConfigDataLoaderContext context, ZookeeperConfigDataLocation location) {
 		try {
-			CuratorFramework curator = context.getBootstrapRegistry().getRegistration(CuratorFramework.class)
-					.get();
+			CuratorFramework curator = context.getBootstrapContext().get(CuratorFramework.class);
 			ZookeeperPropertySource propertySource = new ZookeeperPropertySource(location.getContext(),
 					curator);
 			return new ConfigData(Collections.singletonList(propertySource));
 		}
 		catch (Exception e) {
-			if (location.getProperties().isFailFast() || !location.isOptional()) {
+			ZookeeperConfigProperties properties = context.getBootstrapContext()
+					.get(ZookeeperConfigProperties.class);
+			if (properties.isFailFast() || !location.isOptional()) {
 				throw new ConfigDataLocationNotFoundException(location, e);
 			}
 			else {

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

@@ -23,20 +23,14 @@ import org.springframework.core.style.ToStringCreator;
 
 public class ZookeeperConfigDataLocation extends ConfigDataLocation {
 
-	private final ZookeeperConfigProperties properties;
 	private final String context;
 	private final boolean optional;
 
-	public ZookeeperConfigDataLocation(ZookeeperConfigProperties properties, String context, boolean optional) {
-		this.properties = properties;
+	public ZookeeperConfigDataLocation(String context, boolean optional) {
 		this.context = context;
 		this.optional = optional;
 	}
 
-	public ZookeeperConfigProperties getProperties() {
-		return this.properties;
-	}
-
 	public String getContext() {
 		return this.context;
 	}
@@ -54,14 +48,13 @@ public class ZookeeperConfigDataLocation extends ConfigDataLocation {
 			return false;
 		}
 		ZookeeperConfigDataLocation that = (ZookeeperConfigDataLocation) o;
-		return this.properties.equals(that.properties) &&
-				this.optional == that.optional &&
+		return this.optional == that.optional &&
 				this.context.equals(that.context);
 	}
 
 	@Override
 	public int hashCode() {
-		return Objects.hash(this.optional, this.properties, this.context);
+		return Objects.hash(this.optional, this.context);
 	}
 
 	@Override
@@ -69,7 +62,6 @@ public class ZookeeperConfigDataLocation extends ConfigDataLocation {
 		return new ToStringCreator(this)
 				.append("context", context)
 				.append("optional", optional)
-				.append("properties", properties)
 				.toString();
 
 	}

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

@@ -21,20 +21,26 @@ 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.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.framework.CuratorFrameworkFactory;
-import org.apache.curator.retry.ExponentialBackoffRetry;
 
+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;
 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;
@@ -74,32 +80,41 @@ public class ZookeeperConfigDataLocationResolver implements ConfigDataLocationRe
 	@Override
 	public List<ZookeeperConfigDataLocation> resolveProfileSpecific(ConfigDataLocationResolverContext context,
 			String location, boolean optional, Profiles profiles) throws ConfigDataLocationNotFoundException {
-		UriComponents locationUri = parseLocation(context, location);
+		UriComponents locationUri = parseLocation(location);
 
 		ZookeeperConfigProperties properties = loadConfigProperties(context.getBinder());
+		context.getBootstrapContext().register(ZookeeperConfigProperties.class, InstanceSupplier.of(properties));
+
+		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)));
 
 		List<String> contexts = (locationUri == null || CollectionUtils.isEmpty(locationUri.getPathSegments()))
-				? getAutomaticContexts(profiles, properties) : getCustomContexts(locationUri, properties);
+				? getAutomaticContexts(profiles, properties) : getCustomContexts(locationUri);
 
-		context.getBootstrapRegistry()
-				.register(CuratorFramework.class,
-						() -> curatorFramework(optional, loadProperties(context.getBinder(), locationUri)))
-				.onApplicationContextPrepared((ctxt, curatorFramework) -> {
-					ctxt.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);
-					ctxt.getEnvironment().getPropertySources().addFirst(propertySource);
-				});
+		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);
+			event.getApplicationContext().getEnvironment().getPropertySources().addFirst(propertySource);
+		});
 
 		ArrayList<ZookeeperConfigDataLocation> locations = new ArrayList<>();
 		contexts.forEach(propertySourceContext -> locations
-				.add(new ZookeeperConfigDataLocation(properties, propertySourceContext, optional)));
+				.add(new ZookeeperConfigDataLocation(propertySourceContext, optional)));
 
 		return locations;
 	}
 
-	protected List<String> getCustomContexts(UriComponents uriComponents, ZookeeperConfigProperties properties) {
+	protected List<String> getCustomContexts(UriComponents uriComponents) {
 		if (StringUtils.isEmpty(uriComponents.getPath())) {
 			return Collections.emptyList();
 		}
@@ -129,7 +144,7 @@ public class ZookeeperConfigDataLocationResolver implements ConfigDataLocationRe
 	}
 
 	@Nullable
-	protected UriComponents parseLocation(ConfigDataLocationResolverContext context, String location) {
+	protected UriComponents parseLocation(String location) {
 		String uri = location.substring(PREFIX.length());
 		if (!StringUtils.hasText(uri)) {
 			return null;
@@ -150,26 +165,21 @@ public class ZookeeperConfigDataLocationResolver implements ConfigDataLocationRe
 		}
 	}
 
-	protected CuratorFramework curatorFramework(boolean optional, ZookeeperProperties properties) {
-		CuratorFrameworkFactory.Builder builder = CuratorFrameworkFactory.builder();
+	protected CuratorFramework curatorFramework(ConfigurableBootstrapContext context, ZookeeperProperties properties,
+			boolean optional) {
 
-		builder.connectString(properties.getConnectString())
-				.sessionTimeoutMs((int) properties.getSessionTimeout().toMillis())
-				.connectionTimeoutMs((int) properties.getConnectionTimeout().toMillis())
-				.retryPolicy(retryPolicy(properties));
-
-		CuratorFramework curator = builder.build();
-
-		curator.start();
-		if (log.isTraceEnabled()) {
-			log.trace("blocking until connected to zookeeper for " + properties.getBlockUntilConnectedWait()
-					+ properties.getBlockUntilConnectedUnit());
-		}
 		try {
-			curator.blockUntilConnected(properties.getBlockUntilConnectedWait(),
-					properties.getBlockUntilConnectedUnit());
+			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 (InterruptedException e) {
+		catch (Exception e) {
 			if (!optional) {
 				log.error("Unable to connect to zookeeper", e);
 				throw new ConfigDataLocationNotFoundException("Unable to connect to zookeeper", null, e);
@@ -178,15 +188,11 @@ public class ZookeeperConfigDataLocationResolver implements ConfigDataLocationRe
 				log.debug("Unable to connect to zookeeper", e);
 			}
 		}
-		if (log.isTraceEnabled()) {
-			log.trace("connected to zookeeper");
-		}
-		return curator;
+		return null;
 	}
 
-	protected RetryPolicy retryPolicy(ZookeeperProperties properties) {
-		return new ExponentialBackoffRetry(properties.getBaseSleepTimeMs(), properties.getMaxRetries(),
-				properties.getMaxSleepMs());
+	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) {
@@ -195,7 +201,8 @@ public class ZookeeperConfigDataLocationResolver implements ConfigDataLocationRe
 
 		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());
+				throw new IllegalArgumentException(
+						"zookeeper port must be greater than or equal to zero: " + location.getPort());
 			}
 			properties.setConnectString(location.getHost() + ":" + location.getPort());
 		}

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

@@ -0,0 +1,120 @@
+/*
+ * Copyright 2015-2019 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.config;
+
+import java.util.UUID;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.apache.curator.framework.CuratorFramework;
+import org.apache.curator.framework.CuratorFrameworkFactory;
+import org.apache.curator.framework.imps.CuratorFrameworkImpl;
+import org.apache.curator.test.TestingServer;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import org.springframework.boot.BootstrapContext;
+import org.springframework.boot.WebApplicationType;
+import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
+import org.springframework.boot.builder.SpringApplicationBuilder;
+import org.springframework.cloud.zookeeper.CuratorFactory;
+import org.springframework.cloud.zookeeper.ZookeeperProperties;
+import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.util.ReflectionUtils;
+import org.springframework.util.SocketUtils;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * @author Spencer Gibb
+ */
+public class ZookeeperConfigDataCustomizationIntegrationTests {
+
+	private static final Log log = LogFactory
+			.getLog(ZookeeperConfigDataCustomizationIntegrationTests.class);
+
+	public static final String PREFIX = "test__configdata__";
+
+	public static final String ROOT = "/" + PREFIX + UUID.randomUUID();
+
+	private ConfigurableApplicationContext context;
+
+	private TestingServer testingServer;
+
+	@Before
+	public void setup() throws Exception {
+		int port = SocketUtils.findAvailableTcpPort();
+		this.testingServer = new TestingServer(port);
+		String connectString = "localhost:" + port;
+
+		this.context = new SpringApplicationBuilder(Config.class)
+				.web(WebApplicationType.NONE)
+				.addBootstrapper(ZookeeperBootstrapper.fromBootstrapContext(this::curatorFramework))
+				.run("--spring.config.import=zookeeper:" + connectString,
+						"--spring.application.name=testZkConfigDataIntegration",
+						"--logging.level.org.springframework.cloud.zookeeper=DEBUG",
+						"--spring.cloud.zookeeper.config.root=" + ROOT);
+	}
+
+	@After
+	public void after() throws Exception {
+		if (context != null) {
+			this.context.close();
+		}
+		if (testingServer != null) {
+			this.testingServer.close();
+		}
+	}
+
+	CuratorFramework curatorFramework(BootstrapContext context) {
+		ZookeeperProperties properties = context.get(ZookeeperProperties.class);
+		CuratorFrameworkFactory.Builder builder = CuratorFrameworkFactory.builder()
+				.retryPolicy(CuratorFactory.retryPolicy(properties))
+				.connectString(properties.getConnectString());
+		TestCuratorFramework curator = new TestCuratorFramework(builder);
+		curator.start();
+		try {
+			curator.blockUntilConnected(properties.getBlockUntilConnectedWait(),
+					properties.getBlockUntilConnectedUnit());
+		}
+		catch (InterruptedException e) {
+			ReflectionUtils.rethrowRuntimeException(e);
+		}
+		return curator;
+	}
+
+	@Test
+	public void curatorFrameworkIsCustom() {
+		CuratorFramework curator = context.getBean(CuratorFramework.class);
+		assertThat(curator).isNotNull().isInstanceOf(TestCuratorFramework.class);
+	}
+
+	static class TestCuratorFramework extends CuratorFrameworkImpl {
+		TestCuratorFramework(CuratorFrameworkFactory.Builder builder) {
+			super(builder);
+		}
+	}
+
+	@Configuration
+	@EnableAutoConfiguration
+	static class Config {
+
+	}
+
+}

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

@@ -22,17 +22,16 @@ import java.util.stream.Collectors;
 
 import org.junit.jupiter.api.Test;
 
+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.boot.env.BootstrapRegistry;
 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.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -41,12 +40,12 @@ public class ZookeeperConfigDataLocationResolverTests {
 	@Test
 	public void testParseLocation() {
 		ZookeeperConfigDataLocationResolver resolver = new ZookeeperConfigDataLocationResolver();
-		UriComponents uriComponents = resolver.parseLocation(null,
+		UriComponents uriComponents = resolver.parseLocation(
 				"zookeeper:myhost:2182/mypath1;/mypath2;/mypath3");
 		assertThat(uriComponents.toUri()).hasScheme("zookeeper").hasHost("myhost")
 				.hasPort(2182).hasPath("/mypath1;/mypath2;/mypath3");
 
-		uriComponents = resolver.parseLocation(null, "zookeeper:myhost:2182");
+		uriComponents = resolver.parseLocation("zookeeper:myhost:2182");
 		assertThat(uriComponents.toUri()).hasScheme("zookeeper").hasHost("myhost")
 				.hasPort(2182).hasPath("");
 	}
@@ -62,7 +61,7 @@ public class ZookeeperConfigDataLocationResolverTests {
 
 	@Test
 	public void testResolveProfileSpecificWithAutomaticPaths() {
-		String location = "zookeeper:myhost";
+		String location = "zookeeper:myhost:1234";
 		List<ZookeeperConfigDataLocation> locations = testResolveProfileSpecific(location);
 		assertThat(locations).hasSize(4);
 		assertThat(toContexts(locations)).containsExactly("config/testapp,dev",
@@ -88,13 +87,12 @@ public class ZookeeperConfigDataLocationResolverTests {
 		MockEnvironment env = new MockEnvironment();
 		env.setProperty("spring.application.name", "testapp");
 
-		BootstrapRegistry registry = mock(BootstrapRegistry.class);
-		when(registry.register(any(), any())).thenReturn(mock(BootstrapRegistry.Registration.class));
+		ConfigurableBootstrapContext bootstrapContext = mock(ConfigurableBootstrapContext.class);
 
 		ConfigDataLocationResolverContext context = mock(
 				ConfigDataLocationResolverContext.class);
 
-		when(context.getBootstrapRegistry()).thenReturn(registry);
+		when(context.getBootstrapContext()).thenReturn(bootstrapContext);
 		when(context.getBinder()).thenReturn(Binder.get(env));
 
 		Profiles profiles = mock(Profiles.class);

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

@@ -0,0 +1,86 @@
+/*
+ * 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 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.framework.CuratorFrameworkFactory;
+import org.apache.curator.retry.ExponentialBackoffRetry;
+
+
+public abstract class CuratorFactory {
+	private static final Log log = LogFactory.getLog(ZookeeperAutoConfiguration.class);
+
+	public static CuratorFramework curatorFramework(
+			ZookeeperProperties properties,
+			RetryPolicy retryPolicy,
+			Supplier<Stream<CuratorFrameworkCustomizer>> optionalCuratorFrameworkCustomizerProvider,
+			Supplier<EnsembleProvider> optionalEnsembleProvider,
+			Supplier<TracerDriver> optionalTracerDriverProvider) throws Exception {
+		CuratorFrameworkFactory.Builder builder = CuratorFrameworkFactory.builder();
+
+		EnsembleProvider ensembleProvider = optionalEnsembleProvider.get();
+		if (ensembleProvider != null) {
+			builder.ensembleProvider(ensembleProvider);
+		}
+		else {
+			builder.connectString(properties.getConnectString());
+		}
+		builder.sessionTimeoutMs((int) properties.getSessionTimeout().toMillis())
+				.connectionTimeoutMs((int) properties.getConnectionTimeout().toMillis())
+				.retryPolicy(retryPolicy);
+
+		Stream<CuratorFrameworkCustomizer> customizers = optionalCuratorFrameworkCustomizerProvider
+				.get();
+		if (customizers != null) {
+			customizers
+					.forEach(curatorFrameworkCustomizer -> curatorFrameworkCustomizer
+							.customize(builder));
+		}
+
+		CuratorFramework curator = builder.build();
+		TracerDriver tracerDriver = optionalTracerDriverProvider.get();
+		if (tracerDriver != null && curator.getZookeeperClient() != null) {
+			curator.getZookeeperClient().setTracerDriver(tracerDriver);
+		}
+
+		curator.start();
+		if (log.isTraceEnabled()) {
+			log.trace("blocking until connected to zookeeper for "
+					+ properties.getBlockUntilConnectedWait()
+					+ properties.getBlockUntilConnectedUnit());
+		}
+		curator.blockUntilConnected(properties.getBlockUntilConnectedWait(),
+				properties.getBlockUntilConnectedUnit());
+		if (log.isTraceEnabled()) {
+			log.trace("connected to zookeeper");
+		}
+		return curator;
+	}
+
+	public static RetryPolicy retryPolicy(ZookeeperProperties properties) {
+		return new ExponentialBackoffRetry(properties.getBaseSleepTimeMs(), properties.getMaxRetries(),
+				properties.getMaxSleepMs());
+	}
+}

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

@@ -22,8 +22,6 @@ 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.framework.CuratorFrameworkFactory;
-import org.apache.curator.retry.ExponentialBackoffRetry;
 
 import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
@@ -53,54 +51,19 @@ public class ZookeeperAutoConfiguration {
 
 	@Bean(destroyMethod = "close")
 	@ConditionalOnMissingBean
-	public CuratorFramework curatorFramework(RetryPolicy retryPolicy,
-			ZookeeperProperties properties,
+	public CuratorFramework curatorFramework(ZookeeperProperties properties, RetryPolicy retryPolicy,
 			ObjectProvider<CuratorFrameworkCustomizer> optionalCuratorFrameworkCustomizerProvider,
 			ObjectProvider<EnsembleProvider> optionalEnsembleProvider,
 			ObjectProvider<TracerDriver> optionalTracerDriverProvider) throws Exception {
-		CuratorFrameworkFactory.Builder builder = CuratorFrameworkFactory.builder();
-
-		EnsembleProvider ensembleProvider = optionalEnsembleProvider.getIfAvailable();
-		if (ensembleProvider != null) {
-			builder.ensembleProvider(ensembleProvider);
-		}
-		else {
-			builder.connectString(properties.getConnectString());
-		}
-		builder.sessionTimeoutMs((int) properties.getSessionTimeout().toMillis())
-				.connectionTimeoutMs((int) properties.getConnectionTimeout().toMillis())
-				.retryPolicy(retryPolicy);
-
-		optionalCuratorFrameworkCustomizerProvider.orderedStream()
-				.forEach(curatorFrameworkCustomizer -> curatorFrameworkCustomizer
-						.customize(builder));
-
-		CuratorFramework curator = builder.build();
-		optionalTracerDriverProvider.ifAvailable(tracerDriver -> {
-			if (curator.getZookeeperClient() != null) {
-				curator.getZookeeperClient().setTracerDriver(tracerDriver);
-			}
-		});
-
-		curator.start();
-		if (log.isTraceEnabled()) {
-			log.trace("blocking until connected to zookeeper for "
-					+ properties.getBlockUntilConnectedWait()
-					+ properties.getBlockUntilConnectedUnit());
-		}
-		curator.blockUntilConnected(properties.getBlockUntilConnectedWait(),
-				properties.getBlockUntilConnectedUnit());
-		if (log.isTraceEnabled()) {
-			log.trace("connected to zookeeper");
-		}
-		return curator;
+		return CuratorFactory.curatorFramework(properties, retryPolicy,
+				optionalCuratorFrameworkCustomizerProvider::orderedStream, optionalEnsembleProvider::getIfAvailable,
+				optionalTracerDriverProvider::getIfAvailable);
 	}
 
 	@Bean
 	@ConditionalOnMissingBean
 	public RetryPolicy exponentialBackoffRetry(ZookeeperProperties properties) {
-		return new ExponentialBackoffRetry(properties.getBaseSleepTimeMs(),
-				properties.getMaxRetries(), properties.getMaxSleepMs());
+		return CuratorFactory.retryPolicy(properties);
 	}
 
 }