Address Co-pilot comments and fix the cases for default Sessions

Author: tanishq-chugh

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

@@ -40,4 +40,11 @@ public interface ExternalSessionsRegistry {
    * Closes the external session registry.
    */
   void close();
+
+  /**
+   * Returns true if this registry instance currently holds a claim on the given AM.
+   */
+  default boolean isClaimed(String appId) {
+    return true; // Non-ZK registries case is always true
+  }
 }

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

@@ -198,6 +198,9 @@ public boolean killQuery(String reason) throws HiveException {
 
   @Override
   public DAGClient submitDAG(DAG dag) throws TezException, IOException {
+    if (!registry.isClaimed(externalAppId)) {
+      throw new TezException("Cannot submit DAG as the Tez Session no-longer owns the AM: " + externalAppId);
+    }
     try {
       return getTezClient().submitDAG(dag);
     } catch (TezException e) {
@@ -210,6 +213,9 @@ public DAGClient submitDAG(DAG dag) throws TezException, IOException {
   }
 
   private void tryKillRunningDAGs(TezClient session) throws TezException {
+    if (!registry.isClaimed(externalAppId)) {
+      throw new TezException("Cannot kill running DAG as the Tez Session no-longer owns the AM: " + externalAppId);
+    }
     LOG.info("External session has an AM which is already running a DAG on app ID {}", externalAppId);
     DAGClientAMProtocolBlockingPB proxy = session.sendAMHeartbeat(null);
     if (proxy == null) {

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

@@ -354,7 +354,7 @@ void returnSession(TezSession tezSessionState) {
         defaultSessionPool.returnSession((TezSessionPoolSession)tezSessionState);
       }
 
-      if (useExternalSessions) {
+      if (useExternalSessions && !tezSessionState.isDefault()) {
         if (tezSessionState.getTezClient() != null
             && tezSessionState.getTezClient().getAppMasterApplicationId() != null) {
           try {

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

@@ -472,9 +472,8 @@ void ensureSessionHasResources(
       TezSession session, String[] nonConfResources) throws Exception {
     TezClient client = session.getTezClient();
     // TODO null can also mean that this operation was interrupted. Should we really try to re-create the session in that case ?
-    if (client == null) {
-      // Note: the only sane case where this can happen is the non-pool one. We should get rid
-      //       of it, in non-pool case perf doesn't matter so we might as well open at get time
+    if (client == null || !session.isOpen()) {
+      // Note: We should get rid of it, in non-pool case perf doesn't matter so we might as well open at get time
       //       and then call update like we do in the else.
       // Can happen if the user sets the tez flag after the session was established.
       LOG.info("Tez session hasn't been created yet. Opening session");

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

@@ -61,6 +61,7 @@ public class ZookeeperExternalSessionsRegistryClient implements ExternalSessions
   private InterProcessMutex globalQueue;
   private String claimsPath;
   private volatile boolean isInitialized;
+  private volatile boolean zkConnectionHealthy = false;
 
 
   public ZookeeperExternalSessionsRegistryClient(final HiveConf initConf) {
@@ -99,17 +100,17 @@ private void init() {
       client.start();
 
       client.getConnectionStateListenable().addListener((curatorClient, newState) -> {
-        if (newState != ConnectionState.LOST) {
-          return;
-        }
-        Set<String> sessionsToKill;
-        synchronized (lock) {
-          LOG.error("ZK connection state has changed to lost; killing running DAGs on claimed AMs: {}", taken);
-          sessionsToKill = new HashSet<>(taken);
-          taken.clear();
-        }
-        for (String appId : sessionsToKill) {
-          TezJobMonitor.killRunningDAGsForApplication(appId);
+        if (newState == ConnectionState.CONNECTED || newState == ConnectionState.RECONNECTED) {
+          zkConnectionHealthy = true;
+        } else if (newState == ConnectionState.LOST) {
+          zkConnectionHealthy = false;
+          Set<String> sessionsToKill;
+          synchronized (lock) {
+            LOG.error("ZK connection state has changed to lost; killing running DAGs on claimed AMs: {}", taken);
+            sessionsToKill = new HashSet<>(taken);
+            taken.clear();
+          }
+          sessionsToKill.forEach(TezJobMonitor::killRunningDAGsForApplication);
         }
       });
 
@@ -201,7 +202,9 @@ public String getSession() throws Exception {
           }
           long remainingTimeNs = timeoutNs - (System.nanoTime() - startTimeNs);
           if (remainingTimeNs > 0) {
-            long waitTimeMs = Math.min(1000L, remainingTimeNs / 1_000_000L);
+            // Add one to remainingTime milliseconds computation to prevent the case where
+            // (remainingTimeNs / 1000000L) can return 0 causing the lock to be held indefinitely.
+            long waitTimeMs = Math.min(1000L, (remainingTimeNs / 1000000L) + 1);
             lock.wait(waitTimeMs);
           }
         }
@@ -288,4 +291,14 @@ public void childEvent(final CuratorFramework client, final PathChildrenCacheEve
       }
     }
   }
+
+  @Override
+  public boolean isClaimed(String appId) {
+    if (!zkConnectionHealthy) {
+      return false;
+    }
+    synchronized (lock) {
+      return taken.contains(appId);
+    }
+  }
 }

ql/src/java/org/apache/hadoop/hive/ql/exec/tez/monitoring/TezJobMonitor.java

@@ -507,7 +507,7 @@ public static void killRunningDAGsForApplication(String applicationId) {
             c.tryKillDAG();
           }
         } catch (Exception e) {
-          LOG.error("Error while trying to kill running DAG on tez session {}", applicationId);
+          LOG.error("Error while trying to kill running DAG on tez session {}", applicationId, e);
         }
       }
     }