Skip to content

Commit 59faa12

Browse files
author
liwang
committed
ZOOKEEPER-4982: Automatic HostProvider selection based on connection string format
Author: Li Wang <[email protected]>
1 parent 8f9e786 commit 59faa12

File tree

5 files changed

+373
-24
lines changed

5 files changed

+373
-24
lines changed

‎zookeeper-server/src/main/java/org/apache/zookeeper/ZooKeeper.java‎

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,9 @@
4545
import org.apache.zookeeper.Watcher.WatcherType;
4646
import org.apache.zookeeper.client.Chroot;
4747
import org.apache.zookeeper.client.ConnectStringParser;
48+
import org.apache.zookeeper.client.ConnectionType;
4849
import org.apache.zookeeper.client.HostProvider;
50+
import org.apache.zookeeper.client.HostProviderFactory;
4951
import org.apache.zookeeper.client.StaticHostProvider;
5052
import org.apache.zookeeper.client.ZKClientConfig;
5153
import org.apache.zookeeper.client.ZooKeeperBuilder;
@@ -1139,8 +1141,11 @@ public ZooKeeper(ZooKeeperOptions options) throws IOException {
11391141
HostProvider hostProvider;
11401142
if (options.getHostProvider() != null) {
11411143
hostProvider = options.getHostProvider().apply(connectStringParser.getServerAddresses());
1144+
// Validate compatibility when HostProvider is provided
1145+
validateHostProviderCompatibility(connectString, hostProvider);
11421146
} else {
1143-
hostProvider = new StaticHostProvider(connectStringParser.getServerAddresses());
1147+
// Use HostProviderFactory to auto-detect and create a HostProvider based on connect string
1148+
hostProvider = HostProviderFactory.createHostProvider(connectString);
11441149
}
11451150
this.hostProvider = hostProvider;
11461151

@@ -1329,6 +1334,11 @@ public synchronized void close() throws InterruptedException {
13291334
LOG.debug("Ignoring unexpected exception during close", e);
13301335
}
13311336

1337+
// Close the host provider to release any resources
1338+
if (hostProvider != null) {
1339+
hostProvider.close();
1340+
}
1341+
13321342
LOG.info("Session: 0x{} closed", Long.toHexString(getSessionId()));
13331343
}
13341344

@@ -3199,4 +3209,24 @@ public synchronized List<ClientInfo> whoAmI() throws InterruptedException {
31993209
return response.getClientInfo();
32003210
}
32013211

3212+
/**
3213+
* Validates compatibility between connectString and hostProvider.
3214+
*
3215+
* @param connectString the connection string provided by user
3216+
* @param hostProvider the host provider provided by user
3217+
* @throws IllegalArgumentException if incompatible combination is detected
3218+
*/
3219+
private static void validateHostProviderCompatibility(final String connectString, final HostProvider hostProvider) {
3220+
final ConnectionType connectStringType = ConnectStringParser.getConnectionType(connectString);
3221+
final ConnectionType supportedConnectStringType = hostProvider.getSupportedConnectionType();
3222+
3223+
if (connectStringType != supportedConnectStringType) {
3224+
final String hostProviderName = hostProvider.getClass().getSimpleName();
3225+
LOG.error("Connection string type {} is incompatible with host provider type {}: connectString={}, hostProvider={}",
3226+
connectStringType.getName(), supportedConnectStringType.getName(), connectString, hostProviderName);
3227+
throw new IllegalArgumentException(
3228+
String.format("Connection string type %s is incompatible with host provider type %s",
3229+
connectStringType.getName(), hostProviderName));
3230+
}
3231+
}
32023232
}

‎zookeeper-server/src/main/java/org/apache/zookeeper/client/ConnectStringParser.java‎

Lines changed: 86 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import java.util.List;
2525
import org.apache.zookeeper.common.NetUtils;
2626
import org.apache.zookeeper.common.PathUtils;
27+
import org.apache.zookeeper.common.StringUtils;
2728

