Merge pull request datastax#976 from datastax/SPARKC-377-master

SPARKC-377 shutdown cluster and client connection sequentially

.travis.yml

@@ -7,6 +7,16 @@ scala:
   - 2.10.5
   - 2.11.7
 
+cache:
+  directories:
+    - $HOME/.ivy2/cache
+    - $HOME/.sbt/boot/
+
+before_cache:
+  # Tricks to avoid unnecessary cache updates
+  - find $HOME/.ivy2 -name "ivydata-*.properties" -delete
+  - find $HOME/.sbt -name "*.lock" -delete
+
 script:
   - "sbt ++$TRAVIS_SCALA_VERSION -Dtravis=true test"
   - "sbt ++$TRAVIS_SCALA_VERSION -Dtravis=true it:test"

project/Settings.scala

@@ -52,7 +52,7 @@ object Settings extends Build {
   val osmxBean = ManagementFactory.getOperatingSystemMXBean.asInstanceOf[OperatingSystemMXBean]
   val sysMemoryInMB = osmxBean.getTotalPhysicalMemorySize >> 20
   val singleRunRequiredMem = 3 * 1024 + 512
-  val parallelTasks = if (isTravis) 2 else Math.max(1, ((sysMemoryInMB - 1550) / singleRunRequiredMem).toInt)
+  val parallelTasks = if (isTravis) 1 else Math.max(1, ((sysMemoryInMB - 1550) / singleRunRequiredMem).toInt)
 
   // Due to lack of entrophy on virtual machines we want to use /dev/urandom instead of /dev/random
   val useURandom = Files.exists(Paths.get("/dev/urandom"))
@@ -67,7 +67,6 @@ object Settings extends Build {
   val cassandraTestVersion = sys.props.get("test.cassandra.version").getOrElse(Versions.Cassandra)
 
   lazy val TEST_JAVA_OPTS = Seq(
-    "-Xms256m",
     "-Xmx512m",
     s"-Dtest.cassandra.version=$cassandraTestVersion",
     "-Dsun.io.serialization.extendedDebugInfo=true",

spark-cassandra-connector-embedded/src/main/scala/com/datastax/spark/connector/embedded/EmbeddedCassandra.scala

@@ -137,13 +137,20 @@ object EmbeddedCassandra {
     }
   }
 
-  Runtime.getRuntime.addShutdownHook(new Thread(new Runnable {
+  private val shutdownThread: Thread = new Thread("Shutdown embedded C* hook thread") {
     override def run() = {
-      cassandraRunners.flatten.foreach(_.destroy())
-      release()
+      shutdown()
     }
-  }))
+  }
 
+  Runtime.getRuntime.addShutdownHook(shutdownThread)
 
+  private[connector] def shutdown(): Unit = {
+    cassandraRunners.flatten.foreach(_.destroy())
+    release()
+  }
 
+  private[connector] def removeShutdownHook(): Boolean = {
+    Runtime.getRuntime.removeShutdownHook(shutdownThread)
+  }
 }

spark-cassandra-connector/src/it/scala/com/datastax/spark/connector/SparkCassandraITFlatSpecBase.scala

@@ -9,8 +9,9 @@ import org.apache.commons.lang3.StringUtils
 import org.scalatest._
 
 import com.datastax.driver.core.Session
-import com.datastax.spark.connector.embedded.SparkTemplate
+import com.datastax.spark.connector.embedded.{EmbeddedCassandra, SparkTemplate}
 import com.datastax.spark.connector.testkit.{AbstractSpec, SharedEmbeddedCassandra}
+import com.datastax.spark.connector.util.SerialShutdownHooks
 
 
 trait SparkCassandraITFlatSpecBase extends FlatSpec with SparkCassandraITSpecBase
@@ -57,4 +58,11 @@ trait SparkCassandraITSpecBase extends Suite with Matchers with SharedEmbeddedCa
 object SparkCassandraITSpecBase {
   val executor = Executors.newFixedThreadPool(100)
   val ec = ExecutionContext.fromExecutor(executor)
+
+  EmbeddedCassandra.removeShutdownHook
+  // now embedded C* won't shutdown itself, let's do it in serial fashion
+  SerialShutdownHooks.add("Shutting down all Cassandra runners")(() => {
+    EmbeddedCassandra.shutdown
+  })
+
 }

spark-cassandra-connector/src/main/scala/com/datastax/spark/connector/cql/CassandraConnector.scala

@@ -10,7 +10,7 @@ import org.apache.spark.{Logging, SparkConf}
 
 import com.datastax.driver.core.{Cluster, Host, Session}
 import com.datastax.spark.connector.cql.CassandraConnectorConf.CassandraSSLConf
-
+import com.datastax.spark.connector.util.SerialShutdownHooks
 
 /** Provides and manages connections to Cassandra.
   *
@@ -179,11 +179,9 @@ object CassandraConnector extends Logging {
     hosts.map(h => conf.copy(hosts = Set(h.getAddress))) + conf.copy(hosts = hosts.map(_.getAddress))
   }
 
-  Runtime.getRuntime.addShutdownHook(new Thread(new Runnable {
-    def run() {
-      sessionCache.shutdown()
-    }
-  }))
+  SerialShutdownHooks.add("Clearing session cache for C* connector")(() => {
+    sessionCache.shutdown()
+  })
 
   /** Returns a CassandraConnector created from properties found in the [[org.apache.spark.SparkConf SparkConf]] object */
   def apply(conf: SparkConf): CassandraConnector = {

spark-cassandra-connector/src/main/scala/com/datastax/spark/connector/util/SerialShutdownHooks.scala

@@ -0,0 +1,37 @@
+package com.datastax.spark.connector.util
+
+import scala.collection.mutable
+
+import org.apache.spark.Logging
+
+private[connector] object SerialShutdownHooks extends Logging {
+
+  private val hooks = mutable.Map[String, () => Unit]()
+  @volatile private var isShuttingDown = false
+
+  def add(name: String)(body: () => Unit): Unit = SerialShutdownHooks.synchronized {
+    if (isShuttingDown) {
+      logError(s"Adding shutdown hook ($name) during shutting down is not allowed.")
+    } else {
+      hooks.put(name, body)
+    }
+  }
+
+  Runtime.getRuntime.addShutdownHook(new Thread("Serial shutdown hooks thread") {
+    override def run(): Unit = {
+      SerialShutdownHooks.synchronized {
+        isShuttingDown = true
+      }
+      for ((name, task) <- hooks) {
+        try {
+          logDebug(s"Running shutdown hook: $name")
+          task()
+          logInfo(s"Successfully executed shutdown hook: $name")
+        } catch {
+          case exc: Throwable =>
+            logError(s"Shutdown hook ($name) failed", exc)
+        }
+      }
+    }
+  })
+}