Skip to content
Original file line number Diff line number Diff line change
Expand Up @@ -18,15 +18,19 @@
*/
package org.apache.tez.client.registry.zookeeper;

import java.util.Optional;
import java.util.concurrent.TimeUnit;

import org.apache.commons.lang3.StringUtils;
import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.Path;
import org.apache.tez.dag.api.TezConfiguration;
import org.apache.zookeeper.client.ZKClientConfig;
import org.apache.zookeeper.common.ClientX509Util;

import com.google.common.base.Preconditions;
import com.google.common.base.Strings;
Expand All @@ -49,6 +53,11 @@ public class ZkConfig {
private final int curatorMaxRetries;
private final int sessionTimeoutMs;
private final int connectionTimeoutMs;
private final String sslEnabled;
private final String sslKeystoreLocation;
private final String sslKeystorePassword;
private final String sslTruststoreLocation;
private final String sslTruststorePassword;

public ZkConfig(Configuration conf) {
zkQuorum = conf.get(TezConfiguration.TEZ_AM_ZOOKEEPER_QUORUM);
Expand Down Expand Up @@ -84,6 +93,16 @@ public ZkConfig(Configuration conf) {
TezConfiguration.TEZ_AM_CURATOR_SESSION_TIMEOUT_DEFAULT, TimeUnit.MILLISECONDS));
connectionTimeoutMs = Math.toIntExact(conf.getTimeDuration(TezConfiguration.TEZ_AM_CURATOR_CONNECTION_TIMEOUT,
TezConfiguration.TEZ_AM_CURATOR_CONNECTION_TIMEOUT_DEFAULT, TimeUnit.MILLISECONDS));
sslEnabled = conf.get(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_ENABLE);
Preconditions.checkArgument(
isValidSslEnabledValue(sslEnabled),
"If the optional %s setting is set, then the value should be a boolean value instead of '%s'",
TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_ENABLE,
sslEnabled);
sslKeystoreLocation = conf.get(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_KEYSTORE_LOCATION);
sslKeystorePassword = conf.get(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_KEYSTORE_PASSWORD);
sslTruststoreLocation = conf.get(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_TRUSTSTORE_LOCATION);
sslTruststorePassword = conf.get(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_TRUSTSTORE_PASSWORD);
}

