[SPARK-27958][SQL] Stopping a SparkSession should not always stop Spark Context

sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala

@@ -19,7 +19,7 @@ package org.apache.spark.sql
 
 import java.io.Closeable
 import java.util.concurrent.TimeUnit._
-import java.util.concurrent.atomic.AtomicReference
+import java.util.concurrent.atomic.{AtomicInteger, AtomicReference}
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable
@@ -711,12 +711,15 @@ class SparkSession private(
   // scalastyle:on
 
   /**
-   * Stop the underlying `SparkContext`.
+   * Stop the underlying `SparkContext` if there are are no active sessions remaining.
    *
    * @since 2.0.0
    */
   def stop(): Unit = {
-    sparkContext.stop()
+    SparkSession.clearActiveSession()
+    if (SparkSession.numActiveSessions.get() == 0) {
+      sparkContext.stop()
+    }
   }
 
   /**
@@ -776,6 +779,8 @@ class SparkSession private(
 @Stable
 object SparkSession extends Logging {
 
+  private[spark] val numActiveSessions: AtomicInteger = new AtomicInteger(0)
+
   /**
    * Builder for [[SparkSession]].
    */
@@ -958,6 +963,8 @@ object SparkSession extends Logging {
         sparkContext.addSparkListener(new SparkListener {
           override def onApplicationEnd(applicationEnd: SparkListenerApplicationEnd): Unit = {
             defaultSession.set(null)
+            // Should remove listener after this event fires
+            sparkContext.removeSparkListener(this)
           }
         })
       }
@@ -981,17 +988,30 @@ object SparkSession extends Logging {
    * @since 2.0.0
    */
   def setActiveSession(session: SparkSession): Unit = {
-    activeThreadSession.set(session)
+    if (getActiveSession.isEmpty
+      || (session != getActiveSession.get && getActiveSession.isDefined)) {
+      numActiveSessions.getAndIncrement
+      activeThreadSession.set(session)
+    } else if (session == null) {
+      this.clearActiveSession()
+    }
   }
 
   /**
-   * Clears the active SparkSession for current thread. Subsequent calls to getOrCreate will
-   * return the first created context instead of a thread-local override.
+   * Clears the active SparkSession for current thread assuming it is defined.
+   * Subsequent calls to getOrCreate will return the first created context
+   * instead of a thread-local override.
    *
    * @since 2.0.0
    */
   def clearActiveSession(): Unit = {
-    activeThreadSession.remove()
+    if (getActiveSession.isDefined) {
+      activeThreadSession.remove()
+      numActiveSessions.decrementAndGet()
+    } else {
+      logWarning("Calling clearActiveSession() on a SparkSession " +
+        "without an active session is a noop.")
+    }
   }
 
   /**
@@ -1004,12 +1024,17 @@ object SparkSession extends Logging {
   }
 
   /**
-   * Clears the default SparkSession that is returned by the builder.
-   *
+   * Clears the default SparkSession that is returned by the builder
+   * if it is not null.
    * @since 2.0.0
    */
   def clearDefaultSession(): Unit = {
-    defaultSession.set(null)
+    if (getDefaultSession.isDefined) {
+      defaultSession.set(null)
+    } else {
+      logWarning("Calling clearDefaultSession() on a SparkSession " +
+        "without an default session is a noop.")
+    }
   }
 
   /**

sql/core/src/test/scala/org/apache/spark/sql/SparkSessionLifecycleSuite.scala

@@ -0,0 +1,66 @@
+/*
+ * 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.spark.sql
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.internal.config.UI.UI_ENABLED
+
+
+/**
+ * Test cases for the lifecycle of a [[SparkSession]].
+ */
+class SparkSessionLifecycleSuite extends SparkFunSuite {
+  test("test SparkContext stopped when last SparkSession is stopped ") {
+    val session1 = SparkSession.builder()
+      .master("local")
+      .config(UI_ENABLED.key, value = false)
+      .config("some-config", "a")
+      .getOrCreate()
+
+    assert(!session1.sparkContext.isStopped)
+
+    val session2 = SparkSession.builder()
+      .master("local")
+      .config(UI_ENABLED.key, value = false)
+      .config("some-config", "b")
+      .getOrCreate()
+
+    session1.stop()
+    session2.stop()
+    assert(session1.sparkContext.isStopped)
+  }
+
+  test("test SparkContext is not stopped when other sessions exist") {
+    val session1 = SparkSession.builder()
+      .master("local")
+      .config(UI_ENABLED.key, value = false)
+      .config("some-config", "a")
+      .getOrCreate()
+
+    assert(!session1.sparkContext.isStopped)
+
+    val session2 = SparkSession.builder()
+      .master("local")
+      .config(UI_ENABLED.key, value = false)
+      .config("some-config", "b")
+      .getOrCreate()
+
+    session1.stop()
+    assert(!session1.sparkContext.isStopped)
+  }
+}