[SPARK-46575][SQL][HIVE] Make HiveThriftServer2.startWithContext DevelopApi retriable and fix flakiness of ThriftServerWithSparkContextInHttpSuite

sql/hive-thriftserver/src/main/java/org/apache/hive/service/cli/thrift/ThriftBinaryCLIService.java

@@ -43,6 +43,7 @@
 import org.apache.thrift.transport.TServerSocket;
 import org.apache.thrift.transport.TTransportFactory;
 
+import org.apache.spark.sql.hive.thriftserver.HiveThriftServer2$;
 
 public class ThriftBinaryCLIService extends ThriftCLIService {
 
@@ -137,7 +138,11 @@ public void run() {
       LOG.error(
           "Error starting HiveServer2: could not start "
               + ThriftBinaryCLIService.class.getSimpleName(), t);
-      System.exit(-1);
+      if (HiveThriftServer2$.MODULE$.systemExitOnError().get()) {
+        System.exit(-1);
+      } else {
+        throw new ServiceException(t);
+      }
     }
   }
 

sql/hive-thriftserver/src/main/java/org/apache/hive/service/cli/thrift/ThriftHttpCLIService.java

@@ -33,6 +33,7 @@
 import org.apache.hive.service.rpc.thrift.TCLIService;
 import org.apache.hive.service.rpc.thrift.TCLIService.Iface;
 import org.apache.hive.service.server.ThreadFactoryWithGarbageCleanup;
+import org.apache.spark.sql.hive.thriftserver.HiveThriftServer2$;
 import org.apache.thrift.TProcessor;
 import org.apache.thrift.protocol.TBinaryProtocol;
 import org.apache.thrift.protocol.TProtocolFactory;
@@ -47,7 +48,6 @@
 import org.eclipse.jetty.util.thread.ExecutorThreadPool;
 import org.eclipse.jetty.util.thread.ScheduledExecutorScheduler;
 
-
 public class ThriftHttpCLIService extends ThriftCLIService {
 
   protected org.eclipse.jetty.server.Server httpServer;
@@ -182,7 +182,11 @@ public void run() {
       } else {
         LOG.error("Error starting HiveServer2: could not start "
             + ThriftHttpCLIService.class.getSimpleName(), t);
-        System.exit(-1);
+        if (HiveThriftServer2$.MODULE$.systemExitOnError().get()) {
+          System.exit(-1);
+        } else {
+          throw new ServiceException(t);
+        }
       }
     }
   }

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

@@ -45,13 +45,21 @@ object HiveThriftServer2 extends Logging {
   var uiTab: Option[ThriftServerTab] = None
   var listener: HiveThriftServer2Listener = _
   var eventManager: HiveThriftServer2EventManager = _
+  val systemExitOnError = new AtomicBoolean(true)
 
   /**
    * :: DeveloperApi ::
    * Starts a new thrift server with the given context.
+   *
+   * @param sqlContext SQLContext to use for the server
+   * @param exitOnError Whether to exit the JVM if HiveThriftServer2 fails to initialize. When true,
+   *                    the call logs the error and exits the JVM with exit code -1. When false, the
+   *                    call throws an exception instead.
    */
   @DeveloperApi
-  def startWithContext(sqlContext: SQLContext): HiveThriftServer2 = {
+  def startWithContext(sqlContext: SQLContext, exitOnError: Boolean = true): HiveThriftServer2 = {
+    systemExitOnError.set(exitOnError)
+
     val executionHive = HiveUtils.newClientForExecution(
       sqlContext.sparkContext.conf,
       sqlContext.sessionState.newHadoopConf())

sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/SharedThriftServer.scala

@@ -30,12 +30,13 @@ import org.apache.hadoop.hive.ql.metadata.Hive
 import org.apache.hadoop.hive.ql.session.SessionState
 import org.apache.hive.jdbc.HttpBasicAuthInterceptor
 import org.apache.hive.service.auth.PlainSaslHelper
-import org.apache.hive.service.cli.thrift.{ThriftCLIService, ThriftCLIServiceClient}
+import org.apache.hive.service.cli.thrift.{ThriftBinaryCLIService, ThriftCLIService, ThriftCLIServiceClient}
 import org.apache.hive.service.rpc.thrift.TCLIService.Client
 import org.apache.http.impl.client.HttpClientBuilder
 import org.apache.thrift.protocol.TBinaryProtocol
 import org.apache.thrift.transport.{THttpClient, TSocket}
 
+import org.apache.spark.SparkException
 import org.apache.spark.sql.test.SharedSparkSession
 import org.apache.spark.util.Utils
 
@@ -138,10 +139,18 @@ trait SharedThriftServer extends SharedSparkSession {
     sqlContext.setConf(ConfVars.HIVE_START_CLEANUP_SCRATCHDIR.varname, "true")
 
     try {
-      hiveServer2 = HiveThriftServer2.startWithContext(sqlContext)
+      // Set exitOnError to false to avoid exiting the JVM process and tearing down the SparkContext
+      // instance in case of any exceptions here. Otherwise, the following retries are doomed to
+      // fail on a stopped context.
+      hiveServer2 = HiveThriftServer2.startWithContext(sqlContext, exitOnError = false)
       hiveServer2.getServices.asScala.foreach {
         case t: ThriftCLIService =>
           serverPort = t.getPortNumber
+          if (t.isInstanceOf[ThriftBinaryCLIService] && mode == ServerMode.http) {
+            logError("A previous Hive's SessionState is leaked, aborting this retry")
+            throw SparkException.internalError("HiveThriftServer2 started in binary mode " +
+              "while the test case is expecting HTTP mode")
+          }
           logInfo(s"Started HiveThriftServer2: mode=$mode, port=$serverPort, attempt=$attempt")
         case _ =>
       }