2829
/**
2930
* A parser for ZooKeeper Client connect strings.
@@ -38,8 +39,9 @@
3839
public final class ConnectStringParser {
3940

4041
private static final int DEFAULT_PORT = 2181;
42+
private static final String DNS_SRV_PREFIX = "dns-srv://";
4143

42-
private final String chrootPath;
44+
private String chrootPath;
4345

4446
private final ArrayList<InetSocketAddress> serverAddresses = new ArrayList<>();
4547

@@ -49,27 +51,71 @@ public final class ConnectStringParser {
4951
* @throws IllegalArgumentException
5052
* for an invalid chroot path.
5153
*/
52-
public ConnectStringParser(String connectString) {
53-
// parse out chroot, if any
54-
int off = connectString.indexOf('/');
55-
if (off >= 0) {
56-
String chrootPath = connectString.substring(off);
57-
// ignore "/" chroot spec, same as null
58-
if (chrootPath.length() == 1) {
59-
this.chrootPath = null;
60-
} else {
61-
PathUtils.validatePath(chrootPath);
62-
this.chrootPath = chrootPath;
63-
}
64-
connectString = connectString.substring(0, off);
54+
public ConnectStringParser(final String connectString) {
55+
if (StringUtils.isBlank(connectString)) {
56+
throw new IllegalArgumentException("Connect string cannot be null or blank");
57+
}
58+
59+
final ConnectionType connectionType = getConnectionType(connectString);
60+
if (connectionType == ConnectionType.DNS_SRV) {
61+
parseDnsSrvFormat(connectString);
6562
} else {
66-
this.chrootPath = null;
63+
parseHostPortFormat(connectString);
64+
}
65+
}
66+
67+
public String getChrootPath() {
68+
return chrootPath;
69+
}
70+
71+
public ArrayList<InetSocketAddress> getServerAddresses() {
72+
return serverAddresses;
73+
}
74+
75+
/**
76+
* Gets the connection type for the given connect string.
77+
*
78+
* @param connectString the connection string to analyze
79+
* @return ConnectionType.DNS_SRV if it's a DNS SRV connect string, ConnectionType.HOST_PORT otherwise
80+
*/
81+
public static ConnectionType getConnectionType(final String connectString) {
82+
if (connectString == null) {
83+
throw new IllegalArgumentException("connectString cannot be null");
6784
}
85+
return connectString.startsWith(DNS_SRV_PREFIX)
86+
? ConnectionType.DNS_SRV : ConnectionType.HOST_PORT;
87+
}
88+
89+
/**
90+
* Parse DNS SRV connection string format: dns-srv://service.domain.com/chroot
91+
* @throws IllegalArgumentException for an invalid chroot path.
92+
*/
93+
private void parseDnsSrvFormat(final String connectString) {
94+
final String dnsName = connectString.substring(DNS_SRV_PREFIX.length());
95+
96+
final String[] parts = extractChrootPath(dnsName);
97+
final String serviceName = parts[0];
98+
99+
chrootPath = parts[1];
100+
// The DNS service name is stored as a placeholder address
101+
// The actual resolution will be handled by DnsSrvHostProvider
102+
serverAddresses.add(InetSocketAddress.createUnresolved(serviceName, DEFAULT_PORT));
103+
}
104+
105+
/**
106+
* Parse host and port by splitting client connectString
107+
* with support for IPv6 literals
108+
* @throws IllegalArgumentException for an invalid chroot path.
109+
*/
110+
private void parseHostPortFormat(String connectString) {
111+
final String[] parts = extractChrootPath(connectString);
112+
final String serversPart = parts[0];
113+
chrootPath = parts[1];
68114

69-
List<String> hostsList = split(connectString, ",");
115+
final List<String> hostsList = split(serversPart, ",");
70116
for (String host : hostsList) {
71117
int port = DEFAULT_PORT;
72-
String[] hostAndPort = NetUtils.getIPV6HostAndPort(host);
118+
final String[] hostAndPort = NetUtils.getIPV6HostAndPort(host);
73119
if (hostAndPort.length != 0) {
74120
host = hostAndPort[0];
75121
if (hostAndPort.length == 2) {
@@ -89,12 +135,29 @@ public ConnectStringParser(String connectString) {
89135
}
90136
}
91137

92-
public String getChrootPath() {
93-
return chrootPath;
94-
}
138+
/**
139+
* Extract chroot path from a connection string.
140+
*
141+
* @param connectionString the connection string that may contain a chroot path
142+
* @return array where [0] is the server part (before chroot) and [1] is the chroot path (or null)
143+
* @throws IllegalArgumentException for an invalid chroot path
144+
*/
145+
private String[] extractChrootPath(final String connectionString) {
146+
String serverPart = connectionString;
147+
String chrootPath = null;
95148

96-
public ArrayList<InetSocketAddress> getServerAddresses() {
97-
return serverAddresses;
149+
// parse out chroot, if any
150+
final int chrootIndex = connectionString.indexOf('/');
151+
if (chrootIndex >= 0) {
152+
chrootPath = connectionString.substring(chrootIndex);
153+
// ignore "/" chroot spec, same as null
154+
if (chrootPath.length() == 1) {
155+
chrootPath = null;
156+
} else {
157+
PathUtils.validatePath(chrootPath);
158+
}
159+
serverPart = connectionString.substring(0, chrootIndex);
160+
}
161+
return new String[]{serverPart, chrootPath};
98162
}
99-
100163
}

‎zookeeper-server/src/main/java/org/apache/zookeeper/client/HostProviderFactory.java‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
19+
package org.apache.zookeeper.client;
20+
21+
import org.apache.yetus.audience.InterfaceAudience;
22+
import org.apache.zookeeper.common.StringUtils;
23+
import org.slf4j.Logger;
24+
import org.slf4j.LoggerFactory;
25+
26+
/**
27+
* Factory for creating appropriate HostProvider instances based on connection string format.
28+
* This factory enables zero-code-change migration by automatically detecting the connection
29+
* string format and creating the appropriate HostProvider implementation.
30+
*
31+
* Supported formats:
32+
* - Host:Port: "host1:port1,host2:port2,host3:port3" (StaticHostProvider)
33+
* - DNS SRV: "dns-srv://service.domain.com" (DnsSrvHostProvider)
34+
* - Future formats can be easily added by extending the factory
35+
*/
36+
@InterfaceAudience.Public
37+
public class HostProviderFactory {
38+
39+
private static final Logger LOG = LoggerFactory.getLogger(HostProviderFactory.class);
40+
41+
/**
42+
* Creates a HostProvider based on the connection string format.
43+
*
44+
* @param connectString the connection string (host:port or DNS SRV format)
45+
* @return appropriate HostProvider implementation
46+
* @throws IllegalArgumentException if the connection string format is not supported
47+
*/
48+
public static HostProvider createHostProvider(final String connectString) {
49+
if (StringUtils.isBlank(connectString)) {
50+
throw new IllegalArgumentException("Connection string cannot be null or empty");
51+
}
52+
53+
final String trimmedConnectString = connectString.trim();
54+
try {
55+
final ConnectionType connectionType = ConnectStringParser.getConnectionType(trimmedConnectString);
56+
if (connectionType == ConnectionType.DNS_SRV) {
57+
LOG.info("Detected DNS SRV connection string format: {}", trimmedConnectString);
58+
return createDnsSrvHostProvider(trimmedConnectString);
59+
}
60+
final ConnectStringParser parser = new ConnectStringParser(trimmedConnectString);
61+
return new StaticHostProvider(parser.getServerAddresses());
62+
} catch (final Exception e) {
63+
throw new IllegalArgumentException("Unsupported connection string format: " + trimmedConnectString, e);
64+
}
65+
}
66+
67+
private static DnsSrvHostProvider createDnsSrvHostProvider(final String connectString) {
68+
final ConnectStringParser parser = new ConnectStringParser(connectString);
69+
70+
if (parser.getServerAddresses().isEmpty()) {
71+
throw new IllegalArgumentException("No DNS service name found in connect string: " + connectString);
72+
}
73+
74+
final String dnsServiceName = parser.getServerAddresses().get(0).getHostString();
75+
76+
return new DnsSrvHostProvider(dnsServiceName);
77+
}
78+
}

‎zookeeper-server/src/test/java/org/apache/zookeeper/CustomHostProviderTest.java‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,14 +18,24 @@
1818

1919
package org.apache.zookeeper;
2020

21+
import static org.junit.jupiter.api.Assertions.assertThrows;
2122
import static org.junit.jupiter.api.Assertions.assertTrue;
23+
import static org.mockito.Mockito.mock;
24+
import static org.mockito.Mockito.when;
2225
import java.io.IOException;
26+
import java.net.InetAddress;
2327
import java.net.InetSocketAddress;
2428
import java.util.Collection;
29+
import java.util.Collections;
30+
import java.util.List;
2531
import java.util.concurrent.atomic.AtomicInteger;
32+
import org.apache.zookeeper.client.DnsSrvHostProvider;
2633
import org.apache.zookeeper.client.HostProvider;
34+
import org.apache.zookeeper.client.StaticHostProvider;
2735
import org.apache.zookeeper.test.ClientBase;
2836
import org.junit.jupiter.api.Test;
37+
import org.xbill.DNS.Name;
38+
import org.xbill.DNS.SRVRecord;
2939

3040
public class CustomHostProviderTest extends ZKTestCase {
3141

@@ -83,4 +93,82 @@ public void testZooKeeperWithCustomHostProvider() throws IOException, Interrupte
8393
assertTrue(counter.get() == expectedCounter);
8494
}
8595

96+
@Test
97+
public void testConfigurationMismatchValidation() throws Exception {
98+
final String hostPortConnectString = "zk1:2181,zk2:2181/myapp";
99+
try (final DnsSrvHostProvider dnsSrvHostProvider = createTestDnsSrvHostProvider()) {
100+
assertThrows(IllegalArgumentException.class, () -> new ZooKeeper(
101+
hostPortConnectString,
102+
ClientBase.CONNECTION_TIMEOUT,
103+
DummyWatcher.INSTANCE,
104+
false,
105+
dnsSrvHostProvider));
106+
}
107+
108+
final String dnsConnectString = "dns-srv://zookeeper.example.com/myapp";
109+
final StaticHostProvider staticHostProvider = createTestStaticHostProvider();
110+
assertThrows(IllegalArgumentException.class, () -> new ZooKeeper(
111+
dnsConnectString,
112+
ClientBase.CONNECTION_TIMEOUT,
113+
DummyWatcher.INSTANCE,
114+
false,
115+
staticHostProvider));
116+
}
117+
118+
@Test
119+
public void testValidConfigurations() throws Exception {
120+
// host port connect string with StaticHostProvider
121+
final StaticHostProvider staticHostProvider = createTestStaticHostProvider();
122+
final ZooKeeper zkTraditional = new ZooKeeper(
123+
"localhost:2181,localhost:2182/myapp",
124+
ClientBase.CONNECTION_TIMEOUT,
125+
DummyWatcher.INSTANCE,
126+
false,
127+
staticHostProvider);
128+
129+
zkTraditional.close();
130+
131+
// DNS SRV connect string with DnsSrvHostProvider
132+
try (final DnsSrvHostProvider dnsSrvProvider = createTestDnsSrvHostProvider()) {
133+
final ZooKeeper zkDnsSrv = new ZooKeeper(
134+
"dns-srv://zookeeper.example.com/myapp",
135+
ClientBase.CONNECTION_TIMEOUT,
136+
DummyWatcher.INSTANCE,
137+
false,
138+
dnsSrvProvider);
139+
140+
zkDnsSrv.close();
141+
}
142+
}
143+
144+
private SRVRecord[] createMockSrvRecords(String dnsName) {
145+
try {
146+
Name targetName1 = Name.fromString("server1.example.com.");
147+
Name targetName2 = Name.fromString("server2.example.com.");
148+
149+
Name serviceName = Name.fromString(dnsName + ".");
150+
151+
return new SRVRecord[]{
152+
new SRVRecord(serviceName, 1, 300, 1, 1, 2181, targetName1),
153+
new SRVRecord(serviceName, 1, 300, 1, 1, 2181, targetName2)
154+
};
155+
} catch (Exception e) {
156+
throw new RuntimeException("Failed to create mock SRV records", e);
157+
}
158+
}
159+
160+
private DnsSrvHostProvider createTestDnsSrvHostProvider() throws Exception {
161+
final DnsSrvHostProvider.DnsResolver mockDnsResolver = mock(DnsSrvHostProvider.DnsResolver.class);
162+
final SRVRecord[] mockSrvRecords = createMockSrvRecords("zookeeper._tcp.example.com");
163+
when(mockDnsResolver.lookupSrvRecords("zookeeper._tcp.example.com.")).thenReturn(mockSrvRecords);
164+
165+
return new DnsSrvHostProvider("zookeeper._tcp.example.com", 12345L, mockDnsResolver);
166+
}
167+
168+
private StaticHostProvider createTestStaticHostProvider() throws Exception {
169+
final List<InetSocketAddress> initialServers = Collections.singletonList(new InetSocketAddress(
170+
InetAddress.getByAddress(new byte[]{10, 10, 10, 1}), 2181)
171+
);
172+
return new StaticHostProvider(initialServers);
173+
}
86174
}

0 commit comments

Comments
 (0)