diff --git a/core/src/test/scala/org/apache/spark/TempLocalSparkContext.scala b/core/src/test/scala/org/apache/spark/TempLocalSparkContext.scala new file mode 100644 index 0000000000000..6d5fcd1edfb03 --- /dev/null +++ b/core/src/test/scala/org/apache/spark/TempLocalSparkContext.scala @@ -0,0 +1,100 @@ +/* + * 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 + +import _root_.io.netty.util.internal.logging.{InternalLoggerFactory, Slf4JLoggerFactory} +import org.scalatest.BeforeAndAfterAll +import org.scalatest.BeforeAndAfterEach +import org.scalatest.Suite + +import org.apache.spark.internal.Logging +import org.apache.spark.resource.ResourceProfile + +/** + * Manages a local `sc` `SparkContext` variable, correctly stopping it after each test. + * + * Note: this class is a copy of [[LocalSparkContext]]. Why copy it? Reduce conflict. Because + * many test suites use [[LocalSparkContext]] and overwrite some variable or function (e.g. + * sc of LocalSparkContext), there occurs conflict when we refactor the `sc` as a new function. + * After migrating all test suites that use [[LocalSparkContext]] to use + * [[TempLocalSparkContext]], we will delete the original [[LocalSparkContext]] and rename + * [[TempLocalSparkContext]] to [[LocalSparkContext]]. + */ +trait TempLocalSparkContext extends BeforeAndAfterEach + with BeforeAndAfterAll with Logging { self: Suite => + + private var _conf: SparkConf = defaultSparkConf + + @transient private var _sc: SparkContext = _ + + def conf: SparkConf = _conf + + /** + * Currently, we are focusing on the reconstruction of LocalSparkContext, so this method + * was created temporarily. When the migration work is completed, this method will be + * renamed to `sc` and the variable `sc` will be deleted. + */ + def sc: SparkContext = { + if (_sc == null) { + _sc = new SparkContext(_conf) + } + _sc + } + + override def beforeAll(): Unit = { + super.beforeAll() + InternalLoggerFactory.setDefaultFactory(Slf4JLoggerFactory.INSTANCE) + } + + override def afterEach(): Unit = { + try { + resetSparkContext() + } finally { + super.afterEach() + } + } + + def resetSparkContext(): Unit = { + TempLocalSparkContext.stop(_sc) + ResourceProfile.clearDefaultProfile() + _sc = null + _conf = defaultSparkConf + } + + private def defaultSparkConf: SparkConf = new SparkConf() + .setMaster("local[2]").setAppName(s"${this.getClass.getSimpleName}") +} + +object TempLocalSparkContext { + def stop(sc: SparkContext): Unit = { + if (sc != null) { + sc.stop() + } + // To avoid RPC rebinding to the same port, since it doesn't unbind immediately on shutdown + System.clearProperty("spark.driver.port") + } + + /** Runs `f` by passing in `sc` and ensures that `sc` is stopped. */ + def withSpark[T](sc: SparkContext)(f: SparkContext => T): T = { + try { + f(sc) + } finally { + stop(sc) + } + } +} diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index 2b38aa1cde6cd..c9a36299a7ea4 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -19,7 +19,7 @@ package org.apache.spark.scheduler import java.util.Properties import java.util.concurrent.{CountDownLatch, TimeUnit} -import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger, AtomicLong, AtomicReference} +import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong, AtomicReference} import scala.annotation.meta.param import scala.collection.mutable.{ArrayBuffer, HashMap, HashSet, Map} @@ -125,14 +125,14 @@ class MyRDD( class DAGSchedulerSuiteDummyException extends Exception -class DAGSchedulerSuite extends SparkFunSuite with LocalSparkContext with TimeLimits { +class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with TimeLimits { import DAGSchedulerSuite._ // Necessary to make ScalaTest 3.x interrupt a thread on the JVM like ScalaTest 2.2.x implicit val defaultSignaler: Signaler = ThreadSignaler - val conf = new SparkConf + private var firstInit: Boolean = _ /** Set of TaskSets the DAGScheduler has requested executed. */ val taskSets = scala.collection.mutable.Buffer[TaskSet]() @@ -295,11 +295,19 @@ class DAGSchedulerSuite extends SparkFunSuite with LocalSparkContext with TimeLi override def beforeEach(): Unit = { super.beforeEach() - init(new SparkConf()) + firstInit = true } - private def init(testConf: SparkConf): Unit = { - sc = new SparkContext("local[2]", "DAGSchedulerSuite", testConf) + override def sc: SparkContext = { + val sc = super.sc + if (firstInit) { + init(sc) + firstInit = false + } + sc + } + + private def init(sc: SparkContext): Unit = { sparkListener = new EventInfoRecordingListener failure = null sc.addSparkListener(sparkListener) @@ -308,10 +316,10 @@ class DAGSchedulerSuite extends SparkFunSuite with LocalSparkContext with TimeLi cancelledStages.clear() cacheLocations.clear() results.clear() - securityMgr = new SecurityManager(conf) - broadcastManager = new BroadcastManager(true, conf, securityMgr) - mapOutputTracker = spy(new MyMapOutputTrackerMaster(conf, broadcastManager)) - blockManagerMaster = spy(new MyBlockManagerMaster(conf)) + securityMgr = new SecurityManager(sc.getConf) + broadcastManager = new BroadcastManager(true, sc.getConf, securityMgr) + mapOutputTracker = spy(new MyMapOutputTrackerMaster(sc.getConf, broadcastManager)) + blockManagerMaster = spy(new MyBlockManagerMaster(sc.getConf)) scheduler = new DAGScheduler( sc, taskScheduler, @@ -351,6 +359,8 @@ class DAGSchedulerSuite extends SparkFunSuite with LocalSparkContext with TimeLi * DAGScheduler event loop. */ private def runEvent(event: DAGSchedulerEvent): Unit = { + // Ensure the initialization of various components + sc dagEventProcessLoopTester.post(event) } @@ -489,12 +499,8 @@ class DAGSchedulerSuite extends SparkFunSuite with LocalSparkContext with TimeLi } test("All shuffle files on the storage endpoint should be cleaned up when it is lost") { - // reset the test context with the right shuffle service config - afterEach() - val conf = new SparkConf() conf.set(config.SHUFFLE_SERVICE_ENABLED.key, "true") conf.set("spark.files.fetchFailure.unRegisterOutputOnHost", "true") - init(conf) runEvent(ExecutorAdded("hostA-exec1", "hostA")) runEvent(ExecutorAdded("hostA-exec2", "hostA")) runEvent(ExecutorAdded("hostB-exec", "hostB")) @@ -563,11 +569,7 @@ class DAGSchedulerSuite extends SparkFunSuite with LocalSparkContext with TimeLi } test("SPARK-32003: All shuffle files for executor should be cleaned up on fetch failure") { - // reset the test context with the right shuffle service config - afterEach() - val conf = new SparkConf() conf.set(config.SHUFFLE_SERVICE_ENABLED.key, "true") - init(conf) val shuffleMapRdd = new MyRDD(sc, 3, Nil) val shuffleDep = new ShuffleDependency(shuffleMapRdd, new HashPartitioner(3)) @@ -857,11 +859,7 @@ class DAGSchedulerSuite extends SparkFunSuite with LocalSparkContext with TimeLi "not lost" } test(s"shuffle files $maybeLost when $eventDescription") { - // reset the test context with the right shuffle service config - afterEach() - val conf = new SparkConf() conf.set(config.SHUFFLE_SERVICE_ENABLED.key, shuffleServiceOn.toString) - init(conf) assert(sc.env.blockManager.externalShuffleServiceEnabled == shuffleServiceOn) val shuffleMapRdd = new MyRDD(sc, 2, Nil) @@ -2876,11 +2874,7 @@ class DAGSchedulerSuite extends SparkFunSuite with LocalSparkContext with TimeLi } test("SPARK-25341: abort stage while using old fetch protocol") { - // reset the test context with using old fetch protocol - afterEach() - val conf = new SparkConf() conf.set(config.SHUFFLE_USE_OLD_FETCH_PROTOCOL.key, "true") - init(conf) // Construct the scenario of indeterminate stage fetch failed. constructIndeterminateStageFetchFailed() // The job should fail because Spark can't rollback the shuffle map stage while @@ -3208,10 +3202,7 @@ class DAGSchedulerSuite extends SparkFunSuite with LocalSparkContext with TimeLi } test("test 2 resource profile with merge conflict config true") { - afterEach() - val conf = new SparkConf() conf.set(config.RESOURCE_PROFILE_MERGE_CONFLICTS.key, "true") - init(conf) val ereqs = new ExecutorResourceRequests().cores(4) val treqs = new TaskResourceRequests().cpus(1) @@ -3229,10 +3220,7 @@ class DAGSchedulerSuite extends SparkFunSuite with LocalSparkContext with TimeLi } test("test multiple resource profiles created from merging use same rp") { - afterEach() - val conf = new SparkConf() conf.set(config.RESOURCE_PROFILE_MERGE_CONFLICTS.key, "true") - init(conf) val ereqs = new ExecutorResourceRequests().cores(4) val treqs = new TaskResourceRequests().cpus(1) @@ -3325,10 +3313,7 @@ class DAGSchedulerSuite extends SparkFunSuite with LocalSparkContext with TimeLi } test("test merge 3 resource profiles") { - afterEach() - val conf = new SparkConf() conf.set(config.RESOURCE_PROFILE_MERGE_CONFLICTS.key, "true") - init(conf) val ereqs = new ExecutorResourceRequests().cores(4) val treqs = new TaskResourceRequests().cpus(1) val rp1 = new ResourceProfile(ereqs.requests, treqs.requests)