Change to Java enum and added Java test

Author: tdas

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

@@ -32,7 +32,7 @@ import org.apache.hadoop.mapreduce.lib.input.TextInputFormat
 import org.apache.hadoop.mapreduce.{InputFormat => NewInputFormat}
 
 import org.apache.spark._
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{DeveloperApi, Experimental}
 import org.apache.spark.input.FixedLengthBinaryInputFormat
 import org.apache.spark.rdd.RDD
 import org.apache.spark.storage.StorageLevel
@@ -511,8 +511,11 @@ class StreamingContext private[streaming] (
   }
 
   /**
+   * :: DeveloperApi ::
+   *
    * Return the current state of the context.
    */
+  @DeveloperApi
   def getState(): StreamingContextState = {
     state
   }
@@ -601,25 +604,27 @@ class StreamingContext private[streaming] (
    *                       received data to be completed
    */
   def stop(stopSparkContext: Boolean, stopGracefully: Boolean): Unit = synchronized {
-    state = STOPPED
-    StreamingContext.setActiveContext(null)
-
-    state match {
-      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()
+    try {
+      state match {
+        case INITIALIZED =>
+          logWarning("StreamingContext has not been started yet")
+        case STOPPED =>
+          logWarning("StreamingContext has already been stopped")
+        case STARTED =>
+          scheduler.stop(stopGracefully)
+          uiTab.foreach(_.detach())
+          // Even if the streaming context has not been started, we still need to stop the SparkContext.
+          // Even if we have already stopped, we still need to attempt to stop the SparkContext because
+          // a user might stop(stopSparkContext = false) and then call stop(stopSparkContext = true).
+          StreamingContext.setActiveContext(null)
+          logInfo("StreamingContext stopped successfully")
+          waiter.notifyStop()
+      }
+      if (stopSparkContext) sc.stop()
+    } finally {
+      // The state should always be Stopped after calling `stop()`, even if we haven't started yet
+      state = STOPPED
     }
-    // Even if the streaming context has not been started, we still need to stop the SparkContext.
-    // Even if we have already stopped, we still need to attempt to stop the SparkContext because
-    // a user might stop(stopSparkContext = false) and then call stop(stopSparkContext = true).
-    if (stopSparkContext) sc.stop()
-    uiTab.foreach(_.detach())
-    // The state should always be Stopped after calling `stop()`, even if we haven't started yet:
   }
 }
 

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

@@ -0,0 +1,30 @@
+/*
+ * 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.streaming;
+
+import org.apache.spark.annotation.DeveloperApi;
+
+/**
+ * :: DeveloperApi ::
+ *
+ * Represents the state of the StreamingContext.
+ */
+@DeveloperApi
+public enum StreamingContextState {
+  INITIALIZED, STARTED, STOPPED
+}

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

@@ -578,8 +578,9 @@ class JavaStreamingContext(val ssc: StreamingContext) extends Closeable {
     ssc.addStreamingListener(streamingListener)
   }
 
-
   /**
+   * :: DeveloperApi ::
+   *
    * Return the current state of the context.
    */
   def getState(): StreamingContextState = {

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

@@ -70,6 +70,20 @@ public void testInitialization() {
     Assert.assertNotNull(ssc.sparkContext());
   }
 
+  @SuppressWarnings("unchecked")
+  @Test
+  public void testContextState() {
+    List<List<Integer>> inputData = Arrays.asList(Arrays.asList(1, 2, 3, 4));
+    Assert.assertTrue(ssc.getState() == StreamingContextState.INITIALIZED);
+    JavaDStream<Integer> stream = JavaTestUtils.attachTestInputStream(ssc, inputData, 1);
+    JavaTestUtils.attachTestOutputStream(stream);
+    Assert.assertTrue(ssc.getState() == StreamingContextState.INITIALIZED);
+    ssc.start();
+    Assert.assertTrue(ssc.getState() == StreamingContextState.STARTED);
+    ssc.stop();
+    Assert.assertTrue(ssc.getState() == StreamingContextState.STOPPED);
+  }
+
   @SuppressWarnings("unchecked")
   @Test
   public void testCount() {

streaming/src/test/java/org/apache/spark/streaming/JavaAPISuite.java

@@ -70,6 +70,7 @@ trait JavaTestBase extends TestSuiteBase {
       ssc: JavaStreamingContext, numBatches: Int, numExpectedOutput: Int): JList[JList[V]] = {
     implicit val cm: ClassTag[V] =
       implicitly[ClassTag[AnyRef]].asInstanceOf[ClassTag[V]]
+    ssc.getState()
     val res = runStreams[V](ssc.ssc, numBatches, numExpectedOutput)
     val out = new ArrayList[JList[V]]()
     res.map(entry => out.append(new ArrayList[V](entry)))