Merge branch '2.2.x'

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

@@ -19,12 +19,13 @@ package org.springframework.cloud.zookeeper;
 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.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.context.annotation.Bean;
@@ -44,9 +45,6 @@ public class ZookeeperAutoConfiguration {
 
 	private static final Log log = LogFactory.getLog(ZookeeperAutoConfiguration.class);
 
-	@Autowired(required = false)
-	private EnsembleProvider ensembleProvider;
-
 	@Bean
 	@ConditionalOnMissingBean
 	public ZookeeperProperties zookeeperProperties() {
@@ -55,16 +53,26 @@ public class ZookeeperAutoConfiguration {
 
 	@Bean(destroyMethod = "close")
 	@ConditionalOnMissingBean
-	public CuratorFramework curatorFramework(RetryPolicy retryPolicy,
-			ZookeeperProperties properties) throws Exception {
+	public CuratorFramework curatorFramework(
+			RetryPolicy retryPolicy,
+			ZookeeperProperties properties,
+			ObjectProvider<EnsembleProvider> optionalEnsembleProvider,
+			ObjectProvider<TracerDriver> optionalTracerDriverProvider)
+	throws Exception {
 		CuratorFrameworkFactory.Builder builder = CuratorFrameworkFactory.builder();
-		if (this.ensembleProvider != null) {
-			builder.ensembleProvider(this.ensembleProvider);
+		EnsembleProvider ensembleProvider = optionalEnsembleProvider.getIfAvailable();
+		if (ensembleProvider != null) {
+			builder.ensembleProvider(ensembleProvider);
 		}
 		else {
 			builder.connectString(properties.getConnectString());
 		}
 		CuratorFramework curator = builder.retryPolicy(retryPolicy).build();
+		TracerDriver tracerDriver = optionalTracerDriverProvider.getIfAvailable();
+		if (curator.getZookeeperClient() != null && tracerDriver != null) {
+			curator.getZookeeperClient().setTracerDriver(tracerDriver);
+		}
+
 		curator.start();
 		log.trace("blocking until connected to zookeeper for "
 				+ properties.getBlockUntilConnectedWait()

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

@@ -0,0 +1,76 @@
+/*
+ * 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;
+
+import org.apache.curator.drivers.TracerDriver;
+import org.apache.curator.framework.CuratorFramework;
+import org.apache.curator.test.TestingServer;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.context.annotation.Bean;
+import org.springframework.test.context.ContextConfiguration;
+import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+
+/**
+ * @author Bernardo Gomez Palacio
+ */
+@RunWith(SpringJUnit4ClassRunner.class)
+@ContextConfiguration(classes = {
+		ZookeeperAutoConfigurationTracerDriverTests.TestConfig.class,
+		ZookeeperAutoConfiguration.class })
+public class ZookeeperAutoConfigurationTracerDriverTests {
+
+	@Autowired(required = false)
+	TracerDriver tracerDriver;
+
+	@Autowired(required = false)
+	CuratorFramework curator;
+
+	@Autowired
+	TestingServer testingServer;
+
+	@Test
+	public void should_successfully_inject_Curators_TracerDriver() {
+		assertThat(curator.getZookeeperClient().getTracerDriver()).isEqualTo(tracerDriver);
+	}
+
+	static class TestConfig {
+
+		@Bean
+		ZookeeperProperties zookeeperProperties(TestingServer testingServer) {
+			ZookeeperProperties properties = new ZookeeperProperties();
+			properties.setConnectString(testingServer.getConnectString());
+			return properties;
+		}
+
+		@Bean(destroyMethod = "close")
+		TestingServer testingServer() throws Exception {
+			return new TestingServer();
+		}
+
+		@Bean
+		TracerDriver mockedTracerDriver() {
+			return mock(TracerDriver.class);
+		}
+	}
+
+}