Merge branch 'TRY'

liu1028 · Sep 23, 2020 · e9386a4 · e9386a4
2 parents e946caf + 947388f
commit e9386a4
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git a/external/mlsql-healthy/src/main/java/tech/mlsql/plugins/healthy/App.scala b/external/mlsql-healthy/src/main/java/tech/mlsql/plugins/healthy/App.scala
@@ -12,6 +12,7 @@ import tech.mlsql.version.VersionCompatibility
 class App extends tech.mlsql.app.App with VersionCompatibility {
   override def run(args: Seq[String]): Unit = {
     AppRuntimeStore.store.registerController("/plugins/healthy", classOf[HealthyController].getName)
+    AppRuntimeStore.store.registerRequestCleaner("/plugins/cleaner/healthy",classOf[ShutdownOnSparkContextStoppedCleaner].getName)
   }
 
   override def supportedVersions: Seq[String] = Seq("1.5.0-SNAPSHOT", "1.5.0", "1.6.0-SNAPSHOT", "1.6.0")

diff --git a/...althy/src/main/java/tech/mlsql/plugins/healthy/ShutdownOnSparkContextStoppedCleaner.scala b/...althy/src/main/java/tech/mlsql/plugins/healthy/ShutdownOnSparkContextStoppedCleaner.scala
@@ -0,0 +1,21 @@
+package tech.mlsql.plugins.healthy
+
+import streaming.dsl.ScriptSQLExec
+import tech.mlsql.app.RequestCleaner
+
+/**
+ * 23/9/2020 WilliamZhu([email protected])
+ */
+class ShutdownOnSparkContextStoppedCleaner extends RequestCleaner {
+  override def run(): Unit = {
+    val context = ScriptSQLExec.context()
+    val session = context.execListener.sparkSession
+    val wow = session.conf.get("spark.mlsql.plugins.healthy.shutdownOnSparkContextStoppedCleaner", "false").toBoolean
+    if (wow) {
+      if (session.sparkContext.isStopped) {
+        System.exit(-1)
+      }
+    }
+
+  }
+}