adding test and cleanup

Author: vinooganesh

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

@@ -939,19 +939,12 @@ object SparkSession extends Logging {
         options.foreach { case (k, v) => session.initialSessionOptions.put(k, v) }
         setDefaultSession(session)
         setActiveSession(session)
-        registerSessionListenerOnContext(sparkContext)
+        registerContextListener(sparkContext)
       }
 
       return session
     }
 
-    private def registerSessionListenerOnContext(sparkContext: SparkContext): Unit = {
-      if (!SparkSession.sessionListenerRegistered.get()) {
-        sparkContext.addSparkListener(_sessionListener)
-        SparkSession.sessionListenerRegistered.set(true)
-      }
-    }
-
     private def applyModifiableSettings(session: SparkSession): Unit = {
       val (staticConfs, otherConfs) =
         options.partition(kv => SQLConf.staticConfKeys.contains(kv._1))
@@ -1061,16 +1054,21 @@ object SparkSession extends Logging {
   ////////////////////////////////////////////////////////////////////////////////////////
   // Private methods from now on
   ////////////////////////////////////////////////////////////////////////////////////////
-  /** Default listener on SparkContext  */
-  private val _sessionListener: SparkListener = new SparkListener {
-    override def onApplicationEnd(applicationEnd: SparkListenerApplicationEnd): Unit = {
-      defaultSession.set(null)
+
+  private val listenerRegistered: AtomicBoolean = new AtomicBoolean(false)
+
+  /** Register the AppEnd listener onto the Context  */
+  private def registerContextListener(sparkContext: SparkContext): Unit = {
+    if (!SparkSession.listenerRegistered.get()) {
+      sparkContext.addSparkListener(new SparkListener {
+        override def onApplicationEnd(applicationEnd: SparkListenerApplicationEnd): Unit = {
+          defaultSession.set(null)
+        }
+      })
+      SparkSession.listenerRegistered.set(true)
     }
   }
 
-  /** Whether the app end listener has been registered on the context */
-  private val sessionListenerRegistered: AtomicBoolean = new AtomicBoolean(false)
-
   /** The active SparkSession for the current thread. */
   private val activeThreadSession = new InheritableThreadLocal[SparkSession]
 
@@ -1087,8 +1085,6 @@ object SparkSession extends Logging {
     }
   }
 
-  private[spark] def isSessionListenerRegistered: Boolean = sessionListenerRegistered.get
-
   private def assertOnDriver(): Unit = {
     if (Utils.isTesting && TaskContext.get != null) {
       // we're accessing it during task execution, fail.

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

@@ -169,6 +169,31 @@ class SparkSessionBuilderSuite extends SparkFunSuite with BeforeAndAfterEach {
     assert(session.sessionState.conf.getConf(GLOBAL_TEMP_DATABASE) === "globaltempdb-spark-31234")
   }
 
+  test("SPARK-31354: SparkContext only register one SparkSession ApplicationEnd listener") {
+    val conf = new SparkConf()
+      .setMaster("local")
+      .setAppName("test-app-SPARK-31354-1")
+    val context = new SparkContext(conf)
+    SparkSession
+      .builder()
+      .sparkContext(context)
+      .master("local")
+      .getOrCreate()
+    val postFirstCreation = context.listenerBus.listeners.size()
+    SparkSession.clearActiveSession()
+    SparkSession.clearDefaultSession()
+
+    SparkSession
+      .builder()
+      .sparkContext(context)
+      .master("local")
+      .getOrCreate()
+    val postSecondCreation = context.listenerBus.listeners.size()
+    SparkSession.clearActiveSession()
+    SparkSession.clearDefaultSession()
+    assert(postFirstCreation == postSecondCreation)
+  }
+
   test("SPARK-31532: should not propagate static sql configs to the existing" +
     " active/default SparkSession") {
     val session = SparkSession.builder()