diff --git a/spring-cloud-zookeeper-core/src/main/java/org/springframework/cloud/zookeeper/ZookeeperAutoConfiguration.java b/spring-cloud-zookeeper-core/src/main/java/org/springframework/cloud/zookeeper/ZookeeperAutoConfiguration.java index 407fa3d5..d81b3f34 100644 --- a/spring-cloud-zookeeper-core/src/main/java/org/springframework/cloud/zookeeper/ZookeeperAutoConfiguration.java +++ b/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 optionalEnsembleProvider, + ObjectProvider 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() diff --git a/spring-cloud-zookeeper-core/src/test/java/org/springframework/cloud/zookeeper/ZookeeperAutoConfigurationTracerDriverTests.java b/spring-cloud-zookeeper-core/src/test/java/org/springframework/cloud/zookeeper/ZookeeperAutoConfigurationTracerDriverTests.java new file mode 100644 index 00000000..b1aa7bec --- /dev/null +++ b/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); + } + } + +}