[SPARK-21606][SQL]HiveThriftServer2 catches OOMs on request threads

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

@@ -0,0 +1,55 @@
+/**
+ * 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.hive.service.cli.thrift;
+
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.Future;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+
+final class ThreadPoolExecutorWithOomHook extends ThreadPoolExecutor {
+  private final Runnable oomHook;
+
+  public ThreadPoolExecutorWithOomHook(int corePoolSize, int maximumPoolSize, long keepAliveTime,
+      TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory,
+      Runnable oomHook) {
+    super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory);
+    this.oomHook = oomHook;
+  }
+
+  @Override
+  protected void afterExecute(Runnable r, Throwable t) {
+    super.afterExecute(r, t);
+    if (t == null && r instanceof Future<?>) {
+      try {
+        Future<?> future = (Future<?>) r;
+        if (future.isDone()) {
+          future.get();
+        }
+      } catch (InterruptedException ie) {
+        Thread.currentThread().interrupt();
+      } catch (Throwable t2) {
+        t = t2;
+      }
+    }
+    if (t instanceof OutOfMemoryError) {
+      oomHook.run();
+    }
+  }
+}

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

@@ -18,13 +18,6 @@
 
 package org.apache.hive.service.cli.thrift;
 
-import java.util.ArrayList;
-import java.util.List;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.SynchronousQueue;
-import java.util.concurrent.ThreadPoolExecutor;
-import java.util.concurrent.TimeUnit;
-
 import org.apache.hadoop.hive.conf.HiveConf;
 import org.apache.hadoop.hive.conf.HiveConf.ConfVars;
 import org.apache.hadoop.hive.shims.ShimLoader;
@@ -37,21 +30,29 @@
 import org.apache.thrift.transport.TServerSocket;
 import org.apache.thrift.transport.TTransportFactory;
 
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.SynchronousQueue;
+import java.util.concurrent.TimeUnit;
+
 
 public class ThriftBinaryCLIService extends ThriftCLIService {
+  private final Runnable oomHook;
 
-  public ThriftBinaryCLIService(CLIService cliService) {
+  public ThriftBinaryCLIService(CLIService cliService, Runnable oomHook) {
     super(cliService, ThriftBinaryCLIService.class.getSimpleName());
+    this.oomHook = oomHook;
   }
 
   @Override
   public void run() {
     try {
       // Server thread pool
       String threadPoolName = "HiveServer2-Handler-Pool";
-      ExecutorService executorService = new ThreadPoolExecutor(minWorkerThreads, maxWorkerThreads,
+      ExecutorService executorService = new ThreadPoolExecutorWithOomHook(minWorkerThreads, maxWorkerThreads,
           workerKeepAliveTime, TimeUnit.SECONDS, new SynchronousQueue<Runnable>(),
-          new ThreadFactoryWithGarbageCleanup(threadPoolName));
+          new ThreadFactoryWithGarbageCleanup(threadPoolName), oomHook);
 
       // Thrift configs
       hiveAuthFactory = new HiveAuthFactory(hiveConf);

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

@@ -438,6 +438,9 @@ public TExecuteStatementResp ExecuteStatement(TExecuteStatementReq req) throws T
           resp.setOperationHandle(operationHandle.toTOperationHandle());
           resp.setStatus(OK_STATUS);
     } catch (Exception e) {
+      // Note: it's rather important that this (and other methods) catch Exception, not Throwable;
+      //       in combination with HiveSessionProxy.invoke code, perhaps unintentionally, it used
+      //       to also catch all errors; and now it allows OOMs only to propagate.
       LOG.warn("Error executing statement: ", e);
       resp.setStatus(HiveSQLException.toTStatus(e));
     }

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

@@ -49,9 +49,11 @@
 
 
 public class ThriftHttpCLIService extends ThriftCLIService {
+  private final Runnable oomHook;
 
-  public ThriftHttpCLIService(CLIService cliService) {
+  public ThriftHttpCLIService(CLIService cliService, Runnable oomHook) {
     super(cliService, ThriftHttpCLIService.class.getSimpleName());
+    this.oomHook = oomHook;
   }
 
   /**
@@ -65,9 +67,9 @@ public void run() {
       // Server thread pool
       // Start with minWorkerThreads, expand till maxWorkerThreads and reject subsequent requests
       String threadPoolName = "HiveServer2-HttpHandler-Pool";
-      ExecutorService executorService = new ThreadPoolExecutor(minWorkerThreads, maxWorkerThreads,
+      ExecutorService executorService = new ThreadPoolExecutorWithOomHook(minWorkerThreads, maxWorkerThreads,
           workerKeepAliveTime, TimeUnit.SECONDS, new SynchronousQueue<Runnable>(),
-          new ThreadFactoryWithGarbageCleanup(threadPoolName));
+          new ThreadFactoryWithGarbageCleanup(threadPoolName), oomHook);
       ExecutorThreadPool threadPool = new ExecutorThreadPool(executorService);
 
       // HTTP Server

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

@@ -18,14 +18,7 @@
 
 package org.apache.hive.service.server;
 
-import java.util.Properties;
-
-import org.apache.commons.cli.GnuParser;
-import org.apache.commons.cli.HelpFormatter;
-import org.apache.commons.cli.Option;
-import org.apache.commons.cli.OptionBuilder;
-import org.apache.commons.cli.Options;
-import org.apache.commons.cli.ParseException;
+import org.apache.commons.cli.*;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.hadoop.hive.common.LogUtils;
@@ -39,6 +32,8 @@
 import org.apache.hive.service.cli.thrift.ThriftCLIService;
 import org.apache.hive.service.cli.thrift.ThriftHttpCLIService;
 
+import java.util.Properties;
+
 /**
  * HiveServer2.
  *
@@ -58,16 +53,22 @@ public HiveServer2() {
   public synchronized void init(HiveConf hiveConf) {
     cliService = new CLIService(this);
     addService(cliService);
+    final HiveServer2 hiveServer2 = this;
+    Runnable oomHook = new Runnable() {
+      @Override
+      public void run() {
+        hiveServer2.stop();
+      }
+    };
     if (isHTTPTransportMode(hiveConf)) {
-      thriftCLIService = new ThriftHttpCLIService(cliService);
+      thriftCLIService = new ThriftHttpCLIService(cliService, oomHook);
     } else {
-      thriftCLIService = new ThriftBinaryCLIService(cliService);
+      thriftCLIService = new ThriftBinaryCLIService(cliService, oomHook);
     }
     addService(thriftCLIService);
     super.init(hiveConf);
 
     // Add a shutdown hook for catching SIGTERM & SIGINT
-    final HiveServer2 hiveServer2 = this;
     Runtime.getRuntime().addShutdownHook(new Thread() {
       @Override
       public void run() {

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

@@ -281,10 +281,16 @@ private[hive] class HiveThriftServer2(sqlContext: SQLContext)
     setSuperField(this, "cliService", sparkSqlCliService)
     addService(sparkSqlCliService)
 
+    val oomHook = new Runnable {
+      override def run(): Unit = {
+        System.exit(-1)
+      }
+    }
+
     val thriftCliService = if (isHTTPTransportMode(hiveConf)) {
-      new ThriftHttpCLIService(sparkSqlCliService)
+      new ThriftHttpCLIService(sparkSqlCliService, oomHook)
     } else {
-      new ThriftBinaryCLIService(sparkSqlCliService)
+      new ThriftBinaryCLIService(sparkSqlCliService, oomHook)
     }
 
     setSuperField(this, "thriftCLIService", thriftCliService)