HIVE-29477: Introduce TezSession interface and a new implementation for unmanaged, external Tez sessions

common/src/java/org/apache/hadoop/hive/conf/HiveConf.java

@@ -3998,6 +3998,15 @@
       "true", new StringSet("true", "false", "ignore"),
       "Whether Tez session pool should allow submitting queries to custom queues. The options\n" +
       "are true, false (error out), ignore (accept the query but ignore the queue setting)."),
+    HIVE_SERVER2_TEZ_USE_EXTERNAL_SESSIONS("hive.server2.tez.use.external.sessions", false,
+      "This flag is used in HiveServer2 to use externally started tez sessions"),
+    HIVE_SERVER2_TEZ_EXTERNAL_SESSIONS_NAMESPACE("hive.server2.tez.external.sessions.namespace", "",
+      "ZK namespace to use for tez external sessions"),
+    HIVE_SERVER2_TEZ_EXTERNAL_SESSIONS_WAIT_MAX_ATTEMPTS("hive.server2.tez.external.sessions.wait.max.attempts",
+      60, "Number of attempts before giving up waiting for external sessions (each attempt is 1 sec long)"),
+    HIVE_SERVER2_TEZ_EXTERNAL_SESSIONS_REGISTRY_CLASS("hive.server2.tez.external.sessions.registry.class",
+      "org.apache.hadoop.hive.ql.exec.tez.DummyExternalSessionsRegistry", "Tez external sessions\n" +
+      "registry implementation to use"),
     HIVE_MAPRED_JOB_FOLLOW_TEZ_QUEUE("hive.mapred.job.follow.tez.queue", false,
         "Whether the MR jobs initiated by a query should be enforced to run in the queue denoted by "
             + "'tez.queue.name', e.g. DistCp jobs."),
@@ -5759,7 +5768,7 @@
 
     HIVE_OTEL_RETRY_BACKOFF_MULTIPLIER("hive.otel.retry.backoff.multiplier", 5f,
         "The multiplier applied to the backoff interval for retries in the OTEL exporter."
             + "This determines how much the backoff interval increases after each failed attempt, following an exponential backoff strategy."),
 
     HIVE_OTEL_EXPORT_TEZ_COUNTERS("hive.otel.export.tez.counters", false,
         "Enables the inclusion of Tez counters in OTEL output for a particular query.");

itests/util/src/main/java/org/apache/hadoop/hive/ql/QTestMiniClusters.java

@@ -50,7 +50,7 @@
 import org.apache.hadoop.hive.llap.daemon.MiniLlapCluster;
 import org.apache.hadoop.hive.llap.io.api.LlapProxy;
 import org.apache.hadoop.hive.ql.exec.Utilities;
-import org.apache.hadoop.hive.ql.exec.tez.TezSessionState;
+import org.apache.hadoop.hive.ql.exec.tez.TezSession;
 import org.apache.hadoop.hive.ql.lockmgr.zookeeper.CuratorFrameworkSingleton;
 import org.apache.hadoop.hive.ql.lockmgr.zookeeper.ZooKeeperHiveLockManager;
 import org.apache.hadoop.hive.ql.session.SessionState;