public String getZkQuorum() {
Expand All @@ -110,17 +129,82 @@ public int getConnectionTimeoutMs() {
return connectionTimeoutMs;
}

public String getZookeeperTrustStorePassword() {
return sslTruststorePassword;
}

public String getZookeeperTrustStoreLocation() {
return sslTruststoreLocation;
}

public String getZookeeperKeyStorePassword() {
return sslKeystorePassword;
}

public String getZookeeperKeyStoreLocation() {
return sslKeystoreLocation;
}

/**
* Returns whether the zookeeper connection will be secure or insecure.
* @return An Optional containing the boolean value that indicates whether zookeeper client
* uses a secure zookeeper connection. An empty Optional indicates that it is not specified,
* and in this case the default settings of zookeeper are used, which can be controlled by
* specific JVM properties.
* @see TezConfiguration#TEZ_AM_ZOOKEEPER_SSL_ENABLE
*/
public Optional<Boolean> isSslEnabled() {
if (this.sslEnabled == null || this.sslEnabled.isEmpty()) {
return Optional.empty();
}
return Optional.of(Boolean.parseBoolean(sslEnabled));
}

public RetryPolicy getRetryPolicy() {
return new ExponentialBackoffRetry(getCuratorBackoffSleepMs(), getCuratorMaxRetries());
}

public CuratorFramework createCuratorFramework() {
return CuratorFrameworkFactory.newClient(
getZkQuorum(),
getSessionTimeoutMs(),
getConnectionTimeoutMs(),
getRetryPolicy()
);
if (!isSslEnabled().isPresent()) {
return CuratorFrameworkFactory.newClient(
getZkQuorum(),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not use the built-in .zkClientConfig() method provided by Curator 5.x? It provides native support for managing these SSL properties. SSLZookeeperFactory.java can we entirely removed and just configure it inline inside ZkConfig.java like this:

ZKClientConfig zkClientConfig = new ZKClientConfig();
zkClientConfig.setProperty(ZKClientConfig.SECURE_CLIENT, "true");
.....
.....

return CuratorFrameworkFactory.builder()
    .connectString(getZkQuorum())
    // ... other settings ...
    .zkClientConfig(zkClientConfig) // Built-in Curator support
    .build();

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense. I've added a commit that removed SSLZookeeperFactory.java and uses the build-in curator config.

getSessionTimeoutMs(),
getConnectionTimeoutMs(),
getRetryPolicy()
);
}

ZKClientConfig zkClientConfig = new ZKClientConfig();
zkClientConfig.setProperty(ZKClientConfig.SECURE_CLIENT, Boolean.toString(isSslEnabled().get()));
zkClientConfig.setProperty(ZKClientConfig.ZOOKEEPER_CLIENT_CNXN_SOCKET,
"org.apache.zookeeper.ClientCnxnSocketNetty");
if (isSslEnabled().get()) {
ClientX509Util x509Util = new ClientX509Util();
if (StringUtils.isEmpty(getZookeeperKeyStoreLocation())) {
LOG.warn("Missing keystoreLocation parameter");
}
if (StringUtils.isEmpty(getZookeeperTrustStoreLocation())) {
LOG.warn("Missing trustStoreLocation parameter");
}
zkClientConfig.setProperty(x509Util.getSslKeystoreLocationProperty(), getZookeeperKeyStoreLocation());

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The above 2 if statements feel wrong. We are LOGGING the warning that parameters are missing and then continuing the use them in 2nd arg in setProperty. that can lead to excpetions like NPE or IllegalArgs.

Instead of lgging we should throw the exception there itself. Maybe IllegalArgException.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we throw an exception here, these parameters will become mandatory instead of optional.
The keystore is only needed if zookeeper is configured for mTLS authentication, so I think it should remain optional. The truststore, on the other hand, is also optional, because if the system truststore trusts the zookeeper TLS cert, then this configuration would be redundant.

How about something like this?

    if (isSslEnabled().get()) {
      try (ClientX509Util x509Util = new ClientX509Util()) {
        if (StringUtils.isNotEmpty(getZookeeperKeyStoreLocation())) {
          zkClientConfig.setProperty(x509Util.getSslKeystoreLocationProperty(), getZookeeperKeyStoreLocation());
          zkClientConfig.setProperty(x509Util.getSslKeystorePasswdProperty(), getZookeeperKeyStorePassword());
        } else {
          LOG.info("No keystore location configured, using ZooKeeper client defaults");
        }
        if (StringUtils.isNotEmpty(getZookeeperTrustStoreLocation())) {
          zkClientConfig.setProperty(x509Util.getSslTruststoreLocationProperty(), getZookeeperTrustStoreLocation());
          zkClientConfig.setProperty(x509Util.getSslTruststorePasswdProperty(), getZookeeperTrustStorePassword());
        } else {
          LOG.info("No truststore location configured, using ZooKeeper client defaults");
        }
      }
    }

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

these parameters will become mandatory instead of optional

thats correct, I missed that part. What you are proposing looks good, wondering if we need to put isNotEmpty check on password as well. Making it a too much if nested code :-) . As I'm not a committer to this project, I would request @abstractdog , for suggestions here. Based on that we can push again.

rest all changes are good.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Your above solution, just small refactor

From e6ca2489dc5179e0b0cc9e3bc90385e78c31af62 Mon Sep 17 00:00:00 2001
From: Raghav Aggarwal <raghavaggarwal03.ra@gmail.com>
Date: Wed, 12 Aug 2026 23:56:47 +0530
Subject: [PATCH] TEZ-4749: Refactor

---
 .../client/registry/zookeeper/ZkConfig.java   | 31 ++++++++++++-------
 tez-tests/pom.xml                             |  1 -
 2 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/tez-api/src/main/java/org/apache/tez/client/registry/zookeeper/ZkConfig.java b/tez-api/src/main/java/org/apache/tez/client/registry/zookeeper/ZkConfig.java
index 7c3807c7c..cb8380e51 100644
--- a/tez-api/src/main/java/org/apache/tez/client/registry/zookeeper/ZkConfig.java
+++ b/tez-api/src/main/java/org/apache/tez/client/registry/zookeeper/ZkConfig.java
@@ -165,7 +165,7 @@ public class ZkConfig {
   }
 
   public CuratorFramework createCuratorFramework() {
-    if (!isSslEnabled().isPresent()) {
+    if (isSslEnabled().isEmpty()) {
       return CuratorFrameworkFactory.newClient(
               getZkQuorum(),
               getSessionTimeoutMs(),
@@ -179,17 +179,13 @@ public class ZkConfig {
     zkClientConfig.setProperty(ZKClientConfig.ZOOKEEPER_CLIENT_CNXN_SOCKET,
             "org.apache.zookeeper.ClientCnxnSocketNetty");
     if (isSslEnabled().get()) {
-      ClientX509Util x509Util = new ClientX509Util();
-      if (StringUtils.isEmpty(getZookeeperKeyStoreLocation())) {
-        LOG.warn("Missing keystoreLocation parameter");
+      try (ClientX509Util x509Util = new ClientX509Util()) {
+        setStoreConfig(zkClientConfig, x509Util.getSslKeystoreLocationProperty(), getZookeeperKeyStoreLocation(),
+            x509Util.getSslKeystorePasswdProperty(), getZookeeperKeyStorePassword(), "keystore");
+
+        setStoreConfig(zkClientConfig, x509Util.getSslTruststoreLocationProperty(), getZookeeperTrustStoreLocation(),
+            x509Util.getSslTruststorePasswdProperty(), getZookeeperTrustStorePassword(), "truststore");
       }
-      if (StringUtils.isEmpty(getZookeeperTrustStoreLocation())) {
-        LOG.warn("Missing trustStoreLocation parameter");
-      }
-      zkClientConfig.setProperty(x509Util.getSslKeystoreLocationProperty(), getZookeeperKeyStoreLocation());
-      zkClientConfig.setProperty(x509Util.getSslKeystorePasswdProperty(), getZookeeperKeyStorePassword());
-      zkClientConfig.setProperty(x509Util.getSslTruststoreLocationProperty(), getZookeeperTrustStoreLocation());
-      zkClientConfig.setProperty(x509Util.getSslTruststorePasswdProperty(), getZookeeperTrustStorePassword());
     }
 
     return CuratorFrameworkFactory.builder()
@@ -201,6 +197,19 @@ public class ZkConfig {
             .build();
   }
 
+  private void setStoreConfig(ZKClientConfig config, String locationProp, String locationVal, String passwordProp,
+                              String passwordVal, String storeName) {
+    if (StringUtils.isEmpty(locationVal)) {
+      LOG.info("No {} location configured, using ZooKeeper client defaults", storeName);
+      return;
+    }
+
+    config.setProperty(locationProp, locationVal);
+    if (StringUtils.isNotEmpty(passwordVal)) {
+      config.setProperty(passwordProp, passwordVal);
+    }
+  }
+
   private boolean isValidSslEnabledValue(String value) {
     return value == null || value.isEmpty()
         || value.trim().equalsIgnoreCase("true")
diff --git a/tez-tests/pom.xml b/tez-tests/pom.xml
index 3b76ed090..8920c17e6 100644
--- a/tez-tests/pom.xml
+++ b/tez-tests/pom.xml
@@ -136,7 +136,6 @@
     <dependency>
       <groupId>org.apache.curator</groupId>
       <artifactId>curator-test</artifactId>
-      <version>${curator.version}</version>
       <scope>test</scope>
     </dependency>
   </dependencies>
-- 
2.55.0

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I added these. Thanks!

zkClientConfig.setProperty(x509Util.getSslKeystorePasswdProperty(), getZookeeperKeyStorePassword());
zkClientConfig.setProperty(x509Util.getSslTruststoreLocationProperty(), getZookeeperTrustStoreLocation());
zkClientConfig.setProperty(x509Util.getSslTruststorePasswdProperty(), getZookeeperTrustStorePassword());
}

return CuratorFrameworkFactory.builder()
.connectString(getZkQuorum())
.sessionTimeoutMs(getSessionTimeoutMs())
.connectionTimeoutMs(getConnectionTimeoutMs())
.retryPolicy(getRetryPolicy())
.zkClientConfig(zkClientConfig)
.build();
}

private boolean isValidSslEnabledValue(String value) {
return value == null || value.isEmpty()
|| value.trim().equalsIgnoreCase("true")
|| value.trim().equalsIgnoreCase("false");
}

/**
Expand Down
57 changes: 57 additions & 0 deletions tez-api/src/main/java/org/apache/tez/dag/api/TezConfiguration.java
Original file line number Diff line number Diff line change
Expand Up @@ -2268,6 +2268,63 @@ static Set<String> getPropertySet() {
public static final String TEZ_SHARED_EXECUTOR_MAX_THREADS = "tez.shared-executor.max-threads";
public static final int TEZ_SHARED_EXECUTOR_MAX_THREADS_DEFAULT = -1;

/**
* Optional boolean value represented by string type. A value of "true" enables secure
* Zookeeper connection in ZkAMRegistry and ZkAMRegistryClient classes, while a value
* of "false" disables secure Zookeeper connection.
* If not specified or empty string, then zookeeper enables/disables the secure Zookeeper
* connection based on JVM properties.
* Default: Empty
*/
@ConfigurationScope(Scope.AM)
@ConfigurationProperty
public static final String TEZ_AM_ZOOKEEPER_SSL_ENABLE = TEZ_AM_PREFIX
+ "zookeeper.ssl.client.enable";

/**
* String value
* An optional setting that specifies the path to the keystore used for the secure
* zookeeper connection.
* Default: Empty
*/
@ConfigurationScope(Scope.AM)
@ConfigurationProperty
public static final String TEZ_AM_ZOOKEEPER_SSL_KEYSTORE_LOCATION = TEZ_AM_PREFIX
+ "zookeeper.ssl.keystore.location";

/**
* String value
* An optional setting that specifies the password of the keystore used for the secure
* zookeeper connection.
* Default: Empty
*/
@ConfigurationScope(Scope.AM)
@ConfigurationProperty
public static final String TEZ_AM_ZOOKEEPER_SSL_KEYSTORE_PASSWORD = TEZ_AM_PREFIX
+ "zookeeper.ssl.keystore.password";

/**
* String value
* An optional setting that specifies the path to the truststore used for the secure
* zookeeper connection.
* Default: Empty
*/
@ConfigurationScope(Scope.AM)
@ConfigurationProperty
public static final String TEZ_AM_ZOOKEEPER_SSL_TRUSTSTORE_LOCATION = TEZ_AM_PREFIX
+ "zookeeper.ssl.truststore.location";

/**
* String value
* An optional setting that specifies the password of the truststore used for the secure
* zookeeper connection.
* Default: Empty
*/
@ConfigurationScope(Scope.AM)
@ConfigurationProperty
public static final String TEZ_AM_ZOOKEEPER_SSL_TRUSTSTORE_PASSWORD = TEZ_AM_PREFIX
+ "zookeeper.ssl.truststore.password";

/**
* Acquire all FileSystems info. e.g., all namenodes info of HDFS federation cluster.
*/
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,13 @@
package org.apache.tez.client.registry.zookeeper;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

import java.util.Optional;
import java.util.concurrent.TimeUnit;

import org.apache.curator.RetryPolicy;
Expand Down Expand Up @@ -231,4 +235,73 @@ public void testDefaultNamespace() {
assertEquals("/tez-external-sessions" + TezConfiguration.TEZ_AM_REGISTRY_NAMESPACE_DEFAULT,
zkConfig.getZkNamespace());
}

@Test
public void testZkConfigTezAmZookeeperSslEnableNotSpecified() {
Configuration conf = new Configuration();
conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_QUORUM, "dummyZkQuorum");
ZkConfig zkConf = new ZkConfig(conf);

assertEquals(Optional.empty(), zkConf.isSslEnabled());
assertNull(zkConf.getZookeeperKeyStoreLocation());
assertNull(zkConf.getZookeeperKeyStorePassword());
assertNull(zkConf.getZookeeperTrustStoreLocation());
assertNull(zkConf.getZookeeperTrustStorePassword());
}

@Test
public void testZkConfigTezAmZookeeperSslEnableEmpty() {
Configuration conf = new Configuration();
conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_QUORUM, "dummyZkQuorum");
conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_ENABLE, ""); // empty means not set
ZkConfig zkConf = new ZkConfig(conf);

assertEquals(Optional.empty(), zkConf.isSslEnabled());
assertNull(zkConf.getZookeeperKeyStoreLocation());
assertNull(zkConf.getZookeeperKeyStorePassword());
assertNull(zkConf.getZookeeperTrustStoreLocation());
assertNull(zkConf.getZookeeperTrustStorePassword());
}

@Test
public void testZkConfigSslEnabled() {
Configuration conf = new Configuration();
conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_QUORUM, "dummyZkQuorum");
conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_ENABLE, "true");
conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_KEYSTORE_LOCATION, "/keystore.jks");
conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_KEYSTORE_PASSWORD, "secret");
conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_TRUSTSTORE_LOCATION, "/truststore.jks");
conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_TRUSTSTORE_PASSWORD, "changeit");
ZkConfig zkConf = new ZkConfig(conf);

