Added getState and exposed StreamingContextState

Author: tdas

streaming/src/main/scala/org/apache/spark/streaming/StreamingContext.scala

@@ -35,6 +35,7 @@ import org.apache.spark.annotation.Experimental
 import org.apache.spark.input.FixedLengthBinaryInputFormat
 import org.apache.spark.rdd.RDD
 import org.apache.spark.storage.StorageLevel
+import org.apache.spark.streaming.StreamingContextState._
 import org.apache.spark.streaming.dstream._
 import org.apache.spark.streaming.receiver.{ActorReceiver, ActorSupervisorStrategy, Receiver}
 import org.apache.spark.streaming.scheduler._
@@ -193,14 +194,7 @@ class StreamingContext private[streaming] (
   assert(env.metricsSystem != null)
   env.metricsSystem.registerSource(streamingSource)
 
-  /** Enumeration to identify current state of the StreamingContext */
-  private[streaming] object StreamingContextState extends Enumeration {
-    type CheckpointState = Value
-    val Initialized, Started, Stopped = Value
-  }
-
-  import StreamingContextState._
-  private[streaming] var state = Initialized
+  private var state: StreamingContextState = INITIALIZED
 
   /**
    * Return the associated Spark context
@@ -512,23 +506,32 @@ class StreamingContext private[streaming] (
     )
   }
 
+  /**
+   * Return the current state of the context.
+   */
+  def getState(): StreamingContextState = {
+    state
+  }
+
   /**
    * Start the execution of the streams.
    *
    * @throws SparkException if the context has already been started or stopped.
    */
   def start(): Unit = synchronized {
-    if (state == Started) {
-      throw new SparkException("StreamingContext has already been started")
-    }
-    if (state == Stopped) {
-      throw new SparkException("StreamingContext has already been stopped")
+    state match {
+      case INITIALIZED =>
+        // good to start
+      case STARTED =>
+        throw new SparkException("StreamingContext has already been started")
+      case STOPPED =>
+        throw new SparkException("StreamingContext has already been stopped")
     }
     validate()
     sparkContext.setCallSite(DStream.getCreationSite())
     scheduler.start()
     uiTab.foreach(_.attach())
-    state = Started
+    state = StreamingContextState.STARTED
   }
 
   /**
@@ -585,9 +588,11 @@ class StreamingContext private[streaming] (
    */
   def stop(stopSparkContext: Boolean, stopGracefully: Boolean): Unit = synchronized {
     state match {
-      case Initialized => logWarning("StreamingContext has not been started yet")
-      case Stopped => logWarning("StreamingContext has already been stopped")
-      case Started =>
+      case INITIALIZED =>
+        logWarning("StreamingContext has not been started yet")
+      case STOPPED =>
+        logWarning("StreamingContext has already been stopped")
+      case STARTED =>
         scheduler.stop(stopGracefully)
         logInfo("StreamingContext stopped successfully")
         waiter.notifyStop()
@@ -598,7 +603,7 @@ class StreamingContext private[streaming] (
     if (stopSparkContext) sc.stop()
     uiTab.foreach(_.detach())
     // The state should always be Stopped after calling `stop()`, even if we haven't started yet:
-    state = Stopped
+    state = STOPPED
   }
 }
 

streaming/src/main/scala/org/apache/spark/streaming/StreamingContextState.scala

@@ -0,0 +1,35 @@
+package org.apache.spark.streaming
+
+import org.apache.spark.annotation.DeveloperApi
+
+/**
+ * :: DeveloperApi ::
+ *
+ * Represents the state of the StreamingContext.
+ */
+@DeveloperApi
+class StreamingContextState private (enumValue: Int) {
+
+  override def hashCode: Int = enumValue
+
+  override def equals(other: Any): Boolean = {
+    other match {
+      case otherState: StreamingContextState =>
+        otherState.hashCode == this.hashCode
+      case _ =>
+        false
+    }
+  }
+}
+
+/**
+ * :: DeveloperApi ::
+ *
+ * Object enumerating all the states that a StreamingContext can be.
+ */
+@DeveloperApi
+object StreamingContextState {
+  val INITIALIZED = new StreamingContextState(0)
+  val STARTED = new StreamingContextState(1)
+  val STOPPED = new StreamingContextState(2)
+}

streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingContext.scala

@@ -578,6 +578,11 @@ class JavaStreamingContext(val ssc: StreamingContext) extends Closeable {
     ssc.addStreamingListener(streamingListener)
   }
 
+
+  def getState(): StreamingContextState = {
+    ssc.getState()
+  }
+
   /**
    * Start the execution of the streams.
    */

streaming/src/test/scala/org/apache/spark/streaming/StreamingContextSuite.scala

@@ -109,38 +109,49 @@ class StreamingContextSuite extends FunSuite with BeforeAndAfter with Timeouts w
     assert(ssc.conf.getTimeAsSeconds("spark.cleaner.ttl", "-1") === 10)
   }
 
+  test("state matching") {
+    import StreamingContextState._
+    assert(INITIALIZED === INITIALIZED)
+    assert(INITIALIZED != STARTED)
+  }
+
   test("start and stop state check") {
     ssc = new StreamingContext(master, appName, batchDuration)
     addInputStream(ssc).register()
 
-    assert(ssc.state === ssc.StreamingContextState.Initialized)
+    assert(ssc.getState() === StreamingContextState.INITIALIZED)
     ssc.start()
-    assert(ssc.state === ssc.StreamingContextState.Started)
+    assert(ssc.getState() === StreamingContextState.STARTED)
     ssc.stop()
-    assert(ssc.state === ssc.StreamingContextState.Stopped)
+    assert(ssc.getState() === StreamingContextState.STOPPED)
   }
 
   test("start multiple times") {
     ssc = new StreamingContext(master, appName, batchDuration)
     addInputStream(ssc).register()
     ssc.start()
+    assert(ssc.getState() === StreamingContextState.STARTED)
     intercept[SparkException] {
       ssc.start()
     }
+    assert(ssc.getState() === StreamingContextState.STARTED)
   }
 
   test("stop multiple times") {
     ssc = new StreamingContext(master, appName, batchDuration)
     addInputStream(ssc).register()
     ssc.start()
     ssc.stop()
+    assert(ssc.getState() === StreamingContextState.STOPPED)
     ssc.stop()
+    assert(ssc.getState() === StreamingContextState.STOPPED)
   }
 
   test("stop before start") {
     ssc = new StreamingContext(master, appName, batchDuration)
     addInputStream(ssc).register()
     ssc.stop()  // stop before start should not throw exception
+    assert(ssc.getState() === StreamingContextState.STOPPED)
   }
 
   test("start after stop") {
@@ -151,6 +162,7 @@ class StreamingContextSuite extends FunSuite with BeforeAndAfter with Timeouts w
     intercept[SparkException] {
       ssc.start() // start after stop should throw exception
     }
+    assert(ssc.getState() === StreamingContextState.STOPPED)
   }
 
   test("stop only streaming context") {
@@ -159,6 +171,7 @@ class StreamingContextSuite extends FunSuite with BeforeAndAfter with Timeouts w
     addInputStream(ssc).register()
     ssc.start()
     ssc.stop(stopSparkContext = false)
+    assert(ssc.getState() === StreamingContextState.STOPPED)
     assert(sc.makeRDD(1 to 100).collect().size === 100)
     ssc = new StreamingContext(sc, batchDuration)
     addInputStream(ssc).register()