Abacn commented on code in PR #40090: URL: https://github.com/apache/beam/pull/40090#discussion_r3983385748
########## runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java: ########## @@ -0,0 +1,193 @@ +/* + * 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.beam.runners.spark.structuredstreaming.translation.streaming; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.io.Serializable; +import java.util.HashSet; +import java.util.Set; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.io.streaming.TestUnboundedSource; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.junit.After; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * State transitions for a streaming pipeline. + * + * <p>Pipelines observe RUNNING, DONE once idle, CANCELLED on cancel, and FAILED on query failure. + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class StreamingPipelineLifecycleTest implements Serializable { + + /** Session shared across tests. */ + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); + + /** How long to wait for a query to start before failing. */ + private static final long QUERY_START_TIMEOUT_MILLIS = 60_000L; + + @After + public void tearDown() { + TestUnboundedSource.forget("lifecycle-done"); + TestUnboundedSource.forget("lifecycle-cancel"); + TestUnboundedSource.forget("lifecycle-healthy"); + TestUnboundedSource.forget("lifecycle-poison"); + } + + /** Blocks until at least one streaming query is active on the shared session. */ + private static void awaitQueryStarted() throws InterruptedException { + long deadline = System.currentTimeMillis() + QUERY_START_TIMEOUT_MILLIS; + while (SESSION.getSession().streams().active().length == 0) { + assertTrue( + "no streaming query started within " + QUERY_START_TIMEOUT_MILLIS + "ms", + System.currentTimeMillis() < deadline); + Thread.sleep(50L); + } + } + + @Test + public void idlePipelineGoesFromRunningToDoneOnceIdle() throws Exception { + String tag = "lifecycle-done"; + String collectorId = StreamingTestUtils.newCollectorId(tag); + StreamingTestUtils.clear(collectorId); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + Pipeline pipeline = Pipeline.create(options); + + pipeline + .apply("ReadUnbounded", Read.from(new TestUnboundedSource(tag, 1, 10))) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + PipelineResult result = pipeline.run(); + assertEquals(PipelineResult.State.RUNNING, result.getState()); + + PipelineResult.State finalState = StreamingTestUtils.waitUntilFinish(result); + assertEquals(PipelineResult.State.DONE, finalState); + assertEquals(PipelineResult.State.DONE, result.getState()); + + Set<String> collected = new HashSet<>(StreamingTestUtils.<String>getCollected(collectorId)); + assertEquals(TestUnboundedSource.elements(tag, 1, 10), collected); + } + + @Test + public void cancelStopsTheQueryAndReportsCancelled() throws Exception { + String tag = "lifecycle-cancel"; + String collectorId = StreamingTestUtils.newCollectorId(tag); + StreamingTestUtils.clear(collectorId); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + // Idle stop disabled so the query stops only from explicit cancel. + options.setStreamingStopAfterIdleBatches(-1); + Pipeline pipeline = Pipeline.create(options); + + pipeline + .apply("ReadUnbounded", Read.from(new TestUnboundedSource(tag, 1, 10))) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + PipelineResult result = pipeline.run(); + assertEquals(PipelineResult.State.RUNNING, result.getState()); + + awaitQueryStarted(); + + PipelineResult.State cancelledState = result.cancel(); + assertEquals(PipelineResult.State.CANCELLED, cancelledState); + assertEquals(PipelineResult.State.CANCELLED, result.getState()); + + long deadline = System.currentTimeMillis() + QUERY_START_TIMEOUT_MILLIS; Review Comment: This almost duplicates `awaitQueryStarted()` except the loop condition. Consider generalize `awaitQueryStarted` -> `awaitQuery(...)` ########## runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java: ########## @@ -0,0 +1,106 @@ +/* + * 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.beam.runners.spark.structuredstreaming.translation; + +import java.util.Collection; +import org.apache.beam.runners.spark.SparkCommonPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.translation.batch.PipelineTranslatorBatch; +import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.ReadUnboundedTranslator; +import org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.sdk.transforms.Combine; +import org.apache.beam.sdk.transforms.GroupByKey; +import org.apache.beam.sdk.transforms.Impulse; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.reflect.DoFnSignature; +import org.apache.beam.sdk.transforms.reflect.DoFnSignatures; +import org.apache.beam.sdk.util.construction.SplittableParDo; +import org.apache.beam.sdk.values.PInput; +import org.apache.beam.sdk.values.POutput; +import org.apache.spark.sql.SparkSession; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Pipeline translator for streaming pipelines on Spark 4. It extends the batch translator to reuse + * the stateless single output ParDo, Window.Assign, Flatten and Reshuffle translators, which are + * safe on a streaming Dataset. Every other primitive fails at translation, the batch translators + * for them persist or collect the Dataset, which Spark rejects for streaming plans. + */ +@Internal +public class PipelineTranslatorStreaming extends PipelineTranslatorBatch { + + /** Returns a {@link TransformTranslator} for the given {@link PTransform} if known. */ + @Override + @SuppressWarnings({"rawtypes", "unchecked"}) Review Comment: Please clean up SuppressWarnings ########## runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java: ########## @@ -0,0 +1,106 @@ +/* + * 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.beam.runners.spark.structuredstreaming.translation; + +import java.util.Collection; +import org.apache.beam.runners.spark.SparkCommonPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.translation.batch.PipelineTranslatorBatch; +import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.ReadUnboundedTranslator; +import org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.sdk.transforms.Combine; +import org.apache.beam.sdk.transforms.GroupByKey; +import org.apache.beam.sdk.transforms.Impulse; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.reflect.DoFnSignature; +import org.apache.beam.sdk.transforms.reflect.DoFnSignatures; +import org.apache.beam.sdk.util.construction.SplittableParDo; +import org.apache.beam.sdk.values.PInput; +import org.apache.beam.sdk.values.POutput; +import org.apache.spark.sql.SparkSession; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Pipeline translator for streaming pipelines on Spark 4. It extends the batch translator to reuse + * the stateless single output ParDo, Window.Assign, Flatten and Reshuffle translators, which are + * safe on a streaming Dataset. Every other primitive fails at translation, the batch translators + * for them persist or collect the Dataset, which Spark rejects for streaming plans. + */ +@Internal +public class PipelineTranslatorStreaming extends PipelineTranslatorBatch { + + /** Returns a {@link TransformTranslator} for the given {@link PTransform} if known. */ + @Override + @SuppressWarnings({"rawtypes", "unchecked"}) + @Nullable + protected <InT extends PInput, OutT extends POutput, TransformT extends PTransform<InT, OutT>> + TransformTranslator<InT, OutT, TransformT> getTransformTranslator(TransformT transform) { + + if (transform instanceof SplittableParDo.PrimitiveUnboundedRead) { + return (TransformTranslator) new ReadUnboundedTranslator<>(); + } + + if (transform instanceof SplittableParDo.PrimitiveBoundedRead) { + throw unsupported( + "Bounded Read (Read.from(BoundedSource), Create with two or more elements)"); + } + + if (transform instanceof Impulse) { + throw unsupported("Impulse (Create with fewer than two elements, PAssert)"); + } + + if (transform instanceof GroupByKey) { + throw unsupported("GroupByKey"); + } + + if (transform instanceof Combine.PerKey) { + throw unsupported("Combine.perKey"); + } + + if (transform instanceof ParDo.MultiOutput) { + ParDo.MultiOutput<?, ?> parDo = (ParDo.MultiOutput<?, ?>) transform; + DoFnSignature signature = DoFnSignatures.signatureForDoFn(parDo.getFn()); + if (signature.usesState() || signature.usesTimers()) { + throw unsupported("Stateful ParDo (" + signature.fnClass().getName() + ")"); + } + if (!parDo.getSideInputs().isEmpty()) { + throw unsupported("ParDo with side inputs (" + signature.fnClass().getName() + ")"); + } + if (!parDo.getAdditionalOutputTags().getAll().isEmpty()) { + throw unsupported("ParDo with additional outputs (" + signature.fnClass().getName() + ")"); + } + } + + return super.getTransformTranslator(transform); + } + + private static UnsupportedOperationException unsupported(String what) { Review Comment: This is a thin wrapper. In convention exceptions should be created at the exact moment an error occurs. ########## runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java: ########## @@ -0,0 +1,106 @@ +/* + * 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.beam.runners.spark.structuredstreaming.translation; + +import java.util.Collection; +import org.apache.beam.runners.spark.SparkCommonPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.translation.batch.PipelineTranslatorBatch; +import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.ReadUnboundedTranslator; +import org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.sdk.transforms.Combine; +import org.apache.beam.sdk.transforms.GroupByKey; +import org.apache.beam.sdk.transforms.Impulse; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.reflect.DoFnSignature; +import org.apache.beam.sdk.transforms.reflect.DoFnSignatures; +import org.apache.beam.sdk.util.construction.SplittableParDo; +import org.apache.beam.sdk.values.PInput; +import org.apache.beam.sdk.values.POutput; +import org.apache.spark.sql.SparkSession; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Pipeline translator for streaming pipelines on Spark 4. It extends the batch translator to reuse + * the stateless single output ParDo, Window.Assign, Flatten and Reshuffle translators, which are + * safe on a streaming Dataset. Every other primitive fails at translation, the batch translators + * for them persist or collect the Dataset, which Spark rejects for streaming plans. + */ +@Internal +public class PipelineTranslatorStreaming extends PipelineTranslatorBatch { Review Comment: While it works and resuses some code, future change in PipelineTranslatorBatch may result in surprises on streaming path. For example, new translations meant only for batch now silently port to streaming. On the other hand, refactoring it sounds an risky choice either. At minimum can we rename "PipelineTranslatorBatch" to "PipelineTranslatorCommon" and create a thin subclass "PipelineTranslatorBatch" to it? ########## runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java: ########## @@ -0,0 +1,258 @@ +/* + * 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.beam.runners.spark.structuredstreaming.translation; + +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.beam.runners.spark.SparkCommonPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.sdk.annotations.Internal; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.streaming.StreamingQuery; +import org.apache.spark.sql.streaming.StreamingQueryException; +import org.apache.spark.sql.streaming.StreamingQueryListener; +import org.apache.spark.sql.streaming.Trigger; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Starts one Spark Structured Streaming query per leaf dataset and blocks until all of them reach a + * terminal state. Queries end through {@link #stop()} or the idle stop listener. + * + * <p>Leaf {@code i} checkpoints under {@code <checkpointDir>/i}, in pipeline graph order. A changed + * pipeline needs a new checkpoint directory, as with any Spark streaming query. + */ +@Internal +public class StreamingEvaluationContext extends EvaluationContext { + private static final Logger LOG = LoggerFactory.getLogger(StreamingEvaluationContext.class); + + private static final long AWAIT_POLL_TIMEOUT_MILLIS = 100; + + private final SparkStructuredStreamingPipelineOptions options; + + // Guards queries and stopped. + private final Object lock = new Object(); + private final List<StreamingQuery> queries = new ArrayList<>(); + private boolean stopped = false; + + StreamingEvaluationContext( + Collection<? extends NamedDataset<?>> leaves, + SparkSession session, + SparkCommonPipelineOptions options) { + super(leaves, session); + this.options = options.as(SparkStructuredStreamingPipelineOptions.class); + } + + /** Starts one streaming query per leaf dataset and blocks until all queries terminate. */ + @Override + public void evaluate() { + String checkpointBaseDir = options.getCheckpointDir(); + checkArgument( + checkpointBaseDir != null && !checkpointBaseDir.isEmpty(), + "checkpointDir must be set for a streaming pipeline"); + int idleStopThreshold = options.getStreamingStopAfterIdleBatches(); + + StreamingQueryListener idleStopListener = null; + if (idleStopThreshold >= 0) { + idleStopListener = new IdleStopListener(idleStopThreshold); + getSparkSession().streams().addListener(idleStopListener); + } + + try { + int leafIndex = 0; + for (NamedDataset<?> ds : leaves()) { + Dataset<?> dataset = ds.dataset(); + if (dataset == null) { + continue; + } + synchronized (lock) { + if (stopped) { + break; + } + } + if (!dataset.isStreaming()) { + EvaluationContext.evaluate(ds.name(), dataset); + continue; + } + + StreamingQuery query = startQuery(dataset, checkpointBaseDir, leafIndex++, options); + boolean alreadyStopped; + synchronized (lock) { + queries.add(query); + alreadyStopped = stopped; + } + if (alreadyStopped) { + stopQuery(query); + } + } + + List<StreamingQuery> toAwait; + synchronized (lock) { + toAwait = new ArrayList<>(queries); + } + awaitTermination(toAwait); + } finally { + if (idleStopListener != null) { + getSparkSession().streams().removeListener(idleStopListener); + } + } + } + + /** + * Stops all queries started by {@link #evaluate()}. This method is idempotent and thread safe. + */ + @Override + public void stop() { + List<StreamingQuery> toStop; + synchronized (lock) { + if (stopped) { + return; + } + stopped = true; + toStop = new ArrayList<>(queries); + } + for (StreamingQuery query : toStop) { + stopQuery(query); + } + } + + private StreamingQuery startQuery( + Dataset<?> dataset, + String checkpointBaseDir, + int leafIndex, + SparkStructuredStreamingPipelineOptions options) { + try { + return dataset + .writeStream() + .format("noop") + .outputMode("append") + .option("checkpointLocation", checkpointBaseDir + "/" + leafIndex) + .trigger(Trigger.ProcessingTime(options.getMaxBatchDurationMillis())) + .start(); + } catch (TimeoutException e) { + throw new RuntimeException( + "Failed to start streaming query for leaf dataset index " + leafIndex, e); + } + } + + /** Blocks until every query in toAwait has terminated. Sibling queries stop on failure. */ + private void awaitTermination(List<StreamingQuery> toAwait) { + List<StreamingQuery> active = new ArrayList<>(toAwait); + while (!active.isEmpty()) { + Iterator<StreamingQuery> iterator = active.iterator(); + while (iterator.hasNext()) { + StreamingQuery query = iterator.next(); + try { + if (query.awaitTermination(AWAIT_POLL_TIMEOUT_MILLIS)) { + iterator.remove(); + } + } catch (StreamingQueryException e) { + LOG.error("Streaming query {} terminated with an exception.", query.id(), e); + stop(); + throw new RuntimeException(e); + } + } + } + } + + /** Stops a single query if active. */ + private void stopQuery(StreamingQuery query) { + try { + if (query.isActive()) { + query.stop(); + } + } catch (TimeoutException | RuntimeException e) { + LOG.warn( + "Error while stopping streaming query {}: {}", + query.id(), + String.valueOf(e.getMessage())); Review Comment: redundant String.valueOf(String) ########## runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java: ########## @@ -0,0 +1,193 @@ +/* + * 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.beam.runners.spark.structuredstreaming.translation.streaming; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.io.Serializable; +import java.util.HashSet; +import java.util.Set; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.io.streaming.TestUnboundedSource; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.junit.After; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * State transitions for a streaming pipeline. + * + * <p>Pipelines observe RUNNING, DONE once idle, CANCELLED on cancel, and FAILED on query failure. + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class StreamingPipelineLifecycleTest implements Serializable { + + /** Session shared across tests. */ + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); + + /** How long to wait for a query to start before failing. */ + private static final long QUERY_START_TIMEOUT_MILLIS = 60_000L; + + @After + public void tearDown() { + TestUnboundedSource.forget("lifecycle-done"); + TestUnboundedSource.forget("lifecycle-cancel"); + TestUnboundedSource.forget("lifecycle-healthy"); Review Comment: This is unusual. Generic teardown that runs after every test seems to clear some test specific resources. ########## runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java: ########## @@ -0,0 +1,193 @@ +/* + * 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.beam.runners.spark.structuredstreaming.translation.streaming; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.io.Serializable; +import java.util.HashSet; +import java.util.Set; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.io.streaming.TestUnboundedSource; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.junit.After; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * State transitions for a streaming pipeline. + * + * <p>Pipelines observe RUNNING, DONE once idle, CANCELLED on cancel, and FAILED on query failure. + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class StreamingPipelineLifecycleTest implements Serializable { + + /** Session shared across tests. */ + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); + + /** How long to wait for a query to start before failing. */ + private static final long QUERY_START_TIMEOUT_MILLIS = 60_000L; + + @After + public void tearDown() { + TestUnboundedSource.forget("lifecycle-done"); + TestUnboundedSource.forget("lifecycle-cancel"); + TestUnboundedSource.forget("lifecycle-healthy"); + TestUnboundedSource.forget("lifecycle-poison"); + } + + /** Blocks until at least one streaming query is active on the shared session. */ + private static void awaitQueryStarted() throws InterruptedException { + long deadline = System.currentTimeMillis() + QUERY_START_TIMEOUT_MILLIS; + while (SESSION.getSession().streams().active().length == 0) { + assertTrue( + "no streaming query started within " + QUERY_START_TIMEOUT_MILLIS + "ms", + System.currentTimeMillis() < deadline); + Thread.sleep(50L); + } + } + + @Test + public void idlePipelineGoesFromRunningToDoneOnceIdle() throws Exception { + String tag = "lifecycle-done"; + String collectorId = StreamingTestUtils.newCollectorId(tag); + StreamingTestUtils.clear(collectorId); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + Pipeline pipeline = Pipeline.create(options); + + pipeline + .apply("ReadUnbounded", Read.from(new TestUnboundedSource(tag, 1, 10))) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + PipelineResult result = pipeline.run(); + assertEquals(PipelineResult.State.RUNNING, result.getState()); + + PipelineResult.State finalState = StreamingTestUtils.waitUntilFinish(result); + assertEquals(PipelineResult.State.DONE, finalState); + assertEquals(PipelineResult.State.DONE, result.getState()); + + Set<String> collected = new HashSet<>(StreamingTestUtils.<String>getCollected(collectorId)); + assertEquals(TestUnboundedSource.elements(tag, 1, 10), collected); + } + + @Test + public void cancelStopsTheQueryAndReportsCancelled() throws Exception { + String tag = "lifecycle-cancel"; + String collectorId = StreamingTestUtils.newCollectorId(tag); + StreamingTestUtils.clear(collectorId); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + // Idle stop disabled so the query stops only from explicit cancel. + options.setStreamingStopAfterIdleBatches(-1); + Pipeline pipeline = Pipeline.create(options); + + pipeline + .apply("ReadUnbounded", Read.from(new TestUnboundedSource(tag, 1, 10))) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + PipelineResult result = pipeline.run(); + assertEquals(PipelineResult.State.RUNNING, result.getState()); + + awaitQueryStarted(); + + PipelineResult.State cancelledState = result.cancel(); + assertEquals(PipelineResult.State.CANCELLED, cancelledState); + assertEquals(PipelineResult.State.CANCELLED, result.getState()); + + long deadline = System.currentTimeMillis() + QUERY_START_TIMEOUT_MILLIS; + while (SESSION.getSession().streams().active().length > 0) { + assertTrue( + "the streaming query was still active " + QUERY_START_TIMEOUT_MILLIS + "ms after cancel", + System.currentTimeMillis() < deadline); + Thread.sleep(50L); + } + } + + /** A failure in any leaf query surfaces through waitUntilFinish. */ + @Test + public void failingLeafQueryFailsThePipelineAndStopsHealthySibling() throws Exception { + String healthyTag = "lifecycle-healthy"; + String poisonTag = "lifecycle-poison"; + String collectorId = StreamingTestUtils.newCollectorId(healthyTag); + StreamingTestUtils.clear(collectorId); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + // Idle stop disabled so the healthy query stops only when the sibling failure stops it. + options.setStreamingStopAfterIdleBatches(-1); + Pipeline pipeline = Pipeline.create(options); + + pipeline + .apply("ReadHealthy", Read.from(new TestUnboundedSource(healthyTag, 1, 10))) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + pipeline + .apply("ReadPoisoned", Read.from(new TestUnboundedSource(poisonTag, 1, 10))) + .apply("Throw", ParDo.of(new ThrowOnElementDoFn(5))); + + PipelineResult result = pipeline.run(); + + assertThrows(RuntimeException.class, () -> StreamingTestUtils.waitUntilFinish(result)); + assertEquals(PipelineResult.State.FAILED, result.getState()); + + long deadline = System.currentTimeMillis() + QUERY_START_TIMEOUT_MILLIS; + while (SESSION.getSession().streams().active().length > 0) { Review Comment: same here (awaitQueryStarted) ########## runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java: ########## @@ -0,0 +1,258 @@ +/* + * 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.beam.runners.spark.structuredstreaming.translation; + +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.beam.runners.spark.SparkCommonPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.sdk.annotations.Internal; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.streaming.StreamingQuery; +import org.apache.spark.sql.streaming.StreamingQueryException; +import org.apache.spark.sql.streaming.StreamingQueryListener; +import org.apache.spark.sql.streaming.Trigger; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Starts one Spark Structured Streaming query per leaf dataset and blocks until all of them reach a + * terminal state. Queries end through {@link #stop()} or the idle stop listener. + * + * <p>Leaf {@code i} checkpoints under {@code <checkpointDir>/i}, in pipeline graph order. A changed + * pipeline needs a new checkpoint directory, as with any Spark streaming query. + */ +@Internal +public class StreamingEvaluationContext extends EvaluationContext { + private static final Logger LOG = LoggerFactory.getLogger(StreamingEvaluationContext.class); + + private static final long AWAIT_POLL_TIMEOUT_MILLIS = 100; + + private final SparkStructuredStreamingPipelineOptions options; + + // Guards queries and stopped. + private final Object lock = new Object(); + private final List<StreamingQuery> queries = new ArrayList<>(); + private boolean stopped = false; + + StreamingEvaluationContext( + Collection<? extends NamedDataset<?>> leaves, + SparkSession session, + SparkCommonPipelineOptions options) { + super(leaves, session); + this.options = options.as(SparkStructuredStreamingPipelineOptions.class); + } + + /** Starts one streaming query per leaf dataset and blocks until all queries terminate. */ + @Override + public void evaluate() { + String checkpointBaseDir = options.getCheckpointDir(); + checkArgument( + checkpointBaseDir != null && !checkpointBaseDir.isEmpty(), + "checkpointDir must be set for a streaming pipeline"); + int idleStopThreshold = options.getStreamingStopAfterIdleBatches(); + + StreamingQueryListener idleStopListener = null; + if (idleStopThreshold >= 0) { + idleStopListener = new IdleStopListener(idleStopThreshold); + getSparkSession().streams().addListener(idleStopListener); + } + + try { + int leafIndex = 0; + for (NamedDataset<?> ds : leaves()) { + Dataset<?> dataset = ds.dataset(); + if (dataset == null) { + continue; + } + synchronized (lock) { + if (stopped) { + break; + } + } + if (!dataset.isStreaming()) { + EvaluationContext.evaluate(ds.name(), dataset); + continue; + } + + StreamingQuery query = startQuery(dataset, checkpointBaseDir, leafIndex++, options); + boolean alreadyStopped; + synchronized (lock) { + queries.add(query); + alreadyStopped = stopped; + } + if (alreadyStopped) { + stopQuery(query); + } + } + + List<StreamingQuery> toAwait; + synchronized (lock) { + toAwait = new ArrayList<>(queries); + } + awaitTermination(toAwait); + } finally { Review Comment: Is there a risk of leak query on exception thrown in the try block? should we call `stopQuery` for all remaining queries in a catch? ########## runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java: ########## @@ -0,0 +1,258 @@ +/* + * 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.beam.runners.spark.structuredstreaming.translation; + +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.beam.runners.spark.SparkCommonPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.sdk.annotations.Internal; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.streaming.StreamingQuery; +import org.apache.spark.sql.streaming.StreamingQueryException; +import org.apache.spark.sql.streaming.StreamingQueryListener; +import org.apache.spark.sql.streaming.Trigger; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Starts one Spark Structured Streaming query per leaf dataset and blocks until all of them reach a + * terminal state. Queries end through {@link #stop()} or the idle stop listener. + * + * <p>Leaf {@code i} checkpoints under {@code <checkpointDir>/i}, in pipeline graph order. A changed + * pipeline needs a new checkpoint directory, as with any Spark streaming query. + */ +@Internal +public class StreamingEvaluationContext extends EvaluationContext { + private static final Logger LOG = LoggerFactory.getLogger(StreamingEvaluationContext.class); + + private static final long AWAIT_POLL_TIMEOUT_MILLIS = 100; + + private final SparkStructuredStreamingPipelineOptions options; + + // Guards queries and stopped. + private final Object lock = new Object(); + private final List<StreamingQuery> queries = new ArrayList<>(); + private boolean stopped = false; + + StreamingEvaluationContext( + Collection<? extends NamedDataset<?>> leaves, + SparkSession session, + SparkCommonPipelineOptions options) { + super(leaves, session); + this.options = options.as(SparkStructuredStreamingPipelineOptions.class); + } + + /** Starts one streaming query per leaf dataset and blocks until all queries terminate. */ + @Override + public void evaluate() { + String checkpointBaseDir = options.getCheckpointDir(); + checkArgument( + checkpointBaseDir != null && !checkpointBaseDir.isEmpty(), + "checkpointDir must be set for a streaming pipeline"); + int idleStopThreshold = options.getStreamingStopAfterIdleBatches(); + + StreamingQueryListener idleStopListener = null; + if (idleStopThreshold >= 0) { + idleStopListener = new IdleStopListener(idleStopThreshold); + getSparkSession().streams().addListener(idleStopListener); + } + + try { + int leafIndex = 0; + for (NamedDataset<?> ds : leaves()) { + Dataset<?> dataset = ds.dataset(); + if (dataset == null) { + continue; + } + synchronized (lock) { + if (stopped) { + break; + } + } + if (!dataset.isStreaming()) { + EvaluationContext.evaluate(ds.name(), dataset); + continue; + } + + StreamingQuery query = startQuery(dataset, checkpointBaseDir, leafIndex++, options); + boolean alreadyStopped; + synchronized (lock) { + queries.add(query); + alreadyStopped = stopped; + } + if (alreadyStopped) { + stopQuery(query); + } + } + + List<StreamingQuery> toAwait; + synchronized (lock) { + toAwait = new ArrayList<>(queries); + } + awaitTermination(toAwait); + } finally { + if (idleStopListener != null) { + getSparkSession().streams().removeListener(idleStopListener); + } + } + } + + /** + * Stops all queries started by {@link #evaluate()}. This method is idempotent and thread safe. + */ + @Override + public void stop() { + List<StreamingQuery> toStop; + synchronized (lock) { + if (stopped) { + return; + } + stopped = true; + toStop = new ArrayList<>(queries); + } + for (StreamingQuery query : toStop) { + stopQuery(query); + } + } + + private StreamingQuery startQuery( + Dataset<?> dataset, + String checkpointBaseDir, + int leafIndex, + SparkStructuredStreamingPipelineOptions options) { + try { + return dataset + .writeStream() + .format("noop") + .outputMode("append") + .option("checkpointLocation", checkpointBaseDir + "/" + leafIndex) Review Comment: consider use builtin path join method to handle trailing paths in checkpointBaseDir ########## runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java: ########## @@ -0,0 +1,154 @@ +/* + * 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.beam.runners.spark.structuredstreaming.translation.streaming; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingRunner; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; +import org.joda.time.Duration; +import org.junit.rules.TemporaryFolder; + +/** + * Shared test utilities for the Spark 4 streaming translators. + * + * <p>Collectors are static so they work in local mode only. + * + * <p>The {@link #run} helper bounds each query at five minutes and cancels on expiry. + * + * <p>Tests poll with deadlines instead of {@code @Test(timeout)}, JUnit runs a timed test in a + * separate thread group and Spark's static thread pools inherit it. + */ +public final class StreamingTestUtils { + + private StreamingTestUtils() {} + + /** Driver side, per collector id accumulation of every element a {@link CollectDoFn} saw. */ + private static final Map<String, List<Object>> COLLECTORS = new ConcurrentHashMap<>(); + + /** + * Appends every element to a static collector named {@code collectorId}, then passes it through. + * Safe to use concurrently. Works in Spark local mode only. + */ + public static final class CollectDoFn<T> extends DoFn<T, T> { + private final String collectorId; + + public CollectDoFn(String collectorId) { + this.collectorId = Preconditions.checkNotNull(collectorId); + } + + @ProcessElement + public void processElement(@Element T element, OutputReceiver<T> out) { + append(collectorId, element); + out.output(element); + } + } + + private static void append(String collectorId, Object value) { + COLLECTORS + .computeIfAbsent(collectorId, unused -> Collections.synchronizedList(new ArrayList<>())) + .add(value); + } + + /** Returns a snapshot of everything collected so far under {@code collectorId}. */ + @SuppressWarnings("unchecked") + public static <T> List<T> getCollected(String collectorId) { Review Comment: getCollected defensively new an ArrayList, while it's callers also redundantly new an ArrayList / Set. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