assertTrue(zkConf.isSslEnabled().isPresent());
assertTrue(zkConf.isSslEnabled().get());
assertEquals(zkConf.getZookeeperKeyStoreLocation(), "/keystore.jks");
assertEquals(zkConf.getZookeeperKeyStorePassword(), "secret");
assertEquals(zkConf.getZookeeperTrustStoreLocation(), "/truststore.jks");
assertEquals(zkConf.getZookeeperTrustStorePassword(), "changeit");
}

@Test
public void testZkConfigSslDisabled() {
Configuration conf = new Configuration();
conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_QUORUM, "dummyZkQuorum");
conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_ENABLE, "False");

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

"False" looks strange, even if it's by design, I would keep using "false", and create a separate test case to show valid values

ZkConfig zkConf = new ZkConfig(conf);

assertTrue(zkConf.isSslEnabled().isPresent());
assertFalse(zkConf.isSslEnabled().get());
assertNull(zkConf.getZookeeperKeyStoreLocation());
assertNull(zkConf.getZookeeperKeyStorePassword());
assertNull(zkConf.getZookeeperTrustStoreLocation());
assertNull(zkConf.getZookeeperTrustStorePassword());
}

@Test
public void testZkConfigAmZookeeperSslEnableInvalid() {
Configuration conf = new Configuration();
conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_QUORUM, "dummyZkQuorum");
conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_ENABLE, "invalidValue");
assertThrows(IllegalArgumentException.class, () -> new ZkConfig(conf));
}
}
6 changes: 6 additions & 0 deletions tez-tests/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,12 @@
<groupId>org.junit.jupiter</groupId>
<artifactId>junit-jupiter</artifactId>
</dependency>
<dependency>
<groupId>org.apache.curator</groupId>
<artifactId>curator-test</artifactId>
<version>${curator.version}</version>

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: No need for version tag here as dependencyManagement handles that.

<scope>test</scope>
</dependency>
</dependencies>

<build>
Expand Down
Loading