Added conf to specify graceful shutdown

rxin · May 21, 2015 · e3d5475 · e3d5475
1 parent 4c18652
commit e3d5475
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/StreamingContext.scala b/streaming/src/main/scala/org/apache/spark/streaming/StreamingContext.scala
@@ -565,11 +565,8 @@ class StreamingContext private[streaming] (
           state = StreamingContextState.ACTIVE
           StreamingContext.setActiveContext(this)
         }
-        shutdownHookRef = Utils.addShutdownHook(StreamingContext.SHUTDOWN_HOOK_PRIORITY) { () =>
-          logInfo("Invoking stop() from shutdown hook")
-          // Do not stop SparkContext, let its own shutdown hook stop it
-          stop(stopSparkContext = false)
-        }
+        shutdownHookRef = Utils.addShutdownHook(
+          StreamingContext.SHUTDOWN_HOOK_PRIORITY)(stopOnShutdown)
         logInfo("StreamingContext started")
       case ACTIVE =>
         logWarning("StreamingContext has already been started")
@@ -659,6 +656,13 @@ class StreamingContext private[streaming] (
       state = STOPPED
     }
   }
+
+  private def stopOnShutdown(): Unit = {
+    val stopGracefully = conf.getBoolean("spark.streaming.stopGracefullyOnShutdown", false)
+    logInfo(s"Invoking stop(stopGracefully=$stopGracefully) from shutdown hook")
+    // Do not stop SparkContext, let its own shutdown hook stop it
+    stop(stopSparkContext = false, stopGracefully = stopGracefully)
+  }
 }
 
 /**