HIVE-29302: Wider Tez session open scope should be guarded to prevent local resource leaks

Author: abstractdog

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

@@ -272,9 +272,6 @@ protected void openInternal(String[] additionalFilesNotFromConf,
     this.queueName = confQueueName;
     this.doAsEnabled = conf.getBoolVar(HiveConf.ConfVars.HIVE_SERVER2_ENABLE_DOAS);
 
-    final boolean llapMode = "llap".equalsIgnoreCase(HiveConf.getVar(
-        conf, HiveConf.ConfVars.HIVE_EXECUTION_MODE));
-
     // TODO This - at least for the session pool - will always be the hive user. How does doAs above this affect things ?
     UserGroupInformation ugi = Utils.getUGI();
     user = ugi.getShortUserName();
@@ -292,12 +289,36 @@ protected void openInternal(String[] additionalFilesNotFromConf,
       LOG.info("Created new resources: " + this.resources);
     }
 
-    // unless already installed on all the cluster nodes, we'll have to
-    // localize hive-exec.jar as well.
+    // unless already installed on all the cluster nodes, we'll have to localize hive-exec.jar as well.
     appJarLr = createJarLocalResource(utils.getExecJarPathLocal(conf));
 
-    // configuration for the application master
-    final Map<String, LocalResource> commonLocalResources = new HashMap<String, LocalResource>();
+    try {
+      openInternalUnsafe(isAsync, console);
+    } catch (Exception e) {
+      LOG.info("Failed to open session, deleting scratch dir to prevent LocalResource leak...", e);
+      cleanupScratchDir();
+      throw e;
+    }
+  }
+
+  /**
+   * Opens a Tez session without performing a complete rollback/cleanup on failure.
+   *
+   * <p><strong>Callers MUST guard this method with try/catch and perform cleanup</strong>
+   * of partially initialized state (such as localized files in the scratch directory).
+   * This method is not safe on its own.</p>
+   *
+   * @param isAsync whether to open the Tez session asynchronously in a separate thread
+   * @param console a {@link LogHelper} used to log session startup events
+   *
+   * @throws TezException if the session fails to start (including failures during
+   *                      container launch or session initialization)
+   * @throws IOException if local resource localization or I/O setup fails
+   */
+  private void openInternalUnsafe(boolean isAsync, LogHelper console) throws TezException, IOException {
+    final Map<String, LocalResource> commonLocalResources = new HashMap<>();
+    final boolean llapMode = "llap".equalsIgnoreCase(HiveConf.getVar(conf, HiveConf.ConfVars.HIVE_EXECUTION_MODE));
+
     commonLocalResources.put(DagUtils.getBaseName(appJarLr), appJarLr);
     for (LocalResource lr : this.resources.localizedResources) {
       commonLocalResources.put(DagUtils.getBaseName(lr), lr);
@@ -312,7 +333,7 @@ protected void openInternal(String[] additionalFilesNotFromConf,
     }
 
     // Create environment for AM.
-    Map<String, String> amEnv = new HashMap<String, String>();
+    Map<String, String> amEnv = new HashMap<>();
     MRHelpers.updateEnvBasedOnMRAMEnv(conf, amEnv);
 
     // and finally we're ready to create and start the session
@@ -383,27 +404,24 @@ protected void openInternal(String[] additionalFilesNotFromConf,
       startSessionAndContainers(session, conf, commonLocalResources, tezConfig, false);
       this.session = session;
     } else {
-      FutureTask<TezClient> sessionFuture = new FutureTask<>(new Callable<TezClient>() {
-        @Override
-        public TezClient call() throws Exception {
-          TezClient result = null;
-          try {
-            result = startSessionAndContainers(
-                session, conf, commonLocalResources, tezConfig, true);
-          } catch (Throwable t) {
-            // The caller has already stopped the session.
-            LOG.error("Failed to start Tez session", t);
-            throw (t instanceof Exception) ? (Exception)t : new Exception(t);
-          }
-          // Check interrupt at the last moment in case we get cancelled quickly.
-          // This is not bulletproof but should allow us to close session in most cases.
-          if (Thread.interrupted()) {
-            LOG.info("Interrupted while starting Tez session");
-            closeAndIgnoreExceptions(result);
-            return null;
-          }
-          return result;
+      FutureTask<TezClient> sessionFuture = new FutureTask<>(() -> {
+        TezClient result = null;
+        try {
+          result = startSessionAndContainers(
+              session, conf, commonLocalResources, tezConfig, true);
+        } catch (Throwable t) {
+          // The caller has already stopped the session.
+          LOG.error("Failed to start Tez session", t);
+          throw (t instanceof Exception) ? (Exception)t : new Exception(t);
+        }
+        // Check interrupt at the last moment in case we get cancelled quickly.
+        // This is not bulletproof but should allow us to close session in most cases.
+        if (Thread.interrupted()) {
+          LOG.info("Interrupted while starting Tez session");
+          closeAndIgnoreExceptions(result);
+          return null;
         }
+        return result;
       });
       new Thread(sessionFuture, "Tez session start thread").start();
       // We assume here nobody will try to get session before open() returns.