Skip to content

Commit 6bd3908

Browse files
Ryan Baxterryanjbaxter
Ryan Baxter
andauthored
Support to use classes registered in bootstrap context. Fixes #324 (#325)
Co-authored-by: Ryan Baxter <[email protected]>
1 parent b3a1c2e commit 6bd3908

File tree

2 files changed

+183
-12
lines changed

2 files changed

+183
-12
lines changed

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

+12-12
Original file line numberDiff line numberDiff line change
@@ -168,22 +168,22 @@ public List<ServiceInstance> apply(String serviceId, Binder binder, BindHandler
168168
return Collections.emptyList();
169169
}
170170

171-
ZookeeperProperties properties = binder.bind(ZookeeperProperties.PREFIX, Bindable.of(ZookeeperProperties.class))
172-
.orElse(new ZookeeperProperties());
173-
RetryPolicy retryPolicy = new ExponentialBackoffRetry(properties.getBaseSleepTimeMs(), properties.getMaxRetries(),
174-
properties.getMaxSleepMs());
171+
ZookeeperProperties properties = context.getOrElse(ZookeeperProperties.class, binder.bind(ZookeeperProperties.PREFIX, Bindable.of(ZookeeperProperties.class))
172+
.orElse(new ZookeeperProperties()));
173+
RetryPolicy retryPolicy = context.getOrElse(RetryPolicy.class, new ExponentialBackoffRetry(properties.getBaseSleepTimeMs(), properties.getMaxRetries(),
174+
properties.getMaxSleepMs()));
175175
try {
176-
CuratorFramework curatorFramework = CuratorFactory.curatorFramework(properties, retryPolicy, Stream::of,
177-
() -> null, () -> null);
178-
InstanceSerializer<ZookeeperInstance> serializer = new JsonInstanceSerializer<>(ZookeeperInstance.class);
179-
ZookeeperDiscoveryProperties discoveryProperties = binder.bind(ZookeeperDiscoveryProperties.PREFIX, Bindable
176+
CuratorFramework curatorFramework = context.getOrElse(CuratorFramework.class, CuratorFactory.curatorFramework(properties, retryPolicy, Stream::of,
177+
() -> null, () -> null));
178+
InstanceSerializer<ZookeeperInstance> serializer = context.getOrElse(InstanceSerializer.class, new JsonInstanceSerializer<>(ZookeeperInstance.class));
179+
ZookeeperDiscoveryProperties discoveryProperties = context.getOrElse(ZookeeperDiscoveryProperties.class, binder.bind(ZookeeperDiscoveryProperties.PREFIX, Bindable
180180
.of(ZookeeperDiscoveryProperties.class), bindHandler)
181-
.orElseGet(() -> new ZookeeperDiscoveryProperties(new InetUtils(new InetUtilsProperties())));
182-
DefaultServiceDiscoveryCustomizer customizer = new DefaultServiceDiscoveryCustomizer(curatorFramework, discoveryProperties, serializer);
181+
.orElseGet(() -> new ZookeeperDiscoveryProperties(new InetUtils(new InetUtilsProperties()))));
182+
ServiceDiscoveryCustomizer customizer = context.getOrElse(ServiceDiscoveryCustomizer.class, new DefaultServiceDiscoveryCustomizer(curatorFramework, discoveryProperties, serializer));
183183
ServiceDiscovery<ZookeeperInstance> serviceDiscovery = customizer.customize(ServiceDiscoveryBuilder.builder(ZookeeperInstance.class));
184-
ZookeeperDependencies dependencies = binder.bind(ZookeeperDependencies.PREFIX, Bindable
184+
ZookeeperDependencies dependencies = context.getOrElse(ZookeeperDependencies.class, binder.bind(ZookeeperDependencies.PREFIX, Bindable
185185
.of(ZookeeperDependencies.class), bindHandler)
186-
.orElseGet(ZookeeperDependencies::new);
186+
.orElseGet(ZookeeperDependencies::new));
187187
return new ZookeeperDiscoveryClient(serviceDiscovery, dependencies, discoveryProperties).getInstances(serviceId);
188188
}
189189
catch (Exception e) {
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
/*
2+
* Copyright 2015-2023 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.cloud.zookeeper.discovery.configclient;
18+
19+
import java.util.concurrent.atomic.AtomicReference;
20+
21+
import org.apache.commons.logging.Log;
22+
import org.apache.curator.RetryPolicy;
23+
import org.apache.curator.drivers.TracerDriver;
24+
import org.apache.curator.ensemble.EnsembleProvider;
25+
import org.apache.curator.ensemble.fixed.FixedEnsembleProvider;
26+
import org.apache.curator.framework.CuratorFrameworkFactory;
27+
import org.apache.curator.retry.ExponentialBackoffRetry;
28+
import org.apache.curator.utils.DefaultTracerDriver;
29+
import org.junit.jupiter.api.AfterEach;
30+
import org.junit.jupiter.api.Test;
31+
32+
import org.springframework.boot.BootstrapRegistry;
33+
import org.springframework.boot.BootstrapRegistryInitializer;
34+
import org.springframework.boot.SpringBootConfiguration;
35+
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
36+
import org.springframework.boot.builder.SpringApplicationBuilder;
37+
import org.springframework.boot.context.properties.bind.BindContext;
38+
import org.springframework.boot.context.properties.bind.BindHandler;
39+
import org.springframework.boot.context.properties.bind.Bindable;
40+
import org.springframework.boot.context.properties.bind.Binder;
41+
import org.springframework.boot.context.properties.source.ConfigurationPropertyName;
42+
import org.springframework.cloud.config.client.ConfigServerInstanceProvider;
43+
import org.springframework.cloud.zookeeper.CuratorFrameworkCustomizer;
44+
import org.springframework.cloud.zookeeper.ZookeeperProperties;
45+
import org.springframework.cloud.zookeeper.discovery.ZookeeperDiscoveryClient;
46+
import org.springframework.cloud.zookeeper.test.ZookeeperTestingServer;
47+
import org.springframework.context.ConfigurableApplicationContext;
48+
import org.springframework.core.Ordered;
49+
50+
import static org.assertj.core.api.Assertions.assertThat;
51+
import static org.mockito.Mockito.mock;
52+
53+
/**
54+
* @author Ryan Baxter
55+
*/
56+
public class ZookeeperConfigServerBootstrapperCustomizerTests {
57+
private ConfigurableApplicationContext context;
58+
59+
@AfterEach
60+
public void after() {
61+
if (context != null) {
62+
context.close();
63+
}
64+
}
65+
66+
@Test
67+
public void enabledAddsInstanceProviderFn() {
68+
AtomicReference<ZookeeperDiscoveryClient> bootstrapDiscoveryClient = new AtomicReference<>();
69+
BindHandlerBootstrapper bindHandlerBootstrapper = new BindHandlerBootstrapper();
70+
context = new SpringApplicationBuilder(ZookeeperConfigServerBootstrapperTests.TestConfig.class)
71+
.listeners(new ZookeeperTestingServer())
72+
.properties("--server.port=0", "spring.cloud.config.discovery.enabled=true",
73+
"spring.cloud.zookeeper.discovery.metadata[mymetadataprop]=mymetadataval",
74+
"spring.cloud.service-registry.auto-registration.enabled=false")
75+
.addBootstrapRegistryInitializer(bindHandlerBootstrapper)
76+
.addBootstrapRegistryInitializer(registry -> registry.addCloseListener(event -> {
77+
ConfigServerInstanceProvider.Function providerFn = event.getBootstrapContext()
78+
.get(ConfigServerInstanceProvider.Function.class);
79+
assertThat(providerFn.apply("id", event.getBootstrapContext()
80+
.get(Binder.class), event.getBootstrapContext()
81+
.get(BindHandler.class), mock(Log.class))).as("Should return empty list.")
82+
.isNotNull();
83+
bootstrapDiscoveryClient.set(event.getBootstrapContext().get(ZookeeperDiscoveryClient.class));
84+
CuratorFrameworkCustomizer curatorFrameworkCustomizer = event.getBootstrapContext()
85+
.get(CuratorFrameworkCustomizer.class);
86+
assertThat(curatorFrameworkCustomizer).isInstanceOf(MyCuratorFrameworkCustomizer.class);
87+
RetryPolicy retryPolicy = event.getBootstrapContext().get(RetryPolicy.class);
88+
assertThat(retryPolicy).isInstanceOf(RetryPolicy.class);
89+
EnsembleProvider ensembleProvider = event.getBootstrapContext().get(EnsembleProvider.class);
90+
assertThat(ensembleProvider).isInstanceOf(MyEnsembleProvider.class);
91+
TracerDriver tracerDriver = event.getBootstrapContext().get(TracerDriver.class);
92+
assertThat(tracerDriver).isInstanceOf(MyTracerDriver.class);
93+
})).run();
94+
95+
ZookeeperDiscoveryClient discoveryClient = context.getBean(ZookeeperDiscoveryClient.class);
96+
97+
assertThat(discoveryClient == bootstrapDiscoveryClient.get()).isTrue();
98+
assertThat(bindHandlerBootstrapper.onSuccessCount).isGreaterThan(0);
99+
}
100+
101+
@SpringBootConfiguration
102+
@EnableAutoConfiguration
103+
static class TestConfig {
104+
105+
}
106+
107+
static class BindHandlerBootstrapper implements BootstrapRegistryInitializer, Ordered {
108+
109+
private int onSuccessCount = 0;
110+
111+
private static <T> void registerIfAbsentAndEnabled(
112+
BootstrapRegistry registry, Class<T> type, BootstrapRegistry.InstanceSupplier<T> supplier) {
113+
registry.registerIfAbsent(type, supplier);
114+
}
115+
116+
@Override
117+
public void initialize(BootstrapRegistry registry) {
118+
registry.register(BindHandler.class, context -> new BindHandler() {
119+
@Override
120+
public Object onSuccess(ConfigurationPropertyName name, Bindable<?> target, BindContext context,
121+
Object result) {
122+
onSuccessCount++;
123+
return result;
124+
}
125+
});
126+
127+
registerIfAbsentAndEnabled(registry, RetryPolicy.class, context ->
128+
new MyRetryPolicy(context.get(ZookeeperProperties.class)));
129+
130+
registerIfAbsentAndEnabled(registry, CuratorFrameworkCustomizer.class, context ->
131+
new MyCuratorFrameworkCustomizer());
132+
133+
registerIfAbsentAndEnabled(registry, EnsembleProvider.class, context ->
134+
new MyEnsembleProvider(context.get(ZookeeperProperties.class)));
135+
136+
registerIfAbsentAndEnabled(registry, TracerDriver.class, context ->
137+
new MyTracerDriver());
138+
}
139+
140+
@Override
141+
public int getOrder() {
142+
return Ordered.HIGHEST_PRECEDENCE;
143+
}
144+
}
145+
146+
static class MyRetryPolicy extends ExponentialBackoffRetry {
147+
148+
MyRetryPolicy(ZookeeperProperties properties) {
149+
super(properties.getBaseSleepTimeMs(), properties.getMaxRetries(), properties.getMaxSleepMs());
150+
}
151+
}
152+
153+
static class MyCuratorFrameworkCustomizer implements CuratorFrameworkCustomizer {
154+
@Override
155+
public void customize(CuratorFrameworkFactory.Builder builder) {
156+
157+
}
158+
}
159+
160+
static class MyEnsembleProvider extends FixedEnsembleProvider {
161+
MyEnsembleProvider(ZookeeperProperties properties) {
162+
super(properties.getConnectString());
163+
}
164+
165+
}
166+
167+
static class MyTracerDriver extends DefaultTracerDriver {
168+
169+
}
170+
171+
}

0 commit comments

Comments
 (0)