[MINOR] Reference-count shared Spark context to prevent concurrent shutdown

src/main/java/org/apache/sysds/api/DMLScript.java

@@ -508,6 +508,11 @@ private static void execute(String dmlScriptStr, String fnameOptConfig, Map<Stri
 		ExecutionContext ec = null;
 		try {
 			ec = ExecutionContextFactory.createContext(rtprog);
+			//register as an active user of the shared spark context; balanced by
+			//exitSparkExecution() in the finally block so a concurrent execution
+			//cannot stop the context mid-job
+			if(ec instanceof SparkExecutionContext)
+				SparkExecutionContext.enterSparkExecution();
 			ScriptExecutorUtils.executeRuntimeProgram(rtprog, ec, ConfigurationManager.getDMLConfig(), STATISTICS ? STATISTICS_COUNT : 0, null);
 		}
 		finally {
@@ -516,8 +521,12 @@ private static void execute(String dmlScriptStr, String fnameOptConfig, Map<Stri
 			FederatedData.clearWorkGroup();
 			//stop spark context (after cleanup of federated workers and other pools,
 			//otherwise federated spark cleanups in local tests throw errors in same JVM)
-			if(ec != null && ec instanceof SparkExecutionContext)
+			if(ec != null && ec instanceof SparkExecutionContext) {
+				//release our registration, then stop the context if no other
+				//execution is still using it
+				SparkExecutionContext.exitSparkExecution();
 				((SparkExecutionContext) ec).close();
+			}
 			LOG.info("END DML run " + getDateTime() );
 		}
 	}

src/main/java/org/apache/sysds/runtime/controlprogram/context/SparkExecutionContext.java

@@ -122,6 +122,11 @@ public class SparkExecutionContext extends ExecutionContext
 	//singleton spark context (as there can be only one spark context per JVM)
 	private static JavaSparkContext _spctx = null;
 
+	//registered users of the singleton context (guarded by the
+	//SparkExecutionContext.class monitor); maintained by enterSparkExecution()/
+	//exitSparkExecution(), and close() only stops the context once it hits zero
+	private static int _activeExecutions = 0;
+
 	//registry of parallelized RDDs to enforce that at any time, we spent at most
 	//10% of JVM max heap size for parallelized RDDs; if this is not sufficient,
 	//matrices or frames are exported to HDFS and the RDDs are created from files.
@@ -170,6 +175,9 @@ public synchronized static JavaSparkContext getSparkContextStatic() {
 		initSparkContext();
 		if(_spctx.sc().isStopped()){
 			_spctx = null;
+			//the previous context was stopped; reset the active-execution count so a
+			//stale registration cannot skip a future legitimate stop of the new one
+			_activeExecutions = 0;
 			initSparkContext();
 		}
 		return _spctx;
@@ -186,11 +194,48 @@ public synchronized static boolean isSparkContextCreated() {
 	}
 
 	public static void resetSparkContextStatic() {
-		_spctx = null;
+		synchronized(SparkExecutionContext.class) {
+			_spctx = null;
+			//force-discarding the shared context: drop the active-execution count so
+			//a stale registration cannot skip a future legitimate stop
+			_activeExecutions = 0;
+		}
+	}
+
+	/**
+	 * Registers an active user of the shared spark context. Must be balanced by a
+	 * later {@link #exitSparkExecution()} so a concurrent execution cannot stop the
+	 * context while this one still has in-flight jobs.
+	 */
+	public static void enterSparkExecution() {
+		synchronized(SparkExecutionContext.class) {
+			_activeExecutions++;
+		}
+	}
+
+	/**
+	 * Releases an active user previously registered via {@link #enterSparkExecution()}.
+	 * Only adjusts the count; the actual teardown is left to {@link #close()}, which
+	 * stops the context once no registered execution remains.
+	 */
+	public static void exitSparkExecution() {
+		synchronized(SparkExecutionContext.class) {
+			if(_activeExecutions > 0)
+				_activeExecutions--;
+		}
 	}
 
 	public void close() {
-		synchronized( SparkExecutionContext.class) {
+		synchronized(SparkExecutionContext.class) {
+			//keep the shared context alive while a registered execution still uses
+			//it; close() never changes the count, so an unpaired close() (a caller
+			//that never entered) cannot stop a context another execution is using
+			if(_activeExecutions > 0) {
+				if(LOG.isDebugEnabled())
+					LOG.debug("Keeping shared spark context alive; " + _activeExecutions
+						+ " execution(s) still active");
+				return;
+			}
 			if(_spctx != null) {
 				Logger spL = Logger.getLogger("org.apache.spark.network.client.TransportResponseHandler");
 				spL.setLevel(Level.FATAL);

src/test/java/org/apache/sysds/test/component/context/SparkContextReferenceCountTest.java

@@ -0,0 +1,126 @@
+/*
+ * 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.sysds.test.component.context;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Arrays;
+
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.sysds.runtime.controlprogram.context.ExecutionContextFactory;
+import org.apache.sysds.runtime.controlprogram.context.SparkExecutionContext;
+import org.junit.Test;
+
+@net.jcip.annotations.NotThreadSafe
+public class SparkContextReferenceCountTest {
+
+	/**
+	 * Two DML executions sharing the JVM-wide singleton spark context (as happens
+	 * with surefire parallel tests, threadCount>1). When the first execution
+	 * finishes and calls close(), the shared context must stay alive because the
+	 * second execution still has in-flight work. Before reference counting,
+	 * close() stopped the context unconditionally, which cancelled the second
+	 * execution's spark job and wedged it until the test watchdog.
+	 */
+	@Test
+	public void closeKeepsContextAliveWhileAnotherExecutionIsActive() {
+		SparkExecutionContext ecA = null;
+		SparkExecutionContext ecB = null;
+		try {
+			// execution A: create the context then register (as in DMLScript.execute)
+			ecA = ExecutionContextFactory.createSparkExecutionContext();
+			SparkExecutionContext.enterSparkExecution();
+			JavaSparkContext sc = ecA.getSparkContext();
+
+			// execution B is a second concurrent user with its own context instance,
+			// sharing the same JVM-wide singleton spark context
+			ecB = ExecutionContextFactory.createSparkExecutionContext();
+			SparkExecutionContext.enterSparkExecution();
+
+			// B's in-flight work
+			JavaRDD<Integer> rdd = sc.parallelize(Arrays.asList(1, 2, 3, 4));
+
+			// A finishes first: release its registration; close() must NOT stop the
+			// context that B still uses
+			SparkExecutionContext.exitSparkExecution();
+			ecA.close();
+			assertFalse("shared context must stay alive while another execution is active",
+				sc.sc().isStopped());
+			assertEquals("B's job must still run on the live context",
+				10L, rdd.reduce(Integer::sum).longValue());
+
+			// B finishes last: releasing the final registration lets close() stop it
+			SparkExecutionContext.exitSparkExecution();
+			ecB.close();
+			assertTrue("shared context must be stopped once the last execution closes",
+				sc.sc().isStopped());
+		}
+		finally {
+			// drain any remaining registrations and stop the context so a failed
+			// assertion cannot leak ref-count state into other tests in this JVM
+			// (exit/close are clamped and no-op once already drained/stopped)
+			SparkExecutionContext.exitSparkExecution();
+			SparkExecutionContext.exitSparkExecution();
+			if(ecA != null)
+				ecA.close();
+		}
+	}
+
+	/**
+	 * An unpaired close() (a caller that borrows the shared context but never
+	 * registered via enterSparkExecution()) must not stop a context another
+	 * execution still uses. This fails on the old unconditional-stop code, which
+	 * tore the context down out from under the active execution.
+	 */
+	@Test
+	public void unpairedCloseDoesNotStopAContextStillInUse() {
+		SparkExecutionContext active = null;
+		SparkExecutionContext unregistered = null;
+		try {
+			// a registered, in-flight execution holds the shared context
+			active = ExecutionContextFactory.createSparkExecutionContext();
+			SparkExecutionContext.enterSparkExecution();
+			JavaSparkContext sc = active.getSparkContext();
+
+			// a context that never registered closes (e.g. a caller that only
+			// borrows the shared context): close() must not stop a context in use
+			unregistered = ExecutionContextFactory.createSparkExecutionContext();
+			unregistered.close();
+			assertFalse("unpaired close() must not stop a context still in use",
+				sc.sc().isStopped());
+
+			// the registered execution finishing stops the context as the last user
+			SparkExecutionContext.exitSparkExecution();
+			active.close();
+			assertTrue("context must stop once the last registered execution closes",
+				sc.sc().isStopped());
+		}
+		finally {
+			SparkExecutionContext.exitSparkExecution();
+			if(active != null)
+				active.close();
+			if(unregistered != null)
+				unregistered.close();
+		}
+	}
+}