Merge pull request #23204: Properly close Spark (streaming) context i…

…f Pipeline translation fails
apache · Sep 13, 2022 · bee135e · bee135e
2 parents d940595 + d635d44
commit bee135e
Showing 1 changed file with 15 additions and 9 deletions.
diff --git a/...g/apache/beam/runners/spark/translation/streaming/SparkRunnerStreamingContextFactory.java b/...g/apache/beam/runners/spark/translation/streaming/SparkRunnerStreamingContextFactory.java
@@ -75,22 +75,28 @@ public JavaStreamingContext call() throws Exception {
     Duration batchDuration = new Duration(options.getBatchIntervalMillis());
     LOG.info("Setting Spark streaming batchDuration to {} msec", batchDuration.milliseconds());
 
-    JavaSparkContext jsc = SparkContextFactory.getSparkContext(options);
-    JavaStreamingContext jssc = new JavaStreamingContext(jsc, batchDuration);
+    JavaSparkContext sparkCtx = SparkContextFactory.getSparkContext(options);
+    JavaStreamingContext streamingCtx = new JavaStreamingContext(sparkCtx, batchDuration);
 
     // We must first init accumulators since translators expect them to be instantiated.
-    SparkRunner.initAccumulators(options, jsc);
+    SparkRunner.initAccumulators(options, sparkCtx);
     // do not need to create a MetricsPusher instance here because if is called in SparkRunner.run()
 
-    EvaluationContext ctxt = new EvaluationContext(jsc, pipeline, options, jssc);
+    EvaluationContext evalCtx = new EvaluationContext(sparkCtx, pipeline, options, streamingCtx);
     // update cache candidates
-    SparkRunner.updateCacheCandidates(pipeline, translator, ctxt);
-    pipeline.traverseTopologically(new SparkRunner.Evaluator(translator, ctxt));
-    ctxt.computeOutputs();
+    SparkRunner.updateCacheCandidates(pipeline, translator, evalCtx);
+    try {
+      pipeline.traverseTopologically(new SparkRunner.Evaluator(translator, evalCtx));
+    } catch (RuntimeException e) {
+      streamingCtx.stop(false, false);
+      SparkContextFactory.stopSparkContext(sparkCtx);
+      throw e;
+    }
+    evalCtx.computeOutputs();
 
-    checkpoint(jssc, checkpointDir);
+    checkpoint(streamingCtx, checkpointDir);
 
-    return jssc;
+    return streamingCtx;
   }
 
   private void checkpoint(JavaStreamingContext jssc, CheckpointDir checkpointDir) {