@@ -358,7 +358,7 @@ public void restartSessions(boolean canReuseSession, CliSessionState ss, Session
     if (oldSs != null && canReuseSession
         && clusterType.getCoreClusterType() == CoreClusterType.TEZ) {
       // Copy the tezSessionState from the old CliSessionState.
-      TezSessionState tezSessionState = oldSs.getTezSession();
+      TezSession tezSessionState = oldSs.getTezSession();
       oldSs.setTezSession(null);
       ss.setTezSession(tezSessionState);
       oldSs.close();

ql/pom.xml

@@ -343,7 +343,6 @@
     <dependency>
       <groupId>org.apache.hadoop</groupId>
       <artifactId>hadoop-yarn-api</artifactId>
-      <optional>true</optional>
     </dependency>
     <dependency>
       <groupId>org.apache.hadoop</groupId>

ql/src/java/org/apache/hadoop/hive/ql/ddl/process/kill/KillQueriesOperation.java

@@ -18,6 +18,7 @@
 
 package org.apache.hadoop.hive.ql.ddl.process.kill;
 
+import org.apache.hadoop.hive.conf.HiveConf;
 import org.apache.hadoop.hive.ql.ddl.DDLOperation;
 import org.apache.hadoop.hive.ql.ddl.DDLOperationContext;
 import org.apache.hadoop.hive.ql.metadata.HiveException;
@@ -37,7 +38,12 @@ public KillQueriesOperation(DDLOperationContext context, KillQueriesDesc desc) {
   public int execute() throws HiveException {
     SessionState sessionState = SessionState.get();
     for (String queryId : desc.getQueryIds()) {
-      sessionState.getKillQuery().killQuery(queryId, KILL_QUERY_MESSAGE, context.getDb().getConf());
+      // For now, get the config setting here; we can only have one type of the session present.
+      // Ideally we should check each session separately.
+      boolean isExternal = HiveConf.getBoolVar(context.getDb().getConf(),
+              HiveConf.ConfVars.HIVE_SERVER2_TEZ_USE_EXTERNAL_SESSIONS);
+      sessionState.getKillQuery().killQuery(queryId, KILL_QUERY_MESSAGE, context.getDb().getConf(),
+              !isExternal);
     }
     LOG.info("kill query called ({})", desc.getQueryIds());
     return 0;

ql/src/java/org/apache/hadoop/hive/ql/exec/tez/AbstractTriggerValidator.java

@@ -0,0 +1,53 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hadoop.hive.ql.exec.tez;
+
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+
+import com.google.common.util.concurrent.ThreadFactoryBuilder;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public abstract class AbstractTriggerValidator {
+  private static final Logger LOG = LoggerFactory.getLogger(AbstractTriggerValidator.class);
+
+  private ScheduledExecutorService scheduledExecutorService = null;
+  abstract Runnable getTriggerValidatorRunnable();
+
+  void startTriggerValidator(long triggerValidationIntervalMs) {
+    if (scheduledExecutorService == null) {
+      scheduledExecutorService = Executors.newSingleThreadScheduledExecutor(
+        new ThreadFactoryBuilder().setDaemon(true).setNameFormat("TriggerValidator").build());
+      Runnable triggerValidatorRunnable = getTriggerValidatorRunnable();
+      scheduledExecutorService.scheduleWithFixedDelay(triggerValidatorRunnable, triggerValidationIntervalMs,
+          triggerValidationIntervalMs, TimeUnit.MILLISECONDS);
+      LOG.info("Started trigger validator with interval: {} ms", triggerValidationIntervalMs);
+    }
+  }
+
+  void stopTriggerValidator() {
+    if (scheduledExecutorService != null) {
+      scheduledExecutorService.shutdownNow();
+      scheduledExecutorService = null;
+      LOG.info("Stopped trigger validator");
+    }
+  }
+}

ql/src/java/org/apache/hadoop/hive/ql/exec/tez/DummyExternalSessionsRegistry.java

@@ -0,0 +1,43 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.hadoop.hive.ql.exec.tez;
+
+
+import org.apache.hadoop.hive.conf.HiveConf;
+
+public class DummyExternalSessionsRegistry implements ExternalSessionsRegistry {
+
+  // This constructor is required. Reflective instantiation will invoke this constructor.
+  public DummyExternalSessionsRegistry(HiveConf conf) {
+  }
+
+  @Override
+  public String getSession() throws Exception {
+    throw new UnsupportedOperationException("not supported in dummy external session registry");
+  }
+
+  @Override
+  public void returnSession(final String appId) {
+    throw new UnsupportedOperationException("not supported in dummy external session registry");
+  }
+
+  @Override
+  public void close() {
+    throw new UnsupportedOperationException("not supported in dummy external session registry");
+  }
+}

ql/src/java/org/apache/hadoop/hive/ql/exec/tez/ExternalSessionsRegistry.java

@@ -0,0 +1,43 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.hadoop.hive.ql.exec.tez;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public interface ExternalSessionsRegistry {
+  Logger LOG = LoggerFactory.getLogger(ExternalSessionsRegistry.class);
+
+  /**
+   * Returns application of id of the external session.
+   * @return application id
+   * @throws Exception in case of any exceptions
+   */
+  String getSession() throws Exception;
+
+  /**
+   * Returns external session back to registry.
+   * @param appId application id
+   */
+  void returnSession(String appId);
+
+  /**
+   * Closes the external session registry.
+   */
+  void close();
+}

ql/src/java/org/apache/hadoop/hive/ql/exec/tez/ExternalSessionsRegistryFactory.java

@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.hadoop.hive.ql.exec.tez;
+
+import org.apache.hadoop.hive.conf.HiveConf;
+import org.apache.hadoop.hive.metastore.api.MetaException;
+import org.apache.hadoop.hive.metastore.utils.JavaUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Factory for creating and managing {@code ExternalSessionsRegistry} clients.
+ * <p>
+ * Instances created by this factory are cached and may be reused across calls.
+ * Callers do not need to manage the lifecycle or closure of these clients,
+ * as this factory is responsible for handling those concerns.
+ * </p>
+ */
+public abstract class ExternalSessionsRegistryFactory {
+  private static final Logger LOG = LoggerFactory.getLogger(ExternalSessionsRegistryFactory.class);
+
+  private static final Map<String, ExternalSessionsRegistry> CLIENTS = new HashMap<>();
+
+  static {
+    Runtime.getRuntime().addShutdownHook(new Thread(() -> CLIENTS.values().forEach(ExternalSessionsRegistry::close)));
+  }
+
+  private ExternalSessionsRegistryFactory() {
+  }
+
+  public static ExternalSessionsRegistry getClient(final HiveConf conf) {
+    HiveConf newConf = prepareConf(conf);
+    // TODO: change this to TezConfiguration.TEZ_AM_REGISTRY_NAMESPACE after Tez 1.0.0 is released.
+    String namespace = conf.get("tez.am.registry.namespace");
+
+    ExternalSessionsRegistry registry;
+    synchronized (CLIENTS) {
+      registry = CLIENTS.computeIfAbsent(namespace, ns -> {
+        try {
+          String clazz = HiveConf.getVar(newConf, HiveConf.ConfVars.HIVE_SERVER2_TEZ_EXTERNAL_SESSIONS_REGISTRY_CLASS);
+
+          return JavaUtils.newInstance(JavaUtils.getClass(clazz, ExternalSessionsRegistry.class),
+              new Class<?>[]{HiveConf.class}, new Object[]{newConf});
+        } catch (MetaException e) {
+          throw new RuntimeException(e);
+        }
+      });
+    }
+    LOG.info("Returning tez external AM registry ({}) for namespace '{}'", System.identityHashCode(registry),
+        namespace);
+    return registry;
+  }
+
+  private static HiveConf prepareConf(HiveConf conf) {
+    // HS2 would need to know about all coordinators running on all compute groups for a given compute (namespace)
+    // Setting this config to false in client, will make registry client listen on paths under @compute instead of
+    // @compute/compute-group
+    HiveConf newConf = new HiveConf(conf, ExternalSessionsRegistryFactory.class);
+
+    // TODO: change this to TezConfiguration.TEZ_AM_REGISTRY_ENABLE_COMPUTE_GROUPS after Tez 1.0.0 is released.
+    newConf.setBoolean("tez.am.registry.enable.compute.groups", false);
+    return newConf;
+  }
+}

ql/src/java/org/apache/hadoop/hive/ql/exec/tez/KillTriggerActionHandler.java

@@ -22,7 +22,6 @@
 import org.apache.hadoop.hive.conf.HiveConf;
 import org.apache.hadoop.hive.ql.DriverUtils;
 import org.apache.hadoop.hive.ql.metadata.HiveException;
-import org.apache.hadoop.hive.ql.session.KillQuery;
 import org.apache.hadoop.hive.ql.wm.Action;
 import org.apache.hadoop.hive.ql.wm.Trigger;
 import org.apache.hadoop.hive.ql.wm.TriggerActionHandler;
@@ -33,35 +32,33 @@
 /**
  * Handles only Kill Action.
  */
-public class KillTriggerActionHandler implements TriggerActionHandler<TezSessionState> {
+public class KillTriggerActionHandler implements TriggerActionHandler<TezSession> {
   private static final Logger LOG = LoggerFactory.getLogger(KillTriggerActionHandler.class);
   private final HiveConf conf;
 
   public KillTriggerActionHandler() {
-      this.conf = new HiveConf();
+    this.conf = new HiveConf();
   }
 
   @Override
-  public void applyAction(final Map<TezSessionState, Trigger> queriesViolated) {
-    for (Map.Entry<TezSessionState, Trigger> entry : queriesViolated.entrySet()) {
-        if (entry.getValue().getAction().getType() == Action.Type.KILL_QUERY) {
-            TezSessionState sessionState = entry.getKey();
-            String queryId = sessionState.getWmContext().getQueryId();
-            try {
-                UserGroupInformation ugi = UserGroupInformation.getCurrentUser();
-                DriverUtils.setUpAndStartSessionState(conf, ugi.getShortUserName());
-                KillQuery killQuery = sessionState.getKillQuery();
-                // if kill query is null then session might have been released to pool or closed already
-                if (killQuery != null) {
-                    sessionState.getKillQuery().killQuery(queryId, entry.getValue().getViolationMsg(),
-                            sessionState.getConf());
-                }
-            } catch (HiveException | IOException e) {
-                LOG.warn("Unable to kill query {} for trigger violation", queryId);
-            }
-        } else {
-            throw new RuntimeException("Unsupported action: " + entry.getValue());
+  public void applyAction(Map<TezSession, Trigger> queriesViolated) {
+    for (Map.Entry<TezSession, Trigger> entry : queriesViolated.entrySet()) {
+      if (entry.getValue().getAction().getType() == Action.Type.KILL_QUERY) {
+        TezSession sessionState = entry.getKey();
+        String queryId = sessionState.getWmContext().getQueryId();
+        try {
+          UserGroupInformation ugi = UserGroupInformation.getCurrentUser();
+          DriverUtils.setUpAndStartSessionState(conf, ugi.getShortUserName());
+          boolean wasKilled = sessionState.killQuery(entry.getValue().getViolationMsg());
+          if (!wasKilled) {
+            LOG.info("Didn't kill the query {}", queryId);
+          }
+        } catch (HiveException | IOException e) {
+          LOG.warn("Unable to kill query {} for trigger violation", queryId, e);
         }
+      } else {
+        throw new RuntimeException("Unsupported action: " + entry.getValue());
+      }
     }
   }
 }