Merge pull request #2 from xuluan/am-onyarn

[CARMEL-3186] Run spark thrift server in yarn cluster mode for spark3

Author: mingmwang

core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala

@@ -281,8 +281,6 @@ private[spark] class SparkSubmit extends Logging {
         error("Cluster deploy mode is not applicable to Spark shells.")
       case (_, CLUSTER) if isSqlShell(args.mainClass) =>
         error("Cluster deploy mode is not applicable to Spark SQL shell.")
-      case (_, CLUSTER) if isThriftServer(args.mainClass) =>
-        error("Cluster deploy mode is not applicable to Spark Thrift server.")
       case _ =>
     }
 

sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2.scala

@@ -115,6 +115,10 @@ object HiveThriftServer2 extends Logging {
         logError("SparkContext has stopped even if HiveServer2 has started, so exit")
         System.exit(-1)
       }
+
+      while (server.started.get()) {
+        Thread.sleep(1000)
+      }
     } catch {
       case e: Exception =>
         logError("Error starting HiveThriftServer2", e)

sql/hive-thriftserver/v2.3/src/main/java/org/apache/hive/service/server/HiveServer2.java

@@ -18,7 +18,12 @@
 
 package org.apache.hive.service.server;
 
+import static org.apache.hadoop.hive.ql.util.ZooKeeperHiveHelper.ZOOKEEPER_PATH_SEPARATOR;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
 import java.util.Properties;
+import java.util.concurrent.TimeUnit;
 
 import scala.runtime.AbstractFunction0;
 import scala.runtime.BoxedUnit;
@@ -29,14 +34,26 @@
 import org.apache.commons.cli.OptionBuilder;
 import org.apache.commons.cli.Options;
 import org.apache.commons.cli.ParseException;
+import org.apache.curator.framework.CuratorFramework;
+import org.apache.curator.framework.CuratorFrameworkFactory;
+import org.apache.curator.framework.recipes.nodes.PersistentEphemeralNode;
+import org.apache.curator.framework.recipes.nodes.PersistentEphemeralNode.Mode;
+import org.apache.curator.retry.ExponentialBackoffRetry;
 import org.apache.hadoop.hive.common.JvmPauseMonitor;
 import org.apache.hadoop.hive.conf.HiveConf;
+import org.apache.hadoop.hive.conf.HiveConf.ConfVars;
+import org.apache.hadoop.hive.ql.util.ZooKeeperHiveHelper;
+import org.apache.hadoop.hive.shims.Utils;
+import org.apache.hadoop.security.UserGroupInformation;
 import org.apache.hive.common.util.HiveStringUtils;
 import org.apache.hive.service.CompositeService;
+import org.apache.hive.service.ServiceException;
 import org.apache.hive.service.cli.CLIService;
 import org.apache.hive.service.cli.thrift.ThriftBinaryCLIService;
 import org.apache.hive.service.cli.thrift.ThriftCLIService;
 import org.apache.hive.service.cli.thrift.ThriftHttpCLIService;
+import org.apache.zookeeper.CreateMode;
+import org.apache.zookeeper.KeeperException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -51,6 +68,8 @@ public class HiveServer2 extends CompositeService {
 
   private CLIService cliService;
   private ThriftCLIService thriftCLIService;
+  private PersistentEphemeralNode znode;
+  private CuratorFramework zooKeeperClient;
 
   public HiveServer2() {
     super(HiveServer2.class.getSimpleName());
@@ -103,12 +122,139 @@ public static boolean isHTTPTransportMode(HiveConf hiveConf) {
   @Override
   public synchronized void start() {
     super.start();
+    HiveConf hiveConf = this.getHiveConf();
+    if (hiveConf.getBoolVar(ConfVars.HIVE_SERVER2_SUPPORT_DYNAMIC_SERVICE_DISCOVERY)) {
+      try {
+        addServerInstanceToZooKeeper(hiveConf);
+      } catch (Exception e) {
+        LOG.error("Error adding this HiveServer2 instance to ZooKeeper: ", e);
+        throw new ServiceException(e);
+      }
+    }
   }
 
   @Override
   public synchronized void stop() {
     LOG.info("Shutting down HiveServer2");
+    HiveConf hiveConf = this.getHiveConf();
     super.stop();
+    // Remove this server instance from ZooKeeper if dynamic service discovery is set
+    if (hiveConf != null && hiveConf.getBoolVar(ConfVars.HIVE_SERVER2_SUPPORT_DYNAMIC_SERVICE_DISCOVERY)) {
+      try {
+        removeServerInstanceFromZooKeeper();
+      } catch (Exception e) {
+        LOG.error("Error removing znode for this HiveServer2 instance from ZooKeeper.", e);
+      }
+    }
+  }
+
+  private void removeServerInstanceFromZooKeeper() throws Exception {
+    if (znode != null) {
+      znode.close();
+    }
+    zooKeeperClient.close();
+    LOG.info("Server instance removed from ZooKeeper.");
+  }
+
+  private void addServerInstanceToZooKeeper(HiveConf hiveConf) throws Exception {
+    String appId = System.getProperty("spark.yarn.app.id");
+    if (appId == null || appId.trim().isEmpty()) {
+      LOG.error("Only cluster mode thrift server register to ZK, invalid application is {}", appId);
+      return;
+    }
+    zooKeeperClient = startZookeeperClient(hiveConf);
+    // Create a znode under the rootNamespace parent for this instance of the server
+    String rootNamespace = hiveConf.getVar(HiveConf.ConfVars.HIVE_SERVER2_ZOOKEEPER_NAMESPACE);
+    String queue = hiveConf.get("spark.yarn.queue");
+    if (queue == null) {
+      queue = "hdmi-default";
+    }
+    String ephemeralPath = appId + ":" + getServerInstanceURI();
+    String zkPath = ZOOKEEPER_PATH_SEPARATOR + rootNamespace
+            + ZOOKEEPER_PATH_SEPARATOR + queue + ZOOKEEPER_PATH_SEPARATOR + ephemeralPath;
+    LOG.info("Add server instance to zookeeper on {}", zkPath);
+    try {
+      znode = new PersistentEphemeralNode(zooKeeperClient,
+              Mode.EPHEMERAL, zkPath, ephemeralPath.getBytes(StandardCharsets.UTF_8));
+      znode.start();
+      // We'll wait for 120s for node creation
+      long znodeCreationTimeout = 120;
+      if (!znode.waitForInitialCreate(znodeCreationTimeout, TimeUnit.SECONDS)) {
+        throw new Exception("Max znode creation wait time: " + znodeCreationTimeout + "s exhausted");
+      }
+    } catch (Exception e) {
+      LOG.error("Unable to create a znode for this server instance", e);
+      if (znode != null) {
+        znode.close();
+      }
+      throw (e);
+    }
+  }
+
+  private CuratorFramework startZookeeperClient(HiveConf hiveConf) throws Exception {
+    setUpZooKeeperAuth(hiveConf);
+    String zooKeeperEnsemble = ZooKeeperHiveHelper.getQuorumServers(hiveConf);
+    String rootNamespace = hiveConf.getVar(HiveConf.ConfVars.HIVE_SERVER2_ZOOKEEPER_NAMESPACE);
+    int sessionTimeout =
+            (int) hiveConf.getTimeVar(HiveConf.ConfVars.HIVE_ZOOKEEPER_SESSION_TIMEOUT,
+                    TimeUnit.MILLISECONDS);
+    int baseSleepTime =
+            (int) hiveConf.getTimeVar(HiveConf.ConfVars.HIVE_ZOOKEEPER_CONNECTION_BASESLEEPTIME,
+                    TimeUnit.MILLISECONDS);
+    int maxRetries = hiveConf.getIntVar(HiveConf.ConfVars.HIVE_ZOOKEEPER_CONNECTION_MAX_RETRIES);
+    // Create a CuratorFramework instance to be used as the ZooKeeper client
+    // Use the zooKeeperAclProvider to create appropriate ACLs
+    CuratorFramework zooKeeperClient =
+            CuratorFrameworkFactory.builder().connectString(zooKeeperEnsemble)
+                    .sessionTimeoutMs(sessionTimeout)
+                    .retryPolicy(new ExponentialBackoffRetry(baseSleepTime, maxRetries)).build();
+    zooKeeperClient.start();
+    // Create the parent znodes recursively; ignore if the parent already exists.
+    String queue = hiveConf.get("spark.yarn.queue");
+    if (queue == null) {
+      queue = "hdmi-default";
+    }
+    try {
+      zooKeeperClient.create().creatingParentsIfNeeded().withMode(CreateMode.PERSISTENT)
+              .forPath(ZOOKEEPER_PATH_SEPARATOR + rootNamespace + ZOOKEEPER_PATH_SEPARATOR + queue);
+      LOG.info("Created the root name space: " + rootNamespace + " on ZooKeeper for HiveServer2");
+    } catch (KeeperException e) {
+      if (e.code() != KeeperException.Code.NODEEXISTS) {
+        LOG.error("Unable to create HiveServer2 namespace: " + rootNamespace + " on ZooKeeper", e);
+        throw e;
+      }
+    }
+    return zooKeeperClient;
+  }
+
+  /**
+   * For a kerberized cluster, we dynamically set up the client's JAAS conf.
+   *
+   * @param hiveConf
+   * @return
+   * @throws Exception
+   */
+  private void setUpZooKeeperAuth(HiveConf hiveConf) throws Exception {
+    if (UserGroupInformation.isSecurityEnabled()) {
+      String principal = hiveConf.getVar(ConfVars.HIVE_SERVER2_KERBEROS_PRINCIPAL);
+      if (principal.isEmpty()) {
+        throw new IOException("HiveServer2 Kerberos principal is empty");
+      }
+      String keyTabFile = hiveConf.getVar(ConfVars.HIVE_SERVER2_KERBEROS_KEYTAB);
+      if (keyTabFile.isEmpty()) {
+        throw new IOException("HiveServer2 Kerberos keytab is empty");
+      }
+      // Install the JAAS Configuration for the runtime
+      Utils.setZookeeperClientKerberosJaasConfig(principal, keyTabFile);
+    }
+  }
+
+  private String getServerInstanceURI() throws Exception {
+    if ((thriftCLIService == null) || (thriftCLIService.getServerIPAddress() == null)) {
+      throw new Exception("Unable to get the server address; it hasn't been initialized yet.");
+    }
+    return thriftCLIService.getServerIPAddress().getHostName() + ":"
+            + thriftCLIService.getPortNumber();
   }
 
   private static void startHiveServer2() throws Throwable {