ZOOKEEPER-4819 Optimizing send4LetterWord with specific clientConfig

apache · Oct 17, 2024 · 4c26a3f · 4c26a3f
1 parent 33c19e3
commit 4c26a3f
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 12 deletions.
diff --git a/zookeeper-server/src/main/java/org/apache/zookeeper/ClientCnxn.java b/zookeeper-server/src/main/java/org/apache/zookeeper/ClientCnxn.java
@@ -1300,11 +1300,10 @@ private void cleanAndNotifyState() {
         private void pingRwServer() throws RWServerFoundException {
             String result = null;
             InetSocketAddress addr = hostProvider.next(0);
-            boolean useSecure = clientConfig.getBoolean(ZKClientConfig.SECURE_CLIENT);
 
             LOG.info("Checking server {} for being r/w. Timeout {}", addr, pingRwTimeout);
             try {
-                result = FourLetterWordMain.send4LetterWord(addr.getHostString(), addr.getPort(), "isro", useSecure, 1000);
+                result = FourLetterWordMain.send4LetterWord(addr.getHostString(), addr.getPort(), "isro", 1000, clientConfig);
             } catch (ConnectException e) {
                 // ignore, this just means server is not up
             } catch (IOException | X509Exception.SSLContextException e) {

diff --git a/zookeeper-server/src/main/java/org/apache/zookeeper/client/FourLetterWordMain.java b/zookeeper-server/src/main/java/org/apache/zookeeper/client/FourLetterWordMain.java
@@ -94,6 +94,33 @@ public static String send4LetterWord(
         return send4LetterWord(host, port, cmd, secure, timeout, null);
     }
 
+    /**
+     * Send the 4letterword
+     * @param host the destination host
+     * @param port the destination port
+     * @param cmd the 4letterword
+     * @param timeout in milliseconds, maximum time to wait while connecting/reading data
+     * @param clientConfig client config
+     * @return server response
+     * @throws SSLContextException
+     * @throws IOException
+     */
+    public static String send4LetterWord(
+            String host,
+            int port,
+            String cmd,
+            int timeout,
+            ZKClientConfig clientConfig) throws SSLContextException, IOException {
+        boolean useSecure = clientConfig.getBoolean(ZKClientConfig.SECURE_CLIENT);
+        SSLContext sslContext = null;
+        if (useSecure) {
+            try (X509Util x509Util = new ClientX509Util()) {
+                sslContext = x509Util.createSSLContext(clientConfig);
+            }
+        }
+        return send4LetterWord(host, port, cmd, useSecure, timeout, sslContext);
+    }
+
     /**
      * Send the 4letterword
      * @param host the destination host

diff --git a/zookeeper-server/src/test/java/org/apache/zookeeper/test/ReadOnlyModeWithSSLTest.java b/zookeeper-server/src/test/java/org/apache/zookeeper/test/ReadOnlyModeWithSSLTest.java
@@ -37,24 +37,28 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
-import org.slf4j.LoggerFactory;
 
 public class ReadOnlyModeWithSSLTest extends ZKTestCase {
 
-    private static final org.slf4j.Logger LOG = LoggerFactory.getLogger(ReadOnlyModeWithSSLTest.class);
     private static int CONNECTION_TIMEOUT = QuorumBase.CONNECTION_TIMEOUT;
     private QuorumUtil qu = new QuorumUtil(1);
     private ClientX509Util clientX509Util;
+    private ZKClientConfig clientConfig;
 
     @BeforeEach
     public void setUp() throws Exception {
-        System.setProperty("readonlymode.enabled", "true");
-        System.setProperty(NettyServerCnxnFactory.PORT_UNIFICATION_KEY, Boolean.TRUE.toString());
         clientX509Util = new ClientX509Util();
         String testDataPath = System.getProperty("test.data.dir", "src/test/resources/data");
+        clientConfig = new ZKClientConfig();
+        clientConfig.setProperty(ZKClientConfig.SECURE_CLIENT, "true");
+        clientConfig.setProperty(ZKClientConfig.ZOOKEEPER_CLIENT_CNXN_SOCKET, "org.apache.zookeeper.ClientCnxnSocketNetty");
+        clientConfig.setProperty(clientX509Util.getSslKeystoreLocationProperty(), testDataPath + "/ssl/testKeyStore.jks");
+        clientConfig.setProperty(clientX509Util.getSslKeystorePasswdProperty(), "testpass");
+        clientConfig.setProperty(clientX509Util.getSslTruststoreLocationProperty(), testDataPath + "/ssl/testTrustStore.jks");
+        clientConfig.setProperty(clientX509Util.getSslTruststorePasswdProperty(), "testpass");
+        System.setProperty("readonlymode.enabled", "true");
+        System.setProperty(NettyServerCnxnFactory.PORT_UNIFICATION_KEY, Boolean.TRUE.toString());
         System.setProperty(ServerCnxnFactory.ZOOKEEPER_SERVER_CNXN_FACTORY, "org.apache.zookeeper.server.NettyServerCnxnFactory");
-        System.setProperty(ZKClientConfig.ZOOKEEPER_CLIENT_CNXN_SOCKET, "org.apache.zookeeper.ClientCnxnSocketNetty");
-        System.setProperty(ZKClientConfig.SECURE_CLIENT, "true");
         System.setProperty(clientX509Util.getSslKeystoreLocationProperty(), testDataPath + "/ssl/testKeyStore.jks");
         System.setProperty(clientX509Util.getSslKeystorePasswdProperty(), "testpass");
         System.setProperty(clientX509Util.getSslTruststoreLocationProperty(), testDataPath + "/ssl/testTrustStore.jks");
@@ -65,8 +69,6 @@ public void setUp() throws Exception {
     public void tearDown() throws Exception {
         System.clearProperty(NettyServerCnxnFactory.PORT_UNIFICATION_KEY);
         System.clearProperty(ServerCnxnFactory.ZOOKEEPER_SERVER_CNXN_FACTORY);
-        System.clearProperty(ZKClientConfig.ZOOKEEPER_CLIENT_CNXN_SOCKET);
-        System.clearProperty(ZKClientConfig.SECURE_CLIENT);
         System.clearProperty(clientX509Util.getSslKeystoreLocationProperty());
         System.clearProperty(clientX509Util.getSslKeystorePasswdProperty());
         System.clearProperty(clientX509Util.getSslKeystorePasswdPathProperty());
@@ -96,7 +98,7 @@ public void testSeekForRwServerWithSSL() throws Exception {
 
             qu.shutdown(2);
             ClientBase.CountdownWatcher watcher = new ClientBase.CountdownWatcher();
-            ZooKeeper zk = new ZooKeeper(qu.getConnString(), CONNECTION_TIMEOUT, watcher, true);
+            ZooKeeper zk = new ZooKeeper(qu.getConnString(), CONNECTION_TIMEOUT, watcher, true, clientConfig);
             watcher.waitForConnected(CONNECTION_TIMEOUT);
 
             // if we don't suspend a peer it will rejoin a quorum
@@ -132,4 +134,4 @@ public void testSeekForRwServerWithSSL() throws Exception {
             assertTrue(found, "Majority server wasn't found while connected to r/o server");
         }
     }
-}
